datafusion_physical_expr/window/
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::standard::add_new_ordering_expr_with_partition_by;
26use crate::window::window_expr::{
27    AggregateWindowExpr, WindowEvalContext, WindowFn, filter_array,
28};
29use crate::window::{
30    PartitionBatches, PartitionWindowAggStates, SlidingAggregateWindowExpr, WindowExpr,
31};
32use crate::{EquivalenceProperties, PhysicalExpr};
33
34use arrow::array::ArrayRef;
35use arrow::array::BooleanArray;
36use arrow::datatypes::FieldRef;
37use arrow::record_batch::RecordBatch;
38use datafusion_common::{Result, ScalarValue, exec_datafusion_err};
39use datafusion_expr::{Accumulator, WindowFrame, WindowFrameBound, WindowFrameUnits};
40use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
41
42/// A window expr that takes the form of an aggregate function.
43///
44/// See comments on [`WindowExpr`] for more details.
45#[derive(Debug)]
46pub struct PlainAggregateWindowExpr {
47    aggregate: Arc<AggregateFunctionExpr>,
48    partition_by: Vec<Arc<dyn PhysicalExpr>>,
49    order_by: Vec<PhysicalSortExpr>,
50    window_frame: Arc<WindowFrame>,
51    is_constant_in_partition: bool,
52    filter: Option<Arc<dyn PhysicalExpr>>,
53}
54
55impl PlainAggregateWindowExpr {
56    /// Create a new aggregate window function expression
57    pub fn new(
58        aggregate: Arc<AggregateFunctionExpr>,
59        partition_by: &[Arc<dyn PhysicalExpr>],
60        order_by: &[PhysicalSortExpr],
61        window_frame: Arc<WindowFrame>,
62        filter: Option<Arc<dyn PhysicalExpr>>,
63    ) -> Self {
64        let is_constant_in_partition =
65            Self::is_window_constant_in_partition(order_by, &window_frame);
66        Self {
67            aggregate,
68            partition_by: partition_by.to_vec(),
69            order_by: order_by.to_vec(),
70            window_frame,
71            is_constant_in_partition,
72            filter,
73        }
74    }
75
76    /// Get aggregate expr of AggregateWindowExpr
77    pub fn get_aggregate_expr(&self) -> &AggregateFunctionExpr {
78        &self.aggregate
79    }
80
81    pub fn add_equal_orderings(
82        &self,
83        eq_properties: &mut EquivalenceProperties,
84        window_expr_index: usize,
85    ) -> Result<()> {
86        if let Some(expr) = self
87            .get_aggregate_expr()
88            .get_result_ordering(window_expr_index)
89        {
90            add_new_ordering_expr_with_partition_by(
91                eq_properties,
92                expr,
93                &self.partition_by,
94            )?;
95        }
96        Ok(())
97    }
98
99    // Returns true if every row in the partition has the same window frame. This allows
100    // for preventing bound + function calculation for every row due to the values being the
101    // same.
102    //
103    // This occurs when both bounds fall under either condition below:
104    //  1. Bound is unbounded (`Preceding` or `Following`)
105    //  2. Bound is `CurrentRow` while using `Range` units with no order by clause
106    //  This results in an invalid range specification. Following PostgreSQL’s convention,
107    //  we interpret this as the entire partition being used for the current window frame.
108    fn is_window_constant_in_partition(
109        order_by: &[PhysicalSortExpr],
110        window_frame: &WindowFrame,
111    ) -> bool {
112        let is_constant_bound = |bound: &WindowFrameBound| match bound {
113            WindowFrameBound::CurrentRow => {
114                window_frame.units == WindowFrameUnits::Range && order_by.is_empty()
115            }
116            _ => bound.is_unbounded(),
117        };
118
119        is_constant_bound(&window_frame.start_bound)
120            && is_constant_bound(&window_frame.end_bound)
121    }
122}
123
124/// peer based evaluation based on the fact that batch is pre-sorted given the sort columns
125/// and then per partition point we'll evaluate the peer group (e.g. SUM or MAX gives the same
126/// results for peers) and concatenate the results.
127impl WindowExpr for PlainAggregateWindowExpr {
128    /// Return a reference to Any that can be used for downcasting
129    fn as_any(&self) -> &dyn Any {
130        self
131    }
132
133    fn field(&self) -> Result<FieldRef> {
134        Ok(self.aggregate.field())
135    }
136
137    fn name(&self) -> &str {
138        self.aggregate.name()
139    }
140
141    fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> {
142        self.aggregate.expressions()
143    }
144
145    fn evaluate(&self, batch: &RecordBatch) -> Result<ArrayRef> {
146        self.aggregate_evaluate(batch)
147    }
148
149    fn evaluate_stateful(
150        &self,
151        partition_batches: &PartitionBatches,
152        window_agg_state: &mut PartitionWindowAggStates,
153        eval_ctx: &WindowEvalContext<'_>,
154    ) -> Result<()> {
155        self.aggregate_evaluate_stateful(partition_batches, window_agg_state, eval_ctx)?;
156
157        // Update window frame range for each partition. As we know that
158        // non-sliding aggregations will never call `retract_batch`, this value
159        // can safely increase, and we can remove "old" parts of the state.
160        // This enables us to run queries involving UNBOUNDED PRECEDING frames
161        // using bounded memory for suitable aggregations.
162        for partition_row in partition_batches.keys() {
163            let window_state = window_agg_state
164                .get_mut(partition_row)
165                .ok_or_else(|| exec_datafusion_err!("Cannot find state"))?;
166            let state = &mut window_state.state;
167            if self.window_frame.start_bound.is_unbounded() {
168                state.window_frame_range.start =
169                    state.window_frame_range.end.saturating_sub(1);
170            }
171        }
172        Ok(())
173    }
174
175    fn partition_by(&self) -> &[Arc<dyn PhysicalExpr>] {
176        &self.partition_by
177    }
178
179    fn order_by(&self) -> &[PhysicalSortExpr] {
180        &self.order_by
181    }
182
183    fn get_window_frame(&self) -> &Arc<WindowFrame> {
184        &self.window_frame
185    }
186
187    fn get_reverse_expr(&self) -> Option<Arc<dyn WindowExpr>> {
188        self.aggregate.reverse_expr().map(|reverse_expr| {
189            let reverse_window_frame = self.window_frame.reverse();
190            if reverse_window_frame.is_ever_expanding() {
191                Arc::new(PlainAggregateWindowExpr::new(
192                    Arc::new(reverse_expr),
193                    &self.partition_by.clone(),
194                    &self
195                        .order_by
196                        .iter()
197                        .map(|e| e.reverse())
198                        .collect::<Vec<_>>(),
199                    Arc::new(self.window_frame.reverse()),
200                    self.filter.clone(),
201                )) as _
202            } else {
203                Arc::new(SlidingAggregateWindowExpr::new(
204                    Arc::new(reverse_expr),
205                    &self.partition_by.clone(),
206                    &self
207                        .order_by
208                        .iter()
209                        .map(|e| e.reverse())
210                        .collect::<Vec<_>>(),
211                    Arc::new(self.window_frame.reverse()),
212                    self.filter.clone(),
213                )) as _
214            }
215        })
216    }
217
218    fn uses_bounded_memory(&self) -> bool {
219        !self.window_frame.end_bound.is_unbounded()
220    }
221
222    fn create_window_fn(&self) -> Result<WindowFn> {
223        Ok(WindowFn::Aggregate(self.get_accumulator()?))
224    }
225}
226
227impl AggregateWindowExpr for PlainAggregateWindowExpr {
228    fn get_accumulator(&self) -> Result<Box<dyn Accumulator>> {
229        self.aggregate.create_accumulator()
230    }
231
232    fn filter_expr(&self) -> Option<&Arc<dyn PhysicalExpr>> {
233        self.filter.as_ref()
234    }
235
236    /// For a given range, calculate accumulation result inside the range on
237    /// `value_slice` and update accumulator state.
238    // We assume that `cur_range` contains `last_range` and their start points
239    // are same. In summary if `last_range` is `Range{start: a,end: b}` and
240    // `cur_range` is `Range{start: a1, end: b1}`, it is guaranteed that a1=a and b1>=b.
241    fn get_aggregate_result_inside_range(
242        &self,
243        last_range: &Range<usize>,
244        cur_range: &Range<usize>,
245        value_slice: &[ArrayRef],
246        accumulator: &mut Box<dyn Accumulator>,
247        filter_mask: Option<&BooleanArray>,
248    ) -> Result<ScalarValue> {
249        if cur_range.start == cur_range.end {
250            self.aggregate
251                .default_value(self.aggregate.field().data_type())
252        } else {
253            // Accumulate any new rows that have entered the window:
254            let update_bound = cur_range.end - last_range.end;
255            // A non-sliding aggregation only processes new data, it never
256            // deals with expiring data as its starting point is always the
257            // same point (i.e. the beginning of the table/frame). Hence, we
258            // do not call `retract_batch`.
259            if update_bound > 0 {
260                let slice_mask =
261                    filter_mask.map(|m| m.slice(last_range.end, update_bound));
262                let update: Vec<ArrayRef> = value_slice
263                    .iter()
264                    .map(|v| v.slice(last_range.end, update_bound))
265                    .map(|arr| match &slice_mask {
266                        Some(m) => filter_array(&arr, m),
267                        None => Ok(arr),
268                    })
269                    .collect::<Result<Vec<_>>>()?;
270                accumulator.update_batch(&update)?
271            }
272            accumulator.evaluate()
273        }
274    }
275
276    fn is_constant_in_partition(&self) -> bool {
277        self.is_constant_in_partition
278    }
279}
datafusion_physical_expr/window/aggregate.rs

datafusion_physical_expr/window/
aggregate.rs
