## 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/functions-aggregate-common/src/aggregate/groups_accumulator.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
    //! Utilities for implementing GroupsAccumulator
    19
    //! Adapter that makes [`GroupsAccumulator`] out of [`Accumulator`]
    20
    21
    pub mod accumulate;
    22
    pub mod bool_op;
    23
    pub mod nulls;
    24
    pub mod prim_op;
    25
    26
    use std::mem::{size_of, size_of_val};
    27
    28
    use arrow::array::new_empty_array;
    29
    use arrow::{
    30
        array::{ArrayRef, AsArray, BooleanArray, PrimitiveArray},
    31
        compute,
    32
        compute::take_arrays,
    33
        datatypes::UInt32Type,
    34
    };
    35
    use datafusion_common::{arrow_datafusion_err, DataFusionError, Result, ScalarValue};
    36
    use datafusion_expr_common::accumulator::Accumulator;
    37
    use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator};
    38
    39
    /// An adapter that implements [`GroupsAccumulator`] for any [`Accumulator`]
    40
    ///
    41
    /// While [`Accumulator`] are simpler to implement and can support
    42
    /// more general calculations (like retractable window functions),
    43
    /// they are not as fast as a specialized `GroupsAccumulator`. This
    44
    /// interface bridges the gap so the group by operator only operates
    45
    /// in terms of [`Accumulator`].
    46
    ///
    47
    /// Internally, this adapter creates a new [`Accumulator`] for each group which
    48
    /// stores the state for that group. This both requires an allocation for each
    49
    /// Accumulator, internal indices, as well as whatever internal allocations the
    50
    /// Accumulator itself requires.
    51
    ///
    52
    /// For example, a `MinAccumulator` that computes the minimum string value with
    53
    /// a [`ScalarValue::Utf8`]. That will require at least two allocations per group
    54
    /// (one for the `MinAccumulator` and one for the `ScalarValue::Utf8`).
    55
    ///
    56
    /// ```text
    57
    ///                       ┌─────────────────────────────────┐
    58
    ///                       │MinAccumulator {                 │
    59
    ///                ┌─────▶│ min: ScalarValue::Utf8("A")     │───────┐
    60
    ///                │      │}                                │       │
    61
    ///                │      └─────────────────────────────────┘       └───────▶   "A"
    62
    ///    ┌─────┐     │      ┌─────────────────────────────────┐
    63
    ///    │  0  │─────┘      │MinAccumulator {                 │
    64
    ///    ├─────┤     ┌─────▶│ min: ScalarValue::Utf8("Z")     │───────────────▶   "Z"
    65
    ///    │  1  │─────┘      │}                                │
    66
    ///    └─────┘            └─────────────────────────────────┘                   ...
    67
    ///      ...                 ...
    68
    ///    ┌─────┐            ┌────────────────────────────────┐
    69
    ///    │ N-2 │            │MinAccumulator {                │
    70
    ///    ├─────┤            │  min: ScalarValue::Utf8("A")   │────────────────▶   "A"
    71
    ///    │ N-1 │─────┐      │}                               │
    72
    ///    └─────┘     │      └────────────────────────────────┘
    73
    ///                │      ┌────────────────────────────────┐        ┌───────▶   "Q"
    74
    ///                │      │MinAccumulator {                │        │
    75
    ///                └─────▶│  min: ScalarValue::Utf8("Q")   │────────┘
    76
    ///                       │}                               │
    77
    ///                       └────────────────────────────────┘
    78
    ///
    79
    ///
    80
    ///  Logical group         Current Min/Max value for that group stored
    81
    ///     number             as a ScalarValue which points to an
    82
    ///                        individually allocated String
    83
    /// ```
    84
    ///
    85
    /// # Optimizations
    86
    ///
    87
    /// The adapter minimizes the number of calls to [`Accumulator::update_batch`]
    88
    /// by first collecting the input rows for each group into a contiguous array
    89
    /// using [`compute::take`]
    90
    pub struct GroupsAccumulatorAdapter {
    91
        factory: Box<dyn Fn() -> Result<Box<dyn Accumulator>> + Send>,
    92
    93
        /// state for each group, stored in group_index order
    94
        states: Vec<AccumulatorState>,
    95
    96
        /// Current memory usage, in bytes.
    97
        ///
    98
        /// Note this is incrementally updated with deltas to avoid the
    99
        /// call to size() being a bottleneck. We saw size() being a
    100
        /// bottleneck in earlier implementations when there were many
    101
        /// distinct groups.
    102
        allocation_bytes: usize,
    103
    }
    104
    105
    struct AccumulatorState {
    106
        /// [`Accumulator`] that stores the per-group state
    107
        accumulator: Box<dyn Accumulator>,
    108
    109
        /// scratch space: indexes in the input array that will be fed to
    110
        /// this accumulator. Stores indexes as `u32` to match the arrow
    111
        /// `take` kernel input.
    112
        indices: Vec<u32>,
    113
    }
    114
    115
    impl AccumulatorState {
    116
    0
        fn new(accumulator: Box<dyn Accumulator>) -> Self {
    117
    0
            Self {
    118
    0
                accumulator,
    119
    0
                indices: vec![],
    120
    0
            }
    121
    0
        }
    122
    123
        /// Returns the amount of memory taken by this structure and its accumulator
    124
    0
        fn size(&self) -> usize {
    125
    0
            self.accumulator.size() + size_of_val(self) + self.indices.allocated_size()
    126
    0
        }
    127
    }
    128
    129
    impl GroupsAccumulatorAdapter {
    130
        /// Create a new adapter that will create a new [`Accumulator`]
    131
        /// for each group, using the specified factory function
    132
    0
        pub fn new<F>(factory: F) -> Self
    133
    0
        where
    134
    0
            F: Fn() -> Result<Box<dyn Accumulator>> + Send + 'static,
    135
        {
    136
    0
            Self {
    137
    0
                factory: Box::new(factory),
    138
    0
                states: vec![],
    139
    0
                allocation_bytes: 0,
    140
    0
            }
    141
    0
        }
    142
    143
        /// Ensure that self.accumulators has total_num_groups
    144
    0
        fn make_accumulators_if_needed(&mut self, total_num_groups: usize) -> Result<()> {
    145
            // can't shrink
    146
    0
            assert!(total_num_groups >= self.states.len());
    147
    0
            let vec_size_pre = self.states.allocated_size();
    148
    149
            // instantiate new accumulators
    150
    0
            let new_accumulators = total_num_groups - self.states.len();
    151
    0
            for _ in 0..new_accumulators {
    152
    0
                let accumulator = (self.factory)()?;
    153
    0
                let state = AccumulatorState::new(accumulator);
    154
    0
                self.add_allocation(state.size());
    155
    0
                self.states.push(state);
    156
            }
    157
    158
    0
            self.adjust_allocation(vec_size_pre, self.states.allocated_size());
    159
    0
            Ok(())
    160
    0
        }
    161
    162
        /// invokes f(accumulator, values) for each group that has values
    163
        /// in group_indices.
    164
        ///
    165
        /// This function first reorders the input and filter so that
    166
        /// values for each group_index are contiguous and then invokes f
    167
        /// on the contiguous ranges, to minimize per-row overhead
    168
        ///
    169
        /// ```text
    170
        /// ┌─────────┐   ┌─────────┐   ┌ ─ ─ ─ ─ ┐                       ┌─────────┐   ┌ ─ ─ ─ ─ ┐
    171
        /// │ ┌─────┐ │   │ ┌─────┐ │     ┌─────┐              ┏━━━━━┓    │ ┌─────┐ │     ┌─────┐
    172
        /// │ │  2  │ │   │ │ 200 │ │   │ │  t  │ │            ┃  0  ┃    │ │ 200 │ │   │ │  t  │ │
    173
        /// │ ├─────┤ │   │ ├─────┤ │     ├─────┤              ┣━━━━━┫    │ ├─────┤ │     ├─────┤
    174
        /// │ │  2  │ │   │ │ 100 │ │   │ │  f  │ │            ┃  0  ┃    │ │ 300 │ │   │ │  t  │ │
    175
        /// │ ├─────┤ │   │ ├─────┤ │     ├─────┤              ┣━━━━━┫    │ ├─────┤ │     ├─────┤
    176
        /// │ │  0  │ │   │ │ 200 │ │   │ │  t  │ │            ┃  1  ┃    │ │ 200 │ │   │ │NULL │ │
    177
        /// │ ├─────┤ │   │ ├─────┤ │     ├─────┤   ────────▶  ┣━━━━━┫    │ ├─────┤ │     ├─────┤
    178
        /// │ │  1  │ │   │ │ 200 │ │   │ │NULL │ │            ┃  2  ┃    │ │ 200 │ │   │ │  t  │ │
    179
        /// │ ├─────┤ │   │ ├─────┤ │     ├─────┤              ┣━━━━━┫    │ ├─────┤ │     ├─────┤
    180
        /// │ │  0  │ │   │ │ 300 │ │   │ │  t  │ │            ┃  2  ┃    │ │ 100 │ │   │ │  f  │ │
    181
        /// │ └─────┘ │   │ └─────┘ │     └─────┘              ┗━━━━━┛    │ └─────┘ │     └─────┘
    182
        /// └─────────┘   └─────────┘   └ ─ ─ ─ ─ ┘                       └─────────┘   └ ─ ─ ─ ─ ┘
    183
        ///
    184
        /// logical group   values      opt_filter           logical group  values       opt_filter
    185
        /// ```
    186
    0
        fn invoke_per_accumulator<F>(
    187
    0
            &mut self,
    188
    0
            values: &[ArrayRef],
    189
    0
            group_indices: &[usize],
    190
    0
            opt_filter: Option<&BooleanArray>,
    191
    0
            total_num_groups: usize,
    192
    0
            f: F,
    193
    0
        ) -> Result<()>
    194
    0
        where
    195
    0
            F: Fn(&mut dyn Accumulator, &[ArrayRef]) -> Result<()>,
    196
        {
    197
    0
            self.make_accumulators_if_needed(total_num_groups)?;
    198
    199
    0
            assert_eq!(values[0].len(), group_indices.len());
    200
    201
            // figure out which input rows correspond to which groups.
    202
            // Note that self.state.indices starts empty for all groups
    203
            // (it is cleared out below)
    204
    0
            for (idx, group_index) in group_indices.iter().enumerate() {
    205
    0
                self.states[*group_index].indices.push(idx as u32);
    206
    0
            }
    207
    208
            // groups_with_rows holds a list of group indexes that have
    209
            // any rows that need to be accumulated, stored in order of
    210
            // group_index
    211
    212
    0
            let mut groups_with_rows = vec![];
    213
    214
            // batch_indices holds indices into values, each group is contiguous
    215
    0
            let mut batch_indices = vec![];
    216
    217
            // offsets[i] is index into batch_indices where the rows for
    218
            // group_index i starts
    219
    0
            let mut offsets = vec![0];
    220
    221
    0
            let mut offset_so_far = 0;
    222
    0
            for (group_index, state) in self.states.iter_mut().enumerate() {
    223
    0
                let indices = &state.indices;
    224
    0
                if indices.is_empty() {
    225
    0
                    continue;
    226
    0
                }
    227
    228
    0
                groups_with_rows.push(group_index);
    229
    0
                batch_indices.extend_from_slice(indices);
    230
    0
                offset_so_far += indices.len();
    231
    0
                offsets.push(offset_so_far);
    232
            }
    233
    0
            let batch_indices = batch_indices.into();
    234
    235
            // reorder the values and opt_filter by batch_indices so that
    236
            // all values for each group are contiguous, then invoke the
    237
            // accumulator once per group with values
    238
    0
            let values = take_arrays(values, &batch_indices, None)?;
    239
    0
            let opt_filter = get_filter_at_indices(opt_filter, &batch_indices)?;
    240
    241
            // invoke each accumulator with the appropriate rows, first
    242
            // pulling the input arguments for this group into their own
    243
            // RecordBatch(es)
    244
    0
            let iter = groups_with_rows.iter().zip(offsets.windows(2));
    245
    246
    0
            let mut sizes_pre = 0;
    247
    0
            let mut sizes_post = 0;
    248
    0
            for (&group_idx, offsets) in iter {
    249
    0
                let state = &mut self.states[group_idx];
    250
    0
                sizes_pre += state.size();
    251
    252
    0
                let values_to_accumulate = slice_and_maybe_filter(
    253
    0
                    &values,
    254
    0
                    opt_filter.as_ref().map(|f| f.as_boolean()),
    255
    0
                    offsets,
    256
    0
                )?;
    257
    0
                f(state.accumulator.as_mut(), &values_to_accumulate)?;
    258
    259
                // clear out the state so they are empty for next
    260
                // iteration
    261
    0
                state.indices.clear();
    262
    0
                sizes_post += state.size();
    263
            }
    264
    265
    0
            self.adjust_allocation(sizes_pre, sizes_post);
    266
    0
            Ok(())
    267
    0
        }
    268
    269
        /// Increment the allocation by `n`
    270
        ///
    271
        /// See [`Self::allocation_bytes`] for rationale.
    272
    0
        fn add_allocation(&mut self, size: usize) {
    273
    0
            self.allocation_bytes += size;
    274
    0
        }
    275
    276
        /// Decrease the allocation by `n`
    277
        ///
    278
        /// See [`Self::allocation_bytes`] for rationale.
    279
    0
        fn free_allocation(&mut self, size: usize) {
    280
            // use saturating sub to avoid errors if the accumulators
    281
            // report erroneous sizes
    282
    0
            self.allocation_bytes = self.allocation_bytes.saturating_sub(size)
    283
    0
        }
    284
    285
        /// Adjusts the allocation for something that started with
    286
        /// start_size and now has new_size avoiding overflow
    287
        ///
    288
        /// See [`Self::allocation_bytes`] for rationale.
    289
    0
        fn adjust_allocation(&mut self, old_size: usize, new_size: usize) {
    290
    0
            if new_size > old_size {
    291
    0
                self.add_allocation(new_size - old_size)
    292
            } else {
    293
    0
                self.free_allocation(old_size - new_size)
    294
            }
    295
    0
        }
    296
    }
    297
    298
    impl GroupsAccumulator for GroupsAccumulatorAdapter {
    299
    0
        fn update_batch(
    300
    0
            &mut self,
    301
    0
            values: &[ArrayRef],
    302
    0
            group_indices: &[usize],
    303
    0
            opt_filter: Option<&BooleanArray>,
    304
    0
            total_num_groups: usize,
    305
    0
        ) -> Result<()> {
    306
    0
            self.invoke_per_accumulator(
    307
    0
                values,
    308
    0
                group_indices,
    309
    0
                opt_filter,
    310
    0
                total_num_groups,
    311
    0
                |accumulator, values_to_accumulate| {
    312
    0
                    accumulator.update_batch(values_to_accumulate)
    313
    0
                },
    314
    0
            )?;
    315
    0
            Ok(())
    316
    0
        }
    317
    318
    0
        fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
    319
    0
            let vec_size_pre = self.states.allocated_size();
    320
    321
    0
            let states = emit_to.take_needed(&mut self.states);
    322
    323
    0
            let results: Vec<ScalarValue> = states
    324
    0
                .into_iter()
    325
    0
                .map(|mut state| {
    326
    0
                    self.free_allocation(state.size());
    327
    0
                    state.accumulator.evaluate()
    328
    0
                })
    329
    0
                .collect::<Result<_>>()?;
    330
    331
    0
            let result = ScalarValue::iter_to_array(results);
    332
    333
    0
            self.adjust_allocation(vec_size_pre, self.states.allocated_size());
    334
    335
    0
            result
    336
    0
        }
    337
    338
        // filtered_null_mask(opt_filter, &values);
    339
    0
        fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
    340
    0
            let vec_size_pre = self.states.allocated_size();
    341
    0
            let states = emit_to.take_needed(&mut self.states);
    342
    343
            // each accumulator produces a potential vector of values
    344
            // which we need to form into columns
    345
    0
            let mut results: Vec<Vec<ScalarValue>> = vec![];
    346
    347
    0
            for mut state in states {
    348
    0
                self.free_allocation(state.size());
    349
    0
                let accumulator_state = state.accumulator.state()?;
    350
    0
                results.resize_with(accumulator_state.len(), Vec::new);
    351
    0
                for (idx, state_val) in accumulator_state.into_iter().enumerate() {
    352
    0
                    results[idx].push(state_val);
    353
    0
                }
    354
            }
    355
    356
            // create an array for each intermediate column
    357
    0
            let arrays = results
    358
    0
                .into_iter()
    359
    0
                .map(ScalarValue::iter_to_array)
    360
    0
                .collect::<Result<Vec<_>>>()?;
    361
    362
            // double check each array has the same length (aka the
    363
            // accumulator was implemented correctly
    364
    0
            if let Some(first_col) = arrays.first() {
    365
    0
                for arr in &arrays {
    366
    0
                    assert_eq!(arr.len(), first_col.len())
    367
                }
    368
    0
            }
    369
    0
            self.adjust_allocation(vec_size_pre, self.states.allocated_size());
    370
    371
    0
            Ok(arrays)
    372
    0
        }
    373
    374
    0
        fn merge_batch(
    375
    0
            &mut self,
    376
    0
            values: &[ArrayRef],
    377
    0
            group_indices: &[usize],
    378
    0
            opt_filter: Option<&BooleanArray>,
    379
    0
            total_num_groups: usize,
    380
    0
        ) -> Result<()> {
    381
    0
            self.invoke_per_accumulator(
    382
    0
                values,
    383
    0
                group_indices,
    384
    0
                opt_filter,
    385
    0
                total_num_groups,
    386
    0
                |accumulator, values_to_accumulate| {
    387
    0
                    accumulator.merge_batch(values_to_accumulate)?;
    388
    0
                    Ok(())
    389
    0
                },
    390
    0
            )?;
    391
    0
            Ok(())
    392
    0
        }
    393
    394
    0
        fn size(&self) -> usize {
    395
    0
            self.allocation_bytes
    396
    0
        }
    397
    398
    0
        fn convert_to_state(
    399
    0
            &self,
    400
    0
            values: &[ArrayRef],
    401
    0
            opt_filter: Option<&BooleanArray>,
    402
    0
        ) -> Result<Vec<ArrayRef>> {
    403
    0
            let num_rows = values[0].len();
    404
    405
            // If there are no rows, return empty arrays
    406
    0
            if num_rows == 0 {
    407
                // create empty accumulator to get the state types
    408
    0
                let empty_state = (self.factory)()?.state()?;
    409
    0
                let empty_arrays = empty_state
    410
    0
                    .into_iter()
    411
    0
                    .map(|state_val| new_empty_array(&state_val.data_type()))
    412
    0
                    .collect::<Vec<_>>();
    413
    414
    0
                return Ok(empty_arrays);
    415
    0
            }
    416
    417
            // Each row has its respective group
    418
    0
            let mut results = vec![];
    419
    0
            for row_idx in 0..num_rows {
    420
                // Create the empty accumulator for converting
    421
    0
                let mut converted_accumulator = (self.factory)()?;
    422
    423
                // Convert row to states
    424
    0
                let values_to_accumulate =
    425
    0
                    slice_and_maybe_filter(values, opt_filter, &[row_idx, row_idx + 1])?;
    426
    0
                converted_accumulator.update_batch(&values_to_accumulate)?;
    427
    0
                let states = converted_accumulator.state()?;
    428
    429
                // Resize results to have enough columns according to the converted states
    430
    0
                results.resize_with(states.len(), || Vec::with_capacity(num_rows));
    431
    432
                // Add the states to results
    433
    0
                for (idx, state_val) in states.into_iter().enumerate() {
    434
    0
                    results[idx].push(state_val);
    435
    0
                }
    436
            }
    437
    438
    0
            let arrays = results
    439
    0
                .into_iter()
    440
    0
                .map(ScalarValue::iter_to_array)
    441
    0
                .collect::<Result<Vec<_>>>()?;
    442
    443
    0
            Ok(arrays)
    444
    0
        }
    445
    446
    0
        fn supports_convert_to_state(&self) -> bool {
    447
    0
            true
    448
    0
        }
    449
    }
    450
    451
    /// Extension trait for [`Vec`] to account for allocations.
    452
    pub trait VecAllocExt {
    453
        /// Item type.
    454
        type T;
    455
        /// Return the amount of memory allocated by this Vec (not
    456
        /// recursively counting any heap allocations contained within the
    457
        /// structure). Does not include the size of `self`
    458
        fn allocated_size(&self) -> usize;
    459
    }
    460
    461
    impl<T> VecAllocExt for Vec<T> {
    462
        type T = T;
    463
    0
        fn allocated_size(&self) -> usize {
    464
    0
            size_of::<T>() * self.capacity()
    465
    0
        }
    466
    }
    467
    468
    0
    fn get_filter_at_indices(
    469
    0
        opt_filter: Option<&BooleanArray>,
    470
    0
        indices: &PrimitiveArray<UInt32Type>,
    471
    0
    ) -> Result<Option<ArrayRef>> {
    472
    0
        opt_filter
    473
    0
            .map(|filter| {
    474
    0
                compute::take(
    475
    0
                    &filter, indices, None, // None: no index check
    476
                )
    477
    0
            })
    478
    0
            .transpose()
    479
    0
            .map_err(|e| arrow_datafusion_err!(e))
    480
    0
    }
    481
    482
    // Copied from physical-plan
    483
    0
    pub(crate) fn slice_and_maybe_filter(
    484
    0
        aggr_array: &[ArrayRef],
    485
    0
        filter_opt: Option<&BooleanArray>,
    486
    0
        offsets: &[usize],
    487
    0
    ) -> Result<Vec<ArrayRef>> {
    488
    0
        let (offset, length) = (offsets[0], offsets[1] - offsets[0]);
    489
    0
        let sliced_arrays: Vec<ArrayRef> = aggr_array
    490
    0
            .iter()
    491
    0
            .map(|array| array.slice(offset, length))
    492
    0
            .collect();
    493
    494
    0
        if let Some(f) = filter_opt {
    495
    0
            let filter = f.slice(offset, length);
    496
    497
    0
            sliced_arrays
    498
    0
                .iter()
    499
    0
                .map(|array| {
    500
    0
                    compute::filter(&array, &filter).map_err(|e| arrow_datafusion_err!(e))
    501
    0
                })
    502
    0
                .collect()
    503
        } else {
    504
    0
            Ok(sliced_arrays)
    505
        }
    506
    0
    }