datafusion_physical_expr/window/
sliding_aggregate.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Physical exec for aggregate window function expressions.
19
20use std::any::Any;
21use std::ops::Range;
22use std::sync::Arc;
23
24use crate::aggregate::AggregateFunctionExpr;
25use crate::window::window_expr::{
26    AggregateWindowExpr, WindowEvalContext, WindowFn, filter_array,
27};
28use crate::window::{
29    PartitionBatches, PartitionWindowAggStates, PlainAggregateWindowExpr, WindowExpr,
30};
31use crate::{PhysicalExpr, expressions::PhysicalSortExpr};
32
33use arrow::array::{ArrayRef, BooleanArray};
34use arrow::datatypes::FieldRef;
35use arrow::record_batch::RecordBatch;
36use datafusion_common::{Result, ScalarValue};
37use datafusion_expr::{Accumulator, WindowFrame};
38
39/// A window expr that takes the form of an aggregate function that
40/// can be incrementally computed over sliding windows.
41///
42/// See comments on [`WindowExpr`] for more details.
43#[derive(Debug)]
44pub struct SlidingAggregateWindowExpr {
45    aggregate: Arc<AggregateFunctionExpr>,
46    partition_by: Vec<Arc<dyn PhysicalExpr>>,
47    order_by: Vec<PhysicalSortExpr>,
48    window_frame: Arc<WindowFrame>,
49    filter: Option<Arc<dyn PhysicalExpr>>,
50}
51
52impl SlidingAggregateWindowExpr {
53    /// Create a new (sliding) aggregate window function expression.
54    pub fn new(
55        aggregate: Arc<AggregateFunctionExpr>,
56        partition_by: &[Arc<dyn PhysicalExpr>],
57        order_by: &[PhysicalSortExpr],
58        window_frame: Arc<WindowFrame>,
59        filter: Option<Arc<dyn PhysicalExpr>>,
60    ) -> Self {
61        Self {
62            aggregate,
63            partition_by: partition_by.to_vec(),
64            order_by: order_by.to_vec(),
65            window_frame,
66            filter,
67        }
68    }
69
70    /// Get the [AggregateFunctionExpr] of this object.
71    pub fn get_aggregate_expr(&self) -> &AggregateFunctionExpr {
72        &self.aggregate
73    }
74}
75
76/// Incrementally update window function using the fact that batch is
77/// pre-sorted given the sort columns and then per partition point.
78///
79/// Evaluates the peer group (e.g. `SUM` or `MAX` gives the same results
80/// for peers) and concatenate the results.
81impl WindowExpr for SlidingAggregateWindowExpr {
82    /// Return a reference to Any that can be used for downcasting
83    fn as_any(&self) -> &dyn Any {
84        self
85    }
86
87    fn field(&self) -> Result<FieldRef> {
88        Ok(self.aggregate.field())
89    }
90
91    fn name(&self) -> &str {
92        self.aggregate.name()
93    }
94
95    fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> {
96        self.aggregate.expressions()
97    }
98
99    fn evaluate(&self, batch: &RecordBatch) -> Result<ArrayRef> {
100        self.aggregate_evaluate(batch)
101    }
102
103    fn evaluate_stateful(
104        &self,
105        partition_batches: &PartitionBatches,
106        window_agg_state: &mut PartitionWindowAggStates,
107        eval_ctx: &WindowEvalContext<'_>,
108    ) -> Result<()> {
109        self.aggregate_evaluate_stateful(partition_batches, window_agg_state, eval_ctx)
110    }
111
112    fn partition_by(&self) -> &[Arc<dyn PhysicalExpr>] {
113        &self.partition_by
114    }
115
116    fn order_by(&self) -> &[PhysicalSortExpr] {
117        &self.order_by
118    }
119
120    fn get_window_frame(&self) -> &Arc<WindowFrame> {
121        &self.window_frame
122    }
123
124    fn get_reverse_expr(&self) -> Option<Arc<dyn WindowExpr>> {
125        self.aggregate.reverse_expr().map(|reverse_expr| {
126            let reverse_window_frame = self.window_frame.reverse();
127            if reverse_window_frame.is_ever_expanding() {
128                Arc::new(PlainAggregateWindowExpr::new(
129                    Arc::new(reverse_expr),
130                    &self.partition_by.clone(),
131                    &self
132                        .order_by
133                        .iter()
134                        .map(|e| e.reverse())
135                        .collect::<Vec<_>>(),
136                    Arc::new(self.window_frame.reverse()),
137                    self.filter.clone(),
138                )) as _
139            } else {
140                Arc::new(SlidingAggregateWindowExpr::new(
141                    Arc::new(reverse_expr),
142                    &self.partition_by.clone(),
143                    &self
144                        .order_by
145                        .iter()
146                        .map(|e| e.reverse())
147                        .collect::<Vec<_>>(),
148                    Arc::new(self.window_frame.reverse()),
149                    self.filter.clone(),
150                )) as _
151            }
152        })
153    }
154
155    fn uses_bounded_memory(&self) -> bool {
156        !self.window_frame.end_bound.is_unbounded()
157    }
158
159    fn with_new_expressions(
160        &self,
161        args: Vec<Arc<dyn PhysicalExpr>>,
162        partition_bys: Vec<Arc<dyn PhysicalExpr>>,
163        order_by_exprs: Vec<Arc<dyn PhysicalExpr>>,
164    ) -> Option<Arc<dyn WindowExpr>> {
165        debug_assert_eq!(self.order_by.len(), order_by_exprs.len());
166
167        let new_order_by = self
168            .order_by
169            .iter()
170            .zip(order_by_exprs)
171            .map(|(req, new_expr)| PhysicalSortExpr {
172                expr: new_expr,
173                options: req.options,
174            })
175            .collect();
176        Some(Arc::new(SlidingAggregateWindowExpr {
177            aggregate: self
178                .aggregate
179                .with_new_expressions(args, vec![])
180                .map(Arc::new)?,
181            partition_by: partition_bys,
182            order_by: new_order_by,
183            window_frame: Arc::clone(&self.window_frame),
184            filter: self.filter.clone(),
185        }))
186    }
187
188    fn create_window_fn(&self) -> Result<WindowFn> {
189        Ok(WindowFn::Aggregate(self.get_accumulator()?))
190    }
191}
192
193impl AggregateWindowExpr for SlidingAggregateWindowExpr {
194    fn get_accumulator(&self) -> Result<Box<dyn Accumulator>> {
195        self.aggregate.create_sliding_accumulator()
196    }
197
198    fn filter_expr(&self) -> Option<&Arc<dyn PhysicalExpr>> {
199        self.filter.as_ref()
200    }
201
202    /// Given current range and the last range, calculates the accumulator
203    /// result for the range of interest.
204    fn get_aggregate_result_inside_range(
205        &self,
206        last_range: &Range<usize>,
207        cur_range: &Range<usize>,
208        value_slice: &[ArrayRef],
209        accumulator: &mut Box<dyn Accumulator>,
210        filter_mask: Option<&BooleanArray>,
211    ) -> Result<ScalarValue> {
212        if cur_range.start == cur_range.end {
213            // Keep the accumulator synchronized with `last_range`. RANGE frames
214            // can become empty between two non-empty frames when the ORDER BY
215            // values contain gaps.
216            let retract_bound = last_range.end - last_range.start;
217            if retract_bound > 0 {
218                let slice_mask =
219                    filter_mask.map(|m| m.slice(last_range.start, retract_bound));
220                let retract: Vec<ArrayRef> = value_slice
221                    .iter()
222                    .map(|v| v.slice(last_range.start, retract_bound))
223                    .map(|arr| match &slice_mask {
224                        Some(m) => filter_array(&arr, m),
225                        None => Ok(arr),
226                    })
227                    .collect::<Result<Vec<_>>>()?;
228                accumulator.retract_batch(&retract)?
229            }
230            self.aggregate
231                .default_value(self.aggregate.field().data_type())
232        } else {
233            // Accumulate any new rows that have entered the window:
234            let update_bound = cur_range.end - last_range.end;
235            if update_bound > 0 {
236                let slice_mask =
237                    filter_mask.map(|m| m.slice(last_range.end, update_bound));
238                let update: Vec<ArrayRef> = value_slice
239                    .iter()
240                    .map(|v| v.slice(last_range.end, update_bound))
241                    .map(|arr| match &slice_mask {
242                        Some(m) => filter_array(&arr, m),
243                        None => Ok(arr),
244                    })
245                    .collect::<Result<Vec<_>>>()?;
246                accumulator.update_batch(&update)?
247            }
248
249            // Remove rows that have now left the window:
250            let retract_bound = cur_range.start - last_range.start;
251            if retract_bound > 0 {
252                let slice_mask =
253                    filter_mask.map(|m| m.slice(last_range.start, retract_bound));
254                let retract: Vec<ArrayRef> = value_slice
255                    .iter()
256                    .map(|v| v.slice(last_range.start, retract_bound))
257                    .map(|arr| match &slice_mask {
258                        Some(m) => filter_array(&arr, m),
259                        None => Ok(arr),
260                    })
261                    .collect::<Result<Vec<_>>>()?;
262                accumulator.retract_batch(&retract)?
263            }
264            accumulator.evaluate()
265        }
266    }
267
268    fn is_constant_in_partition(&self) -> bool {
269        false
270    }
271}
datafusion_physical_expr/window/sliding_aggregate.rs

datafusion_physical_expr/window/
sliding_aggregate.rs
