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