## Coverage Report

#### Created: 2025-11-22 06:00

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/datafusion/datafusion/physical-expr/src/window/standard.rs

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