## 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/sliding_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::window_expr::{filter_array, AggregateWindowExpr, WindowFn};
    26
    use crate::window::{
    27
        PartitionBatches, PartitionWindowAggStates, PlainAggregateWindowExpr, WindowExpr,
    28
    };
    29
    use crate::{expressions::PhysicalSortExpr, PhysicalExpr};
    30
    31
    use arrow::array::{ArrayRef, BooleanArray};
    32
    use arrow::datatypes::FieldRef;
    33
    use arrow::record_batch::RecordBatch;
    34
    use datafusion_common::{Result, ScalarValue};
    35
    use datafusion_expr::{Accumulator, WindowFrame};
    36
    37
    /// A window expr that takes the form of an aggregate function that
    38
    /// can be incrementally computed over sliding windows.
    39
    ///
    40
    /// See comments on [`WindowExpr`] for more details.
    41
    #[derive(Debug)]
    42
    pub struct SlidingAggregateWindowExpr {
    43
        aggregate: Arc<AggregateFunctionExpr>,
    44
        partition_by: Vec<Arc<dyn PhysicalExpr>>,
    45
        order_by: Vec<PhysicalSortExpr>,
    46
        window_frame: Arc<WindowFrame>,
    47
        filter: Option<Arc<dyn PhysicalExpr>>,
    48
    }
    49
    50
    impl SlidingAggregateWindowExpr {
    51
        /// Create a new (sliding) aggregate window function expression.
    52
    0
        pub fn new(
    53
    0
            aggregate: Arc<AggregateFunctionExpr>,
    54
    0
            partition_by: &[Arc<dyn PhysicalExpr>],
    55
    0
            order_by: &[PhysicalSortExpr],
    56
    0
            window_frame: Arc<WindowFrame>,
    57
    0
            filter: Option<Arc<dyn PhysicalExpr>>,
    58
    0
        ) -> Self {
    59
    0
            Self {
    60
    0
                aggregate,
    61
    0
                partition_by: partition_by.to_vec(),
    62
    0
                order_by: order_by.to_vec(),
    63
    0
                window_frame,
    64
    0
                filter,
    65
    0
            }
    66
    0
        }
    67
    68
        /// Get the [AggregateFunctionExpr] of this object.
    69
    0
        pub fn get_aggregate_expr(&self) -> &AggregateFunctionExpr {
    70
    0
            &self.aggregate
    71
    0
        }
    72
    }
    73
    74
    /// Incrementally update window function using the fact that batch is
    75
    /// pre-sorted given the sort columns and then per partition point.
    76
    ///
    77
    /// Evaluates the peer group (e.g. `SUM` or `MAX` gives the same results
    78
    /// for peers) and concatenate the results.
    79
    impl WindowExpr for SlidingAggregateWindowExpr {
    80
        /// Return a reference to Any that can be used for downcasting
    81
    0
        fn as_any(&self) -> &dyn Any {
    82
    0
            self
    83
    0
        }
    84
    85
    0
        fn field(&self) -> Result<FieldRef> {
    86
    0
            Ok(self.aggregate.field())
    87
    0
        }
    88
    89
    0
        fn name(&self) -> &str {
    90
    0
            self.aggregate.name()
    91
    0
        }
    92
    93
    0
        fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> {
    94
    0
            self.aggregate.expressions()
    95
    0
        }
    96
    97
    0
        fn evaluate(&self, batch: &RecordBatch) -> Result<ArrayRef> {
    98
    0
            self.aggregate_evaluate(batch)
    99
    0
        }
    100
    101
    0
        fn evaluate_stateful(
    102
    0
            &self,
    103
    0
            partition_batches: &PartitionBatches,
    104
    0
            window_agg_state: &mut PartitionWindowAggStates,
    105
    0
        ) -> Result<()> {
    106
    0
            self.aggregate_evaluate_stateful(partition_batches, window_agg_state)
    107
    0
        }
    108
    109
    0
        fn partition_by(&self) -> &[Arc<dyn PhysicalExpr>] {
    110
    0
            &self.partition_by
    111
    0
        }
    112
    113
    0
        fn order_by(&self) -> &[PhysicalSortExpr] {
    114
    0
            &self.order_by
    115
    0
        }
    116
    117
    0
        fn get_window_frame(&self) -> &Arc<WindowFrame> {
    118
    0
            &self.window_frame
    119
    0
        }
    120
    121
    0
        fn get_reverse_expr(&self) -> Option<Arc<dyn WindowExpr>> {
    122
    0
            self.aggregate.reverse_expr().map(|reverse_expr| {
    123
    0
                let reverse_window_frame = self.window_frame.reverse();
    124
    0
                if reverse_window_frame.is_ever_expanding() {
    125
    0
                    Arc::new(PlainAggregateWindowExpr::new(
    126
    0
                        Arc::new(reverse_expr),
    127
    0
                        &self.partition_by.clone(),
    128
    0
                        &self
    129
    0
                            .order_by
    130
    0
                            .iter()
    131
    0
                            .map(|e| e.reverse())
    132
    0
                            .collect::<Vec<_>>(),
    133
    0
                        Arc::new(self.window_frame.reverse()),
    134
    0
                        self.filter.clone(),
    135
    0
                    )) as _
    136
                } else {
    137
    0
                    Arc::new(SlidingAggregateWindowExpr::new(
    138
    0
                        Arc::new(reverse_expr),
    139
    0
                        &self.partition_by.clone(),
    140
    0
                        &self
    141
    0
                            .order_by
    142
    0
                            .iter()
    143
    0
                            .map(|e| e.reverse())
    144
    0
                            .collect::<Vec<_>>(),
    145
    0
                        Arc::new(self.window_frame.reverse()),
    146
    0
                        self.filter.clone(),
    147
    0
                    )) as _
    148
                }
    149
    0
            })
    150
    0
        }
    151
    152
    0
        fn uses_bounded_memory(&self) -> bool {
    153
    0
            !self.window_frame.end_bound.is_unbounded()
    154
    0
        }
    155
    156
    0
        fn with_new_expressions(
    157
    0
            &self,
    158
    0
            args: Vec<Arc<dyn PhysicalExpr>>,
    159
    0
            partition_bys: Vec<Arc<dyn PhysicalExpr>>,
    160
    0
            order_by_exprs: Vec<Arc<dyn PhysicalExpr>>,
    161
    0
        ) -> Option<Arc<dyn WindowExpr>> {
    162
    0
            debug_assert_eq!(self.order_by.len(), order_by_exprs.len());
    163
    164
    0
            let new_order_by = self
    165
    0
                .order_by
    166
    0
                .iter()
    167
    0
                .zip(order_by_exprs)
    168
    0
                .map(|(req, new_expr)| PhysicalSortExpr {
    169
    0
                    expr: new_expr,
    170
    0
                    options: req.options,
    171
    0
                })
    172
    0
                .collect();
    173
    0
            Some(Arc::new(SlidingAggregateWindowExpr {
    174
    0
                aggregate: self
    175
    0
                    .aggregate
    176
    0
                    .with_new_expressions(args, vec![])
    177
    0
                    .map(Arc::new)?,
    178
    0
                partition_by: partition_bys,
    179
    0
                order_by: new_order_by,
    180
    0
                window_frame: Arc::clone(&self.window_frame),
    181
    0
                filter: self.filter.clone(),
    182
            }))
    183
    0
        }
    184
    185
    0
        fn create_window_fn(&self) -> Result<WindowFn> {
    186
    0
            Ok(WindowFn::Aggregate(self.get_accumulator()?))
    187
    0
        }
    188
    }
    189
    190
    impl AggregateWindowExpr for SlidingAggregateWindowExpr {
    191
    0
        fn get_accumulator(&self) -> Result<Box<dyn Accumulator>> {
    192
    0
            self.aggregate.create_sliding_accumulator()
    193
    0
        }
    194
    195
    0
        fn filter_expr(&self) -> Option<&Arc<dyn PhysicalExpr>> {
    196
    0
            self.filter.as_ref()
    197
    0
        }
    198
    199
        /// Given current range and the last range, calculates the accumulator
    200
        /// result for the range of interest.
    201
    0
        fn get_aggregate_result_inside_range(
    202
    0
            &self,
    203
    0
            last_range: &Range<usize>,
    204
    0
            cur_range: &Range<usize>,
    205
    0
            value_slice: &[ArrayRef],
    206
    0
            accumulator: &mut Box<dyn Accumulator>,
    207
    0
            filter_mask: Option<&BooleanArray>,
    208
    0
        ) -> Result<ScalarValue> {
    209
    0
            if cur_range.start == cur_range.end {
    210
    0
                self.aggregate
    211
    0
                    .default_value(self.aggregate.field().data_type())
    212
            } else {
    213
                // Accumulate any new rows that have entered the window:
    214
    0
                let update_bound = cur_range.end - last_range.end;
    215
    0
                if update_bound > 0 {
    216
    0
                    let slice_mask =
    217
    0
                        filter_mask.map(|m| m.slice(last_range.end, update_bound));
    218
    0
                    let update: Vec<ArrayRef> = value_slice
    219
    0
                        .iter()
    220
    0
                        .map(|v| v.slice(last_range.end, update_bound))
    221
    0
                        .map(|arr| match &slice_mask {
    222
    0
                            Some(m) => filter_array(&arr, m),
    223
    0
                            None => Ok(arr),
    224
    0
                        })
    225
    0
                        .collect::<Result<Vec<_>>>()?;
    226
    0
                    accumulator.update_batch(&update)?
    227
    0
                }
    228
    229
                // Remove rows that have now left the window:
    230
    0
                let retract_bound = cur_range.start - last_range.start;
    231
    0
                if retract_bound > 0 {
    232
    0
                    let slice_mask =
    233
    0
                        filter_mask.map(|m| m.slice(last_range.start, retract_bound));
    234
    0
                    let retract: Vec<ArrayRef> = value_slice
    235
    0
                        .iter()
    236
    0
                        .map(|v| v.slice(last_range.start, retract_bound))
    237
    0
                        .map(|arr| match &slice_mask {
    238
    0
                            Some(m) => filter_array(&arr, m),
    239
    0
                            None => Ok(arr),
    240
    0
                        })
    241
    0
                        .collect::<Result<Vec<_>>>()?;
    242
    0
                    accumulator.retract_batch(&retract)?
    243
    0
                }
    244
    0
                accumulator.evaluate()
    245
            }
    246
    0
        }
    247
    248
    0
        fn is_constant_in_partition(&self) -> bool {
    249
    0
            false
    250
    0
        }
    251
    }