datafusion_physical_expr/window/
standard.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 standard window function expressions.
19
20use std::any::Any;
21use std::ops::Range;
22use std::sync::Arc;
23
24use super::{StandardWindowFunctionExpr, WindowExpr};
25use crate::window::window_expr::{WindowEvalContext, WindowFn, get_orderby_values};
26use crate::window::{PartitionBatches, PartitionWindowAggStates, WindowState};
27use crate::{EquivalenceProperties, PhysicalExpr};
28
29use arrow::array::{ArrayRef, new_empty_array};
30use arrow::datatypes::FieldRef;
31use arrow::record_batch::RecordBatch;
32use datafusion_common::utils::evaluate_partition_ranges;
33use datafusion_common::{Result, ScalarValue};
34use datafusion_expr::WindowFrame;
35use datafusion_expr::window_state::{WindowAggState, WindowFrameContext};
36use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
37
38/// A window expr that takes the form of a [`StandardWindowFunctionExpr`].
39#[derive(Debug)]
40pub struct StandardWindowExpr {
41    expr: Arc<dyn StandardWindowFunctionExpr>,
42    partition_by: Vec<Arc<dyn PhysicalExpr>>,
43    order_by: Vec<PhysicalSortExpr>,
44    window_frame: Arc<WindowFrame>,
45}
46
47impl StandardWindowExpr {
48    /// create a new standard window function expression
49    pub fn new(
50        expr: Arc<dyn StandardWindowFunctionExpr>,
51        partition_by: &[Arc<dyn PhysicalExpr>],
52        order_by: &[PhysicalSortExpr],
53        window_frame: Arc<WindowFrame>,
54    ) -> Self {
55        Self {
56            expr,
57            partition_by: partition_by.to_vec(),
58            order_by: order_by.to_vec(),
59            window_frame,
60        }
61    }
62
63    /// Get StandardWindowFunction expr of StandardWindowExpr
64    pub fn get_standard_func_expr(&self) -> &Arc<dyn StandardWindowFunctionExpr> {
65        &self.expr
66    }
67
68    /// Adds any equivalent orderings generated by `self.expr` to `builder`.
69    ///
70    /// If `self.expr` doesn't have an ordering, ordering equivalence properties
71    /// are not updated. Otherwise, ordering equivalence properties are updated
72    /// by the ordering of `self.expr`.
73    pub fn add_equal_orderings(
74        &self,
75        eq_properties: &mut EquivalenceProperties,
76    ) -> Result<()> {
77        let schema = eq_properties.schema();
78        if let Some(fn_res_ordering) = self.expr.get_result_ordering(schema) {
79            add_new_ordering_expr_with_partition_by(
80                eq_properties,
81                fn_res_ordering,
82                &self.partition_by,
83            )?;
84        }
85        Ok(())
86    }
87}
88
89impl WindowExpr for StandardWindowExpr {
90    /// Return a reference to Any that can be used for downcasting
91    fn as_any(&self) -> &dyn Any {
92        self
93    }
94
95    fn name(&self) -> &str {
96        self.expr.name()
97    }
98
99    fn field(&self) -> Result<FieldRef> {
100        self.expr.field()
101    }
102
103    fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> {
104        self.expr.expressions()
105    }
106
107    fn partition_by(&self) -> &[Arc<dyn PhysicalExpr>] {
108        &self.partition_by
109    }
110
111    fn order_by(&self) -> &[PhysicalSortExpr] {
112        &self.order_by
113    }
114
115    fn evaluate(&self, batch: &RecordBatch) -> Result<ArrayRef> {
116        let mut evaluator = self.expr.create_evaluator()?;
117        let num_rows = batch.num_rows();
118        if evaluator.uses_window_frame() {
119            let sort_options = self.order_by.iter().map(|o| o.options).collect();
120            let mut row_wise_results = vec![];
121
122            let mut values = self.evaluate_args(batch)?;
123            let order_bys = get_orderby_values(self.order_by_columns(batch)?);
124            let n_args = values.len();
125            values.extend(order_bys);
126            let order_bys_ref = &values[n_args..];
127
128            let mut window_frame_ctx =
129                WindowFrameContext::new(Arc::clone(&self.window_frame), sort_options);
130            let mut last_range = Range { start: 0, end: 0 };
131            // We iterate on each row to calculate window frame range and window function result
132            for idx in 0..num_rows {
133                let range = window_frame_ctx.calculate_range(
134                    order_bys_ref,
135                    &last_range,
136                    num_rows,
137                    idx,
138                )?;
139                let value = evaluator.evaluate(&values, &range)?;
140                row_wise_results.push(value);
141                last_range = range;
142            }
143            ScalarValue::iter_to_array(row_wise_results)
144        } else if evaluator.include_rank() {
145            let columns = self.order_by_columns(batch)?;
146            let sort_partition_points = evaluate_partition_ranges(num_rows, &columns)?;
147            evaluator.evaluate_all_with_rank(num_rows, &sort_partition_points)
148        } else {
149            let values = self.evaluate_args(batch)?;
150            evaluator.evaluate_all(&values, num_rows)
151        }
152    }
153
154    /// Evaluate the window function against the batch. This function facilitates
155    /// stateful, bounded-memory implementations.
156    fn evaluate_stateful(
157        &self,
158        partition_batches: &PartitionBatches,
159        window_agg_state: &mut PartitionWindowAggStates,
160        _eval_ctx: &WindowEvalContext<'_>,
161    ) -> Result<()> {
162        let field = self.expr.field()?;
163        let out_type = field.data_type();
164        let sort_options = self.order_by.iter().map(|o| o.options).collect::<Vec<_>>();
165        // create a WindowAggState to clone when `window_agg_state` does not contain the respective
166        // group, which is faster than potentially creating a new one at every iteration
167        let new_state = WindowAggState::new(out_type)?;
168        for (partition_row, partition_batch_state) in partition_batches.iter() {
169            let window_state =
170                if let Some(window_state) = window_agg_state.get_mut(partition_row) {
171                    window_state
172                } else {
173                    let evaluator = self.expr.create_evaluator()?;
174                    window_agg_state
175                        .entry(partition_row.clone())
176                        .or_insert(WindowState {
177                            state: new_state.clone(),
178                            window_fn: WindowFn::Builtin(evaluator),
179                            published: false,
180                        })
181                };
182            let evaluator = match &mut window_state.window_fn {
183                WindowFn::Builtin(evaluator) => evaluator,
184                _ => unreachable!(),
185            };
186            let state = &mut window_state.state;
187
188            let batch_ref = &partition_batch_state.record_batch;
189            let mut values = self.evaluate_args(batch_ref)?;
190            let order_bys = if evaluator.uses_window_frame() || evaluator.include_rank() {
191                get_orderby_values(self.order_by_columns(batch_ref)?)
192            } else {
193                vec![]
194            };
195            let n_args = values.len();
196            values.extend(order_bys);
197            let order_bys_ref = &values[n_args..];
198
199            // We iterate on each row to perform a running calculation.
200            let record_batch = &partition_batch_state.record_batch;
201            let num_rows = record_batch.num_rows();
202            let mut row_wise_results: Vec<ScalarValue> = vec![];
203            let is_causal = if evaluator.uses_window_frame() {
204                self.window_frame.is_causal()
205            } else {
206                evaluator.is_causal()
207            };
208            for idx in state.last_calculated_index..num_rows {
209                let frame_range = if evaluator.uses_window_frame() {
210                    state
211                        .window_frame_ctx
212                        .get_or_insert_with(|| {
213                            WindowFrameContext::new(
214                                Arc::clone(&self.window_frame),
215                                sort_options.clone(),
216                            )
217                        })
218                        .calculate_range(
219                            order_bys_ref,
220                            // Start search from the last range
221                            &state.window_frame_range,
222                            num_rows,
223                            idx,
224                        )
225                } else {
226                    evaluator.get_range(idx, num_rows)
227                }?;
228
229                // Exit if the range is non-causal and extends all the way:
230                if frame_range.end == num_rows
231                    && !is_causal
232                    && !partition_batch_state.is_end
233                {
234                    break;
235                }
236                // Update last range
237                state.window_frame_range = frame_range;
238                row_wise_results
239                    .push(evaluator.evaluate(&values, &state.window_frame_range)?);
240            }
241            let out_col = if row_wise_results.is_empty() {
242                new_empty_array(out_type)
243            } else if row_wise_results.len() == 1 {
244                // fast path when the result only has a single row
245                row_wise_results[0].to_array()?
246            } else {
247                ScalarValue::iter_to_array(row_wise_results)?
248            };
249
250            state.update(&out_col, partition_batch_state)?;
251            if self.window_frame.start_bound.is_unbounded() {
252                evaluator.memoize(state)?;
253            }
254        }
255        Ok(())
256    }
257
258    fn get_window_frame(&self) -> &Arc<WindowFrame> {
259        &self.window_frame
260    }
261
262    fn get_reverse_expr(&self) -> Option<Arc<dyn WindowExpr>> {
263        self.expr.reverse_expr().map(|reverse_expr| {
264            Arc::new(StandardWindowExpr::new(
265                reverse_expr,
266                &self.partition_by.clone(),
267                &self
268                    .order_by
269                    .iter()
270                    .map(|e| e.reverse())
271                    .collect::<Vec<_>>(),
272                Arc::new(self.window_frame.reverse()),
273            )) as _
274        })
275    }
276
277    fn uses_bounded_memory(&self) -> bool {
278        if let Ok(evaluator) = self.expr.create_evaluator() {
279            evaluator.supports_bounded_execution()
280                && (!evaluator.uses_window_frame()
281                    || !self.window_frame.end_bound.is_unbounded())
282        } else {
283            false
284        }
285    }
286
287    fn create_window_fn(&self) -> Result<WindowFn> {
288        Ok(WindowFn::Builtin(self.expr.create_evaluator()?))
289    }
290}
291
292/// Adds a new ordering expression into existing ordering equivalence class(es) based on
293/// PARTITION BY information (if it exists).
294pub(crate) fn add_new_ordering_expr_with_partition_by(
295    eqp: &mut EquivalenceProperties,
296    expr: PhysicalSortExpr,
297    partition_by: &[Arc<dyn PhysicalExpr>],
298) -> Result<()> {
299    if partition_by.is_empty() {
300        // In the absence of a PARTITION BY, ordering of `self.expr` is global:
301        eqp.add_ordering([expr]);
302    } else {
303        // If we have a PARTITION BY, standard functions can not introduce
304        // a global ordering unless the existing ordering is compatible
305        // with PARTITION BY expressions. To elaborate, when PARTITION BY
306        // expressions and existing ordering expressions are equal (w.r.t.
307        // set equality), we can prefix the ordering of `self.expr` with
308        // the existing ordering.
309        let (mut ordering, _) = eqp.find_longest_permutation(partition_by)?;
310        if ordering.len() == partition_by.len() {
311            ordering.push(expr);
312            eqp.add_ordering(ordering);
313        }
314    }
315    Ok(())
316}
datafusion_physical_expr/window/standard.rs

datafusion_physical_expr/window/
standard.rs
