risingwave_expr_impl/window_function/
aggregate.rs

1// Copyright 2024 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::BTreeSet;
16
17use educe::Educe;
18use futures_util::FutureExt;
19use risingwave_common::array::{DataChunk, Op, StreamChunk};
20use risingwave_common::types::{DataType, Datum};
21use risingwave_common::util::iter_util::ZipEqFast;
22use risingwave_common::{bail, must_match};
23use risingwave_common_estimate_size::{EstimateSize, KvSize};
24use risingwave_expr::Result;
25use risingwave_expr::aggregate::{
26    AggCall, AggType, AggregateFunction, AggregateState as AggImplState, BoxedAggregateFunction,
27    PbAggKind, build_append_only,
28};
29use risingwave_expr::sig::FUNCTION_REGISTRY;
30use risingwave_expr::window_function::{
31    BoxedWindowState, FrameBounds, Readiness, StateEvictHint, StateKey, StatePos, WindowFuncCall,
32    WindowFuncKind, WindowState,
33};
34use smallvec::SmallVec;
35
36use super::buffer::{RangeWindow, RowsWindow, SessionWindow, WindowBuffer, WindowImpl};
37
38type StateValue = SmallVec<[Datum; 2]>;
39
40struct AggregateState<W>
41where
42    W: WindowImpl<Key = StateKey, Value = StateValue>,
43{
44    agg_impl: AggImpl,
45    arg_data_types: Vec<DataType>,
46    ignore_nulls: bool,
47    buffer: WindowBuffer<W>,
48    buffer_heap_size: KvSize,
49}
50
51pub(super) fn new(call: &WindowFuncCall) -> Result<BoxedWindowState> {
52    if call.frame.bounds.validate().is_err() {
53        bail!("the window frame must be valid");
54    }
55    let agg_type = must_match!(&call.kind, WindowFuncKind::Aggregate(agg_type) => agg_type);
56    let arg_data_types = call.args.arg_types().to_vec();
57    let agg_call = AggCall {
58        agg_type: agg_type.clone(),
59        args: call.args.clone(),
60        return_type: call.return_type.clone(),
61        column_orders: Vec::new(), // the input is already sorted
62        // TODO(rc): support filter on window function call
63        filter: None,
64        // TODO(rc): support distinct on window function call? PG doesn't support it either.
65        distinct: false,
66        direct_args: vec![],
67    };
68
69    let (agg_impl, enable_delta) = match agg_type {
70        AggType::Builtin(PbAggKind::FirstValue) => (AggImpl::Shortcut(Shortcut::FirstValue), false),
71        AggType::Builtin(PbAggKind::LastValue) => (AggImpl::Shortcut(Shortcut::LastValue), false),
72        AggType::Builtin(kind) => {
73            let agg_func_sig = FUNCTION_REGISTRY
74                .get(*kind, &arg_data_types, &call.return_type)
75                .expect("the agg func must exist");
76            let agg_func = agg_func_sig.build_aggregate(&agg_call)?;
77            let (agg_impl, enable_delta) =
78                if agg_func_sig.is_retractable() && call.frame.exclusion.is_no_others() {
79                    let init_state = agg_func.create_state()?;
80                    (AggImpl::Incremental(agg_func, init_state), true)
81                } else {
82                    (AggImpl::Full(agg_func), false)
83                };
84            (agg_impl, enable_delta)
85        }
86        AggType::UserDefined(_) => {
87            // TODO(rc): utilize `retract` method of embedded UDAF to do incremental aggregation
88            (AggImpl::Full(build_append_only(&agg_call)?), false)
89        }
90        AggType::WrapScalar(_) => {
91            // we have to feed the wrapped scalar function with all the rows in the window,
92            // instead of doing incremental aggregation
93            (AggImpl::Full(build_append_only(&agg_call)?), false)
94        }
95    };
96
97    let this = match &call.frame.bounds {
98        FrameBounds::Rows(frame_bounds) => Box::new(AggregateState {
99            agg_impl,
100            arg_data_types,
101            ignore_nulls: call.ignore_nulls,
102            buffer: WindowBuffer::<RowsWindow<StateKey, StateValue>>::new(
103                RowsWindow::new(frame_bounds.clone()),
104                call.frame.exclusion,
105                enable_delta,
106            ),
107            buffer_heap_size: KvSize::new(),
108        }) as BoxedWindowState,
109        FrameBounds::Range(frame_bounds) => Box::new(AggregateState {
110            agg_impl,
111            arg_data_types,
112            ignore_nulls: call.ignore_nulls,
113            buffer: WindowBuffer::<RangeWindow<StateValue>>::new(
114                RangeWindow::new(frame_bounds.clone()),
115                call.frame.exclusion,
116                enable_delta,
117            ),
118            buffer_heap_size: KvSize::new(),
119        }) as BoxedWindowState,
120        FrameBounds::Session(frame_bounds) => Box::new(AggregateState {
121            agg_impl,
122            arg_data_types,
123            ignore_nulls: call.ignore_nulls,
124            buffer: WindowBuffer::<SessionWindow<StateValue>>::new(
125                SessionWindow::new(frame_bounds.clone()),
126                call.frame.exclusion,
127                enable_delta,
128            ),
129            buffer_heap_size: KvSize::new(),
130        }) as BoxedWindowState,
131    };
132    Ok(this)
133}
134
135impl<W> AggregateState<W>
136where
137    W: WindowImpl<Key = StateKey, Value = StateValue>,
138{
139    fn slide_inner(&mut self) -> StateEvictHint {
140        let removed_keys: BTreeSet<_> = self
141            .buffer
142            .slide()
143            .map(|(k, v)| {
144                v.iter().for_each(|arg| {
145                    self.buffer_heap_size.sub_val(arg);
146                });
147                self.buffer_heap_size.sub_val(&k);
148                k
149            })
150            .collect();
151        if removed_keys.is_empty() {
152            StateEvictHint::CannotEvict(
153                self.buffer
154                    .smallest_key()
155                    .expect("sliding without removing, must have some entry in the buffer")
156                    .clone(),
157            )
158        } else {
159            StateEvictHint::CanEvict(removed_keys)
160        }
161    }
162}
163
164impl<W> WindowState for AggregateState<W>
165where
166    W: WindowImpl<Key = StateKey, Value = StateValue>,
167{
168    fn append(&mut self, key: StateKey, args: SmallVec<[Datum; 2]>) {
169        args.iter().for_each(|arg| {
170            self.buffer_heap_size.add_val(arg);
171        });
172        self.buffer_heap_size.add_val(&key);
173        self.buffer.append(key, args);
174    }
175
176    fn curr_window(&self) -> StatePos<'_> {
177        let window = self.buffer.curr_window();
178        StatePos {
179            key: window.key,
180            is_ready: window.following_saturated,
181        }
182    }
183
184    fn readiness(&self) -> Readiness {
185        self.buffer.readiness()
186    }
187
188    fn is_empty(&self) -> bool {
189        self.buffer.is_empty()
190    }
191
192    fn slide(&mut self) -> Result<(Datum, StateEvictHint)> {
193        let output = match self.agg_impl {
194            AggImpl::Full(ref agg_func) => {
195                let wrapper = AggregatorWrapper {
196                    agg_func: agg_func.as_ref(),
197                    arg_data_types: &self.arg_data_types,
198                };
199                wrapper.aggregate(self.buffer.curr_window_values())
200            }
201            AggImpl::Incremental(ref agg_func, ref mut state) => {
202                let wrapper = AggregatorWrapper {
203                    agg_func: agg_func.as_ref(),
204                    arg_data_types: &self.arg_data_types,
205                };
206                wrapper.update(state, self.buffer.consume_curr_window_values_delta())
207            }
208            AggImpl::Shortcut(shortcut) => match shortcut {
209                Shortcut::FirstValue => Ok(if !self.ignore_nulls {
210                    // no `IGNORE NULLS`
211                    self.buffer
212                        .curr_window_values()
213                        .next()
214                        .and_then(|args| args[0].clone())
215                } else {
216                    // filter out NULLs
217                    self.buffer
218                        .curr_window_values()
219                        .find(|args| args[0].is_some())
220                        .and_then(|args| args[0].clone())
221                }),
222                Shortcut::LastValue => Ok(if !self.ignore_nulls {
223                    self.buffer
224                        .curr_window_values()
225                        .next_back()
226                        .and_then(|args| args[0].clone())
227                } else {
228                    self.buffer
229                        .curr_window_values()
230                        .rev()
231                        .find(|args| args[0].is_some())
232                        .and_then(|args| args[0].clone())
233                }),
234            },
235        }?;
236        let evict_hint = self.slide_inner();
237        Ok((output, evict_hint))
238    }
239
240    fn slide_no_output(&mut self) -> Result<StateEvictHint> {
241        match self.agg_impl {
242            AggImpl::Full(..) => {}
243            AggImpl::Incremental(ref agg_func, ref mut state) => {
244                // for incremental agg, we need to update the state even if the caller doesn't need
245                // the output
246                let wrapper = AggregatorWrapper {
247                    agg_func: agg_func.as_ref(),
248                    arg_data_types: &self.arg_data_types,
249                };
250                wrapper.update(state, self.buffer.consume_curr_window_values_delta())?;
251            }
252            AggImpl::Shortcut(..) => {}
253        };
254        Ok(self.slide_inner())
255    }
256}
257
258impl<W> EstimateSize for AggregateState<W>
259where
260    W: WindowImpl<Key = StateKey, Value = StateValue>,
261{
262    fn estimated_heap_size(&self) -> usize {
263        // Does not count the overhead of the `VecDeque` in the buffer itself.
264        self.arg_data_types.estimated_heap_size() + self.buffer_heap_size.size()
265    }
266}
267
268#[derive(Educe)]
269#[educe(Debug)]
270enum AggImpl {
271    Incremental(#[educe(Debug(ignore))] BoxedAggregateFunction, AggImplState),
272    Full(#[educe(Debug(ignore))] BoxedAggregateFunction),
273    Shortcut(Shortcut),
274}
275
276#[derive(Debug, Clone, Copy)]
277enum Shortcut {
278    FirstValue,
279    LastValue,
280}
281
282struct AggregatorWrapper<'a> {
283    agg_func: &'a dyn AggregateFunction,
284    arg_data_types: &'a [DataType],
285}
286
287impl AggregatorWrapper<'_> {
288    fn aggregate<V>(&self, values: impl IntoIterator<Item = V>) -> Result<Datum>
289    where
290        V: AsRef<[Datum]>,
291    {
292        let mut state = self.agg_func.create_state()?;
293        self.update(
294            &mut state,
295            values.into_iter().map(|args| (Op::Insert, args)),
296        )
297    }
298
299    fn update<V>(
300        &self,
301        state: &mut AggImplState,
302        delta: impl IntoIterator<Item = (Op, V)>,
303    ) -> Result<Datum>
304    where
305        V: AsRef<[Datum]>,
306    {
307        let mut args_builders = self
308            .arg_data_types
309            .iter()
310            .map(|data_type| data_type.create_array_builder(0 /* bad! */))
311            .collect::<Vec<_>>();
312        let mut ops = Vec::new();
313        let mut n_rows = 0;
314        for (op, value) in delta {
315            n_rows += 1;
316            ops.push(op);
317            for (builder, datum) in args_builders.iter_mut().zip_eq_fast(value.as_ref()) {
318                builder.append(datum);
319            }
320        }
321        let columns = args_builders
322            .into_iter()
323            .map(|builder| builder.finish().into())
324            .collect::<Vec<_>>();
325        let chunk = StreamChunk::from_parts(ops, DataChunk::new(columns, n_rows));
326
327        self.agg_func
328            .update(state, &chunk)
329            .now_or_never()
330            .expect("we don't support UDAF currently, so the function should return immediately")?;
331        self.agg_func
332            .get_result(state)
333            .now_or_never()
334            .expect("we don't support UDAF currently, so the function should return immediately")
335    }
336}
337
338#[cfg(test)]
339mod tests {
340    // TODO(rc): need to add some unit tests
341}
risingwave_expr_impl/window_function/aggregate.rs

risingwave_expr_impl/window_function/
aggregate.rs
