## 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/accumulate.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
    //! [`GroupsAccumulator`] helpers: [`NullState`] and [`accumulate_indices`]
    19
    //!
    20
    //! [`GroupsAccumulator`]: datafusion_expr_common::groups_accumulator::GroupsAccumulator
    21
    22
    use arrow::array::{Array, BooleanArray, BooleanBufferBuilder, PrimitiveArray};
    23
    use arrow::buffer::{BooleanBuffer, NullBuffer};
    24
    use arrow::datatypes::ArrowPrimitiveType;
    25
    26
    use datafusion_expr_common::groups_accumulator::EmitTo;
    27
    /// Track the accumulator null state per row: if any values for that
    28
    /// group were null and if any values have been seen at all for that group.
    29
    ///
    30
    /// This is part of the inner loop for many [`GroupsAccumulator`]s,
    31
    /// and thus the performance is critical and so there are multiple
    32
    /// specialized implementations, invoked depending on the specific
    33
    /// combinations of the input.
    34
    ///
    35
    /// Typically there are 4 potential combinations of inputs must be
    36
    /// special cased for performance:
    37
    ///
    38
    /// * With / Without filter
    39
    /// * With / Without nulls in the input
    40
    ///
    41
    /// If the input has nulls, then the accumulator must potentially
    42
    /// handle each input null value specially (e.g. for `SUM` to mark the
    43
    /// corresponding sum as null)
    44
    ///
    45
    /// If there are filters present, `NullState` tracks if it has seen
    46
    /// *any* value for that group (as some values may be filtered
    47
    /// out). Without a filter, the accumulator is only passed groups that
    48
    /// had at least one value to accumulate so they do not need to track
    49
    /// if they have seen values for a particular group.
    50
    ///
    51
    /// [`GroupsAccumulator`]: datafusion_expr_common::groups_accumulator::GroupsAccumulator
    52
    #[derive(Debug)]
    53
    pub struct NullState {
    54
        /// Have we seen any non-filtered input values for `group_index`?
    55
        ///
    56
        /// If `seen_values[i]` is true, have seen at least one non null
    57
        /// value for group `i`
    58
        ///
    59
        /// If `seen_values[i]` is false, have not seen any values that
    60
        /// pass the filter yet for group `i`
    61
        seen_values: BooleanBufferBuilder,
    62
    }
    63
    64
    impl Default for NullState {
    65
        fn default() -> Self {
    66
            Self::new()
    67
        }
    68
    }
    69
    70
    impl NullState {
    71
        pub fn new() -> Self {
    72
            Self {
    73
                seen_values: BooleanBufferBuilder::new(0),
    74
            }
    75
        }
    76
    77
        /// return the size of all buffers allocated by this null state, not including self
    78
        pub fn size(&self) -> usize {
    79
            // capacity is in bits, so convert to bytes
    80
            self.seen_values.capacity() / 8
    81
        }
    82
    83
        /// Invokes `value_fn(group_index, value)` for each non null, non
    84
        /// filtered value of `value`, while tracking which groups have
    85
        /// seen null inputs and which groups have seen any inputs if necessary
    86
        //
    87
        /// # Arguments:
    88
        ///
    89
        /// * `values`: the input arguments to the accumulator
    90
        /// * `group_indices`:  To which groups do the rows in `values` belong, (aka group_index)
    91
        /// * `opt_filter`: if present, only rows for which is Some(true) are included
    92
        /// * `value_fn`: function invoked for  (group_index, value) where value is non null
    93
        ///
    94
        /// See [`accumulate`], for more details on how value_fn is called
    95
        ///
    96
        /// When value_fn is called it also sets
    97
        ///
    98
        /// 1. `self.seen_values[group_index]` to true for all rows that had a non null value
    99
    0
        pub fn accumulate<T, F>(
    100
    0
            &mut self,
    101
    0
            group_indices: &[usize],
    102
    0
            values: &PrimitiveArray<T>,
    103
    0
            opt_filter: Option<&BooleanArray>,
    104
    0
            total_num_groups: usize,
    105
    0
            mut value_fn: F,
    106
    0
        ) where
    107
    0
            T: ArrowPrimitiveType + Send,
    108
    0
            F: FnMut(usize, T::Native) + Send,
    109
        {
    110
            // ensure the seen_values is big enough (start everything at
    111
            // "not seen" valid)
    112
    0
            let seen_values =
    113
    0
                initialize_builder(&mut self.seen_values, total_num_groups, false);
    114
    0
            accumulate(group_indices, values, opt_filter, |group_index, value| {
    115
    0
                seen_values.set_bit(group_index, true);
    116
    0
                value_fn(group_index, value);
    117
    0
            });
    118
    0
        }
    119
    120
        /// Invokes `value_fn(group_index, value)` for each non null, non
    121
        /// filtered value in `values`, while tracking which groups have
    122
        /// seen null inputs and which groups have seen any inputs, for
    123
        /// [`BooleanArray`]s.
    124
        ///
    125
        /// Since `BooleanArray` is not a [`PrimitiveArray`] it must be
    126
        /// handled specially.
    127
        ///
    128
        /// See [`Self::accumulate`], which handles `PrimitiveArray`s, for
    129
        /// more details on other arguments.
    130
    0
        pub fn accumulate_boolean<F>(
    131
    0
            &mut self,
    132
    0
            group_indices: &[usize],
    133
    0
            values: &BooleanArray,
    134
    0
            opt_filter: Option<&BooleanArray>,
    135
    0
            total_num_groups: usize,
    136
    0
            mut value_fn: F,
    137
    0
        ) where
    138
    0
            F: FnMut(usize, bool) + Send,
    139
        {
    140
    0
            let data = values.values();
    141
    0
            assert_eq!(data.len(), group_indices.len());
    142
    143
            // ensure the seen_values is big enough (start everything at
    144
            // "not seen" valid)
    145
    0
            let seen_values =
    146
    0
                initialize_builder(&mut self.seen_values, total_num_groups, false);
    147
    148
            // These could be made more performant by iterating in chunks of 64 bits at a time
    149
    0
            match (values.null_count() > 0, opt_filter) {
    150
                // no nulls, no filter,
    151
                (false, None) => {
    152
                    // if we have previously seen nulls, ensure the null
    153
                    // buffer is big enough (start everything at valid)
    154
    0
                    group_indices.iter().zip(data.iter()).for_each(
    155
    0
                        |(&group_index, new_value)| {
    156
    0
                            seen_values.set_bit(group_index, true);
    157
    0
                            value_fn(group_index, new_value)
    158
    0
                        },
    159
                    )
    160
                }
    161
                // nulls, no filter
    162
                (true, None) => {
    163
    0
                    let nulls = values.nulls().unwrap();
    164
    0
                    group_indices
    165
    0
                        .iter()
    166
    0
                        .zip(data.iter())
    167
    0
                        .zip(nulls.iter())
    168
    0
                        .for_each(|((&group_index, new_value), is_valid)| {
    169
    0
                            if is_valid {
    170
    0
                                seen_values.set_bit(group_index, true);
    171
    0
                                value_fn(group_index, new_value);
    172
    0
                            }
    173
    0
                        })
    174
                }
    175
                // no nulls, but a filter
    176
    0
                (false, Some(filter)) => {
    177
    0
                    assert_eq!(filter.len(), group_indices.len());
    178
    179
    0
                    group_indices
    180
    0
                        .iter()
    181
    0
                        .zip(data.iter())
    182
    0
                        .zip(filter.iter())
    183
    0
                        .for_each(|((&group_index, new_value), filter_value)| {
    184
    0
                            if let Some(true) = filter_value {
    185
    0
                                seen_values.set_bit(group_index, true);
    186
    0
                                value_fn(group_index, new_value);
    187
    0
                            }
    188
    0
                        })
    189
                }
    190
                // both null values and filters
    191
    0
                (true, Some(filter)) => {
    192
    0
                    assert_eq!(filter.len(), group_indices.len());
    193
    0
                    filter
    194
    0
                        .iter()
    195
    0
                        .zip(group_indices.iter())
    196
    0
                        .zip(values.iter())
    197
    0
                        .for_each(|((filter_value, &group_index), new_value)| {
    198
    0
                            if let Some(true) = filter_value {
    199
    0
                                if let Some(new_value) = new_value {
    200
    0
                                    seen_values.set_bit(group_index, true);
    201
    0
                                    value_fn(group_index, new_value)
    202
    0
                                }
    203
    0
                            }
    204
    0
                        })
    205
                }
    206
            }
    207
    0
        }
    208
    209
        /// Creates the a [`NullBuffer`] representing which group_indices
    210
        /// should have null values (because they never saw any values)
    211
        /// for the `emit_to` rows.
    212
        ///
    213
        /// resets the internal state appropriately
    214
        pub fn build(&mut self, emit_to: EmitTo) -> NullBuffer {
    215
            let nulls: BooleanBuffer = self.seen_values.finish();
    216
    217
            let nulls = match emit_to {
    218
                EmitTo::All => nulls,
    219
                EmitTo::First(n) => {
    220
                    // split off the first N values in seen_values
    221
                    let first_n_null: BooleanBuffer = nulls.slice(0, n);
    222
                    // reset the existing seen buffer
    223
                    self.seen_values
    224
                        .append_buffer(&nulls.slice(n, nulls.len() - n));
    225
                    first_n_null
    226
                }
    227
            };
    228
            NullBuffer::new(nulls)
    229
        }
    230
    }
    231
    232
    /// Invokes `value_fn(group_index, value)` for each non null, non
    233
    /// filtered value of `value`,
    234
    ///
    235
    /// # Arguments:
    236
    ///
    237
    /// * `group_indices`:  To which groups do the rows in `values` belong, (aka group_index)
    238
    /// * `values`: the input arguments to the accumulator
    239
    /// * `opt_filter`: if present, only rows for which is Some(true) are included
    240
    /// * `value_fn`: function invoked for  (group_index, value) where value is non null
    241
    ///
    242
    /// # Example
    243
    ///
    244
    /// ```text
    245
    ///  ┌─────────┐   ┌─────────┐   ┌ ─ ─ ─ ─ ┐
    246
    ///  │ ┌─────┐ │   │ ┌─────┐ │     ┌─────┐
    247
    ///  │ │  2  │ │   │ │ 200 │ │   │ │  t  │ │
    248
    ///  │ ├─────┤ │   │ ├─────┤ │     ├─────┤
    249
    ///  │ │  2  │ │   │ │ 100 │ │   │ │  f  │ │
    250
    ///  │ ├─────┤ │   │ ├─────┤ │     ├─────┤
    251
    ///  │ │  0  │ │   │ │ 200 │ │   │ │  t  │ │
    252
    ///  │ ├─────┤ │   │ ├─────┤ │     ├─────┤
    253
    ///  │ │  1  │ │   │ │ 200 │ │   │ │NULL │ │
    254
    ///  │ ├─────┤ │   │ ├─────┤ │     ├─────┤
    255
    ///  │ │  0  │ │   │ │ 300 │ │   │ │  t  │ │
    256
    ///  │ └─────┘ │   │ └─────┘ │     └─────┘
    257
    ///  └─────────┘   └─────────┘   └ ─ ─ ─ ─ ┘
    258
    ///
    259
    /// group_indices   values        opt_filter
    260
    /// ```
    261
    ///
    262
    /// In the example above, `value_fn` is invoked for each (group_index,
    263
    /// value) pair where `opt_filter[i]` is true and values is non null
    264
    ///
    265
    /// ```text
    266
    /// value_fn(2, 200)
    267
    /// value_fn(0, 200)
    268
    /// value_fn(0, 300)
    269
    /// ```
    270
    0
    pub fn accumulate<T, F>(
    271
    0
        group_indices: &[usize],
    272
    0
        values: &PrimitiveArray<T>,
    273
    0
        opt_filter: Option<&BooleanArray>,
    274
    0
        mut value_fn: F,
    275
    0
    ) where
    276
    0
        T: ArrowPrimitiveType + Send,
    277
    0
        F: FnMut(usize, T::Native) + Send,
    278
    {
    279
    0
        let data: &[T::Native] = values.values();
    280
    0
        assert_eq!(data.len(), group_indices.len());
    281
    282
    0
        match (values.null_count() > 0, opt_filter) {
    283
            // no nulls, no filter,
    284
            (false, None) => {
    285
    0
                let iter = group_indices.iter().zip(data.iter());
    286
    0
                for (&group_index, &new_value) in iter {
    287
    0
                    value_fn(group_index, new_value);
    288
    0
                }
    289
            }
    290
            // nulls, no filter
    291
            (true, None) => {
    292
    0
                let nulls = values.nulls().unwrap();
    293
                // This is based on (ahem, COPY/PASTE) arrow::compute::aggregate::sum
    294
                // iterate over in chunks of 64 bits for more efficient null checking
    295
    0
                let group_indices_chunks = group_indices.chunks_exact(64);
    296
    0
                let data_chunks = data.chunks_exact(64);
    297
    0
                let bit_chunks = nulls.inner().bit_chunks();
    298
    299
    0
                let group_indices_remainder = group_indices_chunks.remainder();
    300
    0
                let data_remainder = data_chunks.remainder();
    301
    302
    0
                group_indices_chunks
    303
    0
                    .zip(data_chunks)
    304
    0
                    .zip(bit_chunks.iter())
    305
    0
                    .for_each(|((group_index_chunk, data_chunk), mask)| {
    306
                        // index_mask has value 1 << i in the loop
    307
    0
                        let mut index_mask = 1;
    308
    0
                        group_index_chunk.iter().zip(data_chunk.iter()).for_each(
    309
    0
                            |(&group_index, &new_value)| {
    310
                                // valid bit was set, real value
    311
    0
                                let is_valid = (mask & index_mask) != 0;
    312
    0
                                if is_valid {
    313
    0
                                    value_fn(group_index, new_value);
    314
    0
                                }
    315
    0
                                index_mask <<= 1;
    316
    0
                            },
    317
                        )
    318
    0
                    });
    319
    320
                // handle any remaining bits (after the initial 64)
    321
    0
                let remainder_bits = bit_chunks.remainder_bits();
    322
    0
                group_indices_remainder
    323
    0
                    .iter()
    324
    0
                    .zip(data_remainder.iter())
    325
    0
                    .enumerate()
    326
    0
                    .for_each(|(i, (&group_index, &new_value))| {
    327
    0
                        let is_valid = remainder_bits & (1 << i) != 0;
    328
    0
                        if is_valid {
    329
    0
                            value_fn(group_index, new_value);
    330
    0
                        }
    331
    0
                    });
    332
            }
    333
            // no nulls, but a filter
    334
    0
            (false, Some(filter)) => {
    335
    0
                assert_eq!(filter.len(), group_indices.len());
    336
                // The performance with a filter could be improved by
    337
                // iterating over the filter in chunks, rather than a single
    338
                // iterator. TODO file a ticket
    339
    0
                group_indices
    340
    0
                    .iter()
    341
    0
                    .zip(data.iter())
    342
    0
                    .zip(filter.iter())
    343
    0
                    .for_each(|((&group_index, &new_value), filter_value)| {
    344
    0
                        if let Some(true) = filter_value {
    345
    0
                            value_fn(group_index, new_value);
    346
    0
                        }
    347
    0
                    })
    348
            }
    349
            // both null values and filters
    350
    0
            (true, Some(filter)) => {
    351
    0
                assert_eq!(filter.len(), group_indices.len());
    352
                // The performance with a filter could be improved by
    353
                // iterating over the filter in chunks, rather than using
    354
                // iterators. TODO file a ticket
    355
    0
                filter
    356
    0
                    .iter()
    357
    0
                    .zip(group_indices.iter())
    358
    0
                    .zip(values.iter())
    359
    0
                    .for_each(|((filter_value, &group_index), new_value)| {
    360
    0
                        if let Some(true) = filter_value {
    361
    0
                            if let Some(new_value) = new_value {
    362
    0
                                value_fn(group_index, new_value)
    363
    0
                            }
    364
    0
                        }
    365
    0
                    })
    366
            }
    367
        }
    368
    0
    }
    369
    370
    /// Accumulates with multiple accumulate(value) columns. (e.g. `corr(c1, c2)`)
    371
    ///
    372
    /// This method assumes that for any input record index, if any of the value column
    373
    /// is null, or it's filtered out by `opt_filter`, then the record would be ignored.
    374
    /// (won't be accumulated by `value_fn`)
    375
    ///
    376
    /// # Arguments
    377
    ///
    378
    /// * `group_indices` - To which groups do the rows in `value_columns` belong
    379
    /// * `value_columns` - The input arrays to accumulate
    380
    /// * `opt_filter` - Optional filter array. If present, only rows where filter is `Some(true)` are included
    381
    /// * `value_fn` - Callback function for each valid row, with parameters:
    382
    ///     * `group_idx`: The group index for the current row
    383
    ///     * `batch_idx`: The index of the current row in the input arrays
    384
    ///     * `columns`: Reference to all input arrays for accessing values
    385
    0
    pub fn accumulate_multiple<T, F>(
    386
    0
        group_indices: &[usize],
    387
    0
        value_columns: &[&PrimitiveArray<T>],
    388
    0
        opt_filter: Option<&BooleanArray>,
    389
    0
        mut value_fn: F,
    390
    0
    ) where
    391
    0
        T: ArrowPrimitiveType + Send,
    392
    0
        F: FnMut(usize, usize, &[&PrimitiveArray<T>]) + Send,
    393
    {
    394
        // Calculate `valid_indices` to accumulate, non-valid indices are ignored.
    395
        // `valid_indices` is a bit mask corresponding to the `group_indices`. An index
    396
        // is considered valid if:
    397
        // 1. All columns are non-null at this index.
    398
        // 2. Not filtered out by `opt_filter`
    399
    400
        // Take AND from all null buffers of `value_columns`.
    401
    0
        let combined_nulls = value_columns
    402
    0
            .iter()
    403
    0
            .map(|arr| arr.logical_nulls())
    404
    0
            .fold(None, |acc, nulls| {
    405
    0
                NullBuffer::union(acc.as_ref(), nulls.as_ref())
    406
    0
            });
    407
    408
        // Take AND from previous combined nulls and `opt_filter`.
    409
    0
        let valid_indices = match (combined_nulls, opt_filter) {
    410
    0
            (None, None) => None,
    411
    0
            (None, Some(filter)) => Some(filter.clone()),
    412
    0
            (Some(nulls), None) => Some(BooleanArray::new(nulls.inner().clone(), None)),
    413
    0
            (Some(nulls), Some(filter)) => {
    414
    0
                let combined = nulls.inner() & filter.values();
    415
    0
                Some(BooleanArray::new(combined, None))
    416
            }
    417
        };
    418
    419
    0
        for col in value_columns.iter() {
    420
    0
            debug_assert_eq!(col.len(), group_indices.len());
    421
        }
    422
    423
    0
        match valid_indices {
    424
            None => {
    425
    0
                for (batch_idx, &group_idx) in group_indices.iter().enumerate() {
    426
    0
                    value_fn(group_idx, batch_idx, value_columns);
    427
    0
                }
    428
            }
    429
    0
            Some(valid_indices) => {
    430
    0
                for (batch_idx, &group_idx) in group_indices.iter().enumerate() {
    431
    0
                    if valid_indices.value(batch_idx) {
    432
    0
                        value_fn(group_idx, batch_idx, value_columns);
    433
    0
                    }
    434
                }
    435
            }
    436
        }
    437
    0
    }
    438
    439
    /// This function is called to update the accumulator state per row
    440
    /// when the value is not needed (e.g. COUNT)
    441
    ///
    442
    /// `F`: Invoked like `value_fn(group_index) for all non null values
    443
    /// passing the filter. Note that no tracking is done for null inputs
    444
    /// or which groups have seen any values
    445
    ///
    446
    /// See [`NullState::accumulate`], for more details on other
    447
    /// arguments.
    448
    0
    pub fn accumulate_indices<F>(
    449
    0
        group_indices: &[usize],
    450
    0
        nulls: Option<&NullBuffer>,
    451
    0
        opt_filter: Option<&BooleanArray>,
    452
    0
        mut index_fn: F,
    453
    0
    ) where
    454
    0
        F: FnMut(usize) + Send,
    455
    {
    456
    0
        match (nulls, opt_filter) {
    457
            (None, None) => {
    458
    0
                for &group_index in group_indices.iter() {
    459
    0
                    index_fn(group_index)
    460
                }
    461
            }
    462
    0
            (None, Some(filter)) => {
    463
    0
                debug_assert_eq!(filter.len(), group_indices.len());
    464
    0
                let group_indices_chunks = group_indices.chunks_exact(64);
    465
    0
                let bit_chunks = filter.values().bit_chunks();
    466
    467
    0
                let group_indices_remainder = group_indices_chunks.remainder();
    468
    469
    0
                group_indices_chunks.zip(bit_chunks.iter()).for_each(
    470
    0
                    |(group_index_chunk, mask)| {
    471
                        // index_mask has value 1 << i in the loop
    472
    0
                        let mut index_mask = 1;
    473
    0
                        group_index_chunk.iter().for_each(|&group_index| {
    474
                            // valid bit was set, real vale
    475
    0
                            let is_valid = (mask & index_mask) != 0;
    476
    0
                            if is_valid {
    477
    0
                                index_fn(group_index);
    478
    0
                            }
    479
    0
                            index_mask <<= 1;
    480
    0
                        })
    481
    0
                    },
    482
                );
    483
    484
                // handle any remaining bits (after the initial 64)
    485
    0
                let remainder_bits = bit_chunks.remainder_bits();
    486
    0
                group_indices_remainder
    487
    0
                    .iter()
    488
    0
                    .enumerate()
    489
    0
                    .for_each(|(i, &group_index)| {
    490
    0
                        let is_valid = remainder_bits & (1 << i) != 0;
    491
    0
                        if is_valid {
    492
    0
                            index_fn(group_index)
    493
    0
                        }
    494
    0
                    });
    495
            }
    496
    0
            (Some(valids), None) => {
    497
    0
                debug_assert_eq!(valids.len(), group_indices.len());
    498
                // This is based on (ahem, COPY/PASTA) arrow::compute::aggregate::sum
    499
                // iterate over in chunks of 64 bits for more efficient null checking
    500
    0
                let group_indices_chunks = group_indices.chunks_exact(64);
    501
    0
                let bit_chunks = valids.inner().bit_chunks();
    502
    503
    0
                let group_indices_remainder = group_indices_chunks.remainder();
    504
    505
    0
                group_indices_chunks.zip(bit_chunks.iter()).for_each(
    506
    0
                    |(group_index_chunk, mask)| {
    507
                        // index_mask has value 1 << i in the loop
    508
    0
                        let mut index_mask = 1;
    509
    0
                        group_index_chunk.iter().for_each(|&group_index| {
    510
                            // valid bit was set, real vale
    511
    0
                            let is_valid = (mask & index_mask) != 0;
    512
    0
                            if is_valid {
    513
    0
                                index_fn(group_index);
    514
    0
                            }
    515
    0
                            index_mask <<= 1;
    516
    0
                        })
    517
    0
                    },
    518
                );
    519
    520
                // handle any remaining bits (after the initial 64)
    521
    0
                let remainder_bits = bit_chunks.remainder_bits();
    522
    0
                group_indices_remainder
    523
    0
                    .iter()
    524
    0
                    .enumerate()
    525
    0
                    .for_each(|(i, &group_index)| {
    526
    0
                        let is_valid = remainder_bits & (1 << i) != 0;
    527
    0
                        if is_valid {
    528
    0
                            index_fn(group_index)
    529
    0
                        }
    530
    0
                    });
    531
            }
    532
    533
    0
            (Some(valids), Some(filter)) => {
    534
    0
                debug_assert_eq!(filter.len(), group_indices.len());
    535
    0
                debug_assert_eq!(valids.len(), group_indices.len());
    536
    537
    0
                let group_indices_chunks = group_indices.chunks_exact(64);
    538
    0
                let valid_bit_chunks = valids.inner().bit_chunks();
    539
    0
                let filter_bit_chunks = filter.values().bit_chunks();
    540
    541
    0
                let group_indices_remainder = group_indices_chunks.remainder();
    542
    543
    0
                group_indices_chunks
    544
    0
                    .zip(valid_bit_chunks.iter())
    545
    0
                    .zip(filter_bit_chunks.iter())
    546
    0
                    .for_each(|((group_index_chunk, valid_mask), filter_mask)| {
    547
                        // index_mask has value 1 << i in the loop
    548
    0
                        let mut index_mask = 1;
    549
    0
                        group_index_chunk.iter().for_each(|&group_index| {
    550
                            // valid bit was set, real vale
    551
    0
                            let is_valid = (valid_mask & filter_mask & index_mask) != 0;
    552
    0
                            if is_valid {
    553
    0
                                index_fn(group_index);
    554
    0
                            }
    555
    0
                            index_mask <<= 1;
    556
    0
                        })
    557
    0
                    });
    558
    559
                // handle any remaining bits (after the initial 64)
    560
    0
                let remainder_valid_bits = valid_bit_chunks.remainder_bits();
    561
    0
                let remainder_filter_bits = filter_bit_chunks.remainder_bits();
    562
    0
                group_indices_remainder
    563
    0
                    .iter()
    564
    0
                    .enumerate()
    565
    0
                    .for_each(|(i, &group_index)| {
    566
    0
                        let is_valid =
    567
    0
                            remainder_valid_bits & remainder_filter_bits & (1 << i) != 0;
    568
    0
                        if is_valid {
    569
    0
                            index_fn(group_index)
    570
    0
                        }
    571
    0
                    });
    572
            }
    573
        }
    574
    0
    }
    575
    576
    /// Ensures that `builder` contains a `BooleanBufferBuilder with at
    577
    /// least `total_num_groups`.
    578
    ///
    579
    /// All new entries are initialized to `default_value`
    580
    fn initialize_builder(
    581
        builder: &mut BooleanBufferBuilder,
    582
        total_num_groups: usize,
    583
        default_value: bool,
    584
    ) -> &mut BooleanBufferBuilder {
    585
        if builder.len() < total_num_groups {
    586
            let new_groups = total_num_groups - builder.len();
    587
            builder.append_n(new_groups, default_value);
    588
        }
    589
        builder
    590
    }
    591
    592
    #[cfg(test)]
    593
    mod test {
    594
        use super::*;
    595
    596
        use arrow::array::{Int32Array, UInt32Array};
    597
        use rand::{rngs::ThreadRng, Rng};
    598
        use std::collections::HashSet;
    599
    600
        #[test]
    601
        fn accumulate() {
    602
            let group_indices = (0..100).collect();
    603
            let values = (0..100).map(|i| (i + 1) * 10).collect();
    604
            let values_with_nulls = (0..100)
    605
                .map(|i| if i % 3 == 0 { None } else { Some((i + 1) * 10) })
    606
                .collect();
    607
    608
            // default to every fifth value being false, every even
    609
            // being null
    610
            let filter: BooleanArray = (0..100)
    611
                .map(|i| {
    612
                    let is_even = i % 2 == 0;
    613
                    let is_fifth = i % 5 == 0;
    614
                    if is_even {
    615
                        None
    616
                    } else if is_fifth {
    617
                        Some(false)
    618
                    } else {
    619
                        Some(true)
    620
                    }
    621
                })
    622
                .collect();
    623
    624
            Fixture {
    625
                group_indices,
    626
                values,
    627
                values_with_nulls,
    628
                filter,
    629
            }
    630
            .run()
    631
        }
    632
    633
        #[test]
    634
        fn accumulate_fuzz() {
    635
            let mut rng = rand::rng();
    636
            for _ in 0..100 {
    637
                Fixture::new_random(&mut rng).run();
    638
            }
    639
        }
    640
    641
        /// Values for testing (there are enough values to exercise the 64 bit chunks
    642
        struct Fixture {
    643
            /// 100..0
    644
            group_indices: Vec<usize>,
    645
    646
            /// 10, 20, ... 1010
    647
            values: Vec<u32>,
    648
    649
            /// same as values, but every third is null:
    650
            /// None, Some(20), Some(30), None ...
    651
            values_with_nulls: Vec<Option<u32>>,
    652
    653
            /// filter (defaults to None)
    654
            filter: BooleanArray,
    655
        }
    656
    657
        impl Fixture {
    658
            fn new_random(rng: &mut ThreadRng) -> Self {
    659
                // Number of input values in a batch
    660
                let num_values: usize = rng.random_range(1..200);
    661
                // number of distinct groups
    662
                let num_groups: usize = rng.random_range(2..1000);
    663
                let max_group = num_groups - 1;
    664
    665
                let group_indices: Vec<usize> = (0..num_values)
    666
                    .map(|_| rng.random_range(0..max_group))
    667
                    .collect();
    668
    669
                let values: Vec<u32> = (0..num_values).map(|_| rng.random()).collect();
    670
    671
                // 10% chance of false
    672
                // 10% change of null
    673
                // 80% chance of true
    674
                let filter: BooleanArray = (0..num_values)
    675
                    .map(|_| {
    676
                        let filter_value = rng.random_range(0.0..1.0);
    677
                        if filter_value < 0.1 {
    678
                            Some(false)
    679
                        } else if filter_value < 0.2 {
    680
                            None
    681
                        } else {
    682
                            Some(true)
    683
                        }
    684
                    })
    685
                    .collect();
    686
    687
                // random values with random number and location of nulls
    688
                // random null percentage
    689
                let null_pct: f32 = rng.random_range(0.0..1.0);
    690
                let values_with_nulls: Vec<Option<u32>> = (0..num_values)
    691
                    .map(|_| {
    692
                        let is_null = null_pct < rng.random_range(0.0..1.0);
    693
                        if is_null {
    694
                            None
    695
                        } else {
    696
                            Some(rng.random())
    697
                        }
    698
                    })
    699
                    .collect();
    700
    701
                Self {
    702
                    group_indices,
    703
                    values,
    704
                    values_with_nulls,
    705
                    filter,
    706
                }
    707
            }
    708
    709
            /// returns `Self::values` an Array
    710
            fn values_array(&self) -> UInt32Array {
    711
                UInt32Array::from(self.values.clone())
    712
            }
    713
    714
            /// returns `Self::values_with_nulls` as an Array
    715
            fn values_with_nulls_array(&self) -> UInt32Array {
    716
                UInt32Array::from(self.values_with_nulls.clone())
    717
            }
    718
    719
            /// Calls `NullState::accumulate` and `accumulate_indices`
    720
            /// with all combinations of nulls and filter values
    721
            fn run(&self) {
    722
                let total_num_groups = *self.group_indices.iter().max().unwrap() + 1;
    723
    724
                let group_indices = &self.group_indices;
    725
                let values_array = self.values_array();
    726
                let values_with_nulls_array = self.values_with_nulls_array();
    727
                let filter = &self.filter;
    728
    729
                // no null, no filters
    730
                Self::accumulate_test(group_indices, &values_array, None, total_num_groups);
    731
    732
                // nulls, no filters
    733
                Self::accumulate_test(
    734
                    group_indices,
    735
                    &values_with_nulls_array,
    736
                    None,
    737
                    total_num_groups,
    738
                );
    739
    740
                // no nulls, filters
    741
                Self::accumulate_test(
    742
                    group_indices,
    743
                    &values_array,
    744
                    Some(filter),
    745
                    total_num_groups,
    746
                );
    747
    748
                // nulls, filters
    749
                Self::accumulate_test(
    750
                    group_indices,
    751
                    &values_with_nulls_array,
    752
                    Some(filter),
    753
                    total_num_groups,
    754
                );
    755
            }
    756
    757
            /// Calls `NullState::accumulate` and `accumulate_indices` to
    758
            /// ensure it generates the correct values.
    759
            fn accumulate_test(
    760
                group_indices: &[usize],
    761
                values: &UInt32Array,
    762
                opt_filter: Option<&BooleanArray>,
    763
                total_num_groups: usize,
    764
            ) {
    765
                Self::accumulate_values_test(
    766
                    group_indices,
    767
                    values,
    768
                    opt_filter,
    769
                    total_num_groups,
    770
                );
    771
                Self::accumulate_indices_test(group_indices, values.nulls(), opt_filter);
    772
    773
                // Convert values into a boolean array (anything above the
    774
                // average is true, otherwise false)
    775
                let avg: usize = values.iter().filter_map(|v| v.map(|v| v as usize)).sum();
    776
                let boolean_values: BooleanArray =
    777
                    values.iter().map(|v| v.map(|v| v as usize > avg)).collect();
    778
                Self::accumulate_boolean_test(
    779
                    group_indices,
    780
                    &boolean_values,
    781
                    opt_filter,
    782
                    total_num_groups,
    783
                );
    784
            }
    785
    786
            /// This is effectively a different implementation of
    787
            /// accumulate that we compare with the above implementation
    788
            fn accumulate_values_test(
    789
                group_indices: &[usize],
    790
                values: &UInt32Array,
    791
                opt_filter: Option<&BooleanArray>,
    792
                total_num_groups: usize,
    793
            ) {
    794
                let mut accumulated_values = vec![];
    795
                let mut null_state = NullState::new();
    796
    797
                null_state.accumulate(
    798
                    group_indices,
    799
                    values,
    800
                    opt_filter,
    801
                    total_num_groups,
    802
                    |group_index, value| {
    803
                        accumulated_values.push((group_index, value));
    804
                    },
    805
                );
    806
    807
                // Figure out the expected values
    808
                let mut expected_values = vec![];
    809
                let mut mock = MockNullState::new();
    810
    811
                match opt_filter {
    812
                    None => group_indices.iter().zip(values.iter()).for_each(
    813
                        |(&group_index, value)| {
    814
                            if let Some(value) = value {
    815
                                mock.saw_value(group_index);
    816
                                expected_values.push((group_index, value));
    817
                            }
    818
                        },
    819
                    ),
    820
                    Some(filter) => {
    821
                        group_indices
    822
                            .iter()
    823
                            .zip(values.iter())
    824
                            .zip(filter.iter())
    825
                            .for_each(|((&group_index, value), is_included)| {
    826
                                // if value passed filter
    827
                                if let Some(true) = is_included {
    828
                                    if let Some(value) = value {
    829
                                        mock.saw_value(group_index);
    830
                                        expected_values.push((group_index, value));
    831
                                    }
    832
                                }
    833
                            });
    834
                    }
    835
                }
    836
    837
                assert_eq!(accumulated_values, expected_values,
    838
                           "\n\naccumulated_values:{accumulated_values:#?}\n\nexpected_values:{expected_values:#?}");
    839
                let seen_values = null_state.seen_values.finish_cloned();
    840
                mock.validate_seen_values(&seen_values);
    841
    842
                // Validate the final buffer (one value per group)
    843
                let expected_null_buffer = mock.expected_null_buffer(total_num_groups);
    844
    845
                let null_buffer = null_state.build(EmitTo::All);
    846
    847
                assert_eq!(null_buffer, expected_null_buffer);
    848
            }
    849
    850
            // Calls `accumulate_indices`
    851
            // and opt_filter and ensures it calls the right values
    852
            fn accumulate_indices_test(
    853
                group_indices: &[usize],
    854
                nulls: Option<&NullBuffer>,
    855
                opt_filter: Option<&BooleanArray>,
    856
            ) {
    857
                let mut accumulated_values = vec![];
    858
    859
                accumulate_indices(group_indices, nulls, opt_filter, |group_index| {
    860
                    accumulated_values.push(group_index);
    861
                });
    862
    863
                // Figure out the expected values
    864
                let mut expected_values = vec![];
    865
    866
                match (nulls, opt_filter) {
    867
                    (None, None) => group_indices.iter().for_each(|&group_index| {
    868
                        expected_values.push(group_index);
    869
                    }),
    870
                    (Some(nulls), None) => group_indices.iter().zip(nulls.iter()).for_each(
    871
                        |(&group_index, is_valid)| {
    872
                            if is_valid {
    873
                                expected_values.push(group_index);
    874
                            }
    875
                        },
    876
                    ),
    877
                    (None, Some(filter)) => group_indices.iter().zip(filter.iter()).for_each(
    878
                        |(&group_index, is_included)| {
    879
                            if let Some(true) = is_included {
    880
                                expected_values.push(group_index);
    881
                            }
    882
                        },
    883
                    ),
    884
                    (Some(nulls), Some(filter)) => {
    885
                        group_indices
    886
                            .iter()
    887
                            .zip(nulls.iter())
    888
                            .zip(filter.iter())
    889
                            .for_each(|((&group_index, is_valid), is_included)| {
    890
                                // if value passed filter
    891
                                if let (true, Some(true)) = (is_valid, is_included) {
    892
                                    expected_values.push(group_index);
    893
                                }
    894
                            });
    895
                    }
    896
                }
    897
    898
                assert_eq!(accumulated_values, expected_values,
    899
                           "\n\naccumulated_values:{accumulated_values:#?}\n\nexpected_values:{expected_values:#?}");
    900
            }
    901
    902
            /// This is effectively a different implementation of
    903
            /// accumulate_boolean that we compare with the above implementation
    904
            fn accumulate_boolean_test(
    905
                group_indices: &[usize],
    906
                values: &BooleanArray,
    907
                opt_filter: Option<&BooleanArray>,
    908
                total_num_groups: usize,
    909
            ) {
    910
                let mut accumulated_values = vec![];
    911
                let mut null_state = NullState::new();
    912
    913
                null_state.accumulate_boolean(
    914
                    group_indices,
    915
                    values,
    916
                    opt_filter,
    917
                    total_num_groups,
    918
                    |group_index, value| {
    919
                        accumulated_values.push((group_index, value));
    920
                    },
    921
                );
    922
    923
                // Figure out the expected values
    924
                let mut expected_values = vec![];
    925
                let mut mock = MockNullState::new();
    926
    927
                match opt_filter {
    928
                    None => group_indices.iter().zip(values.iter()).for_each(
    929
                        |(&group_index, value)| {
    930
                            if let Some(value) = value {
    931
                                mock.saw_value(group_index);
    932
                                expected_values.push((group_index, value));
    933
                            }
    934
                        },
    935
                    ),
    936
                    Some(filter) => {
    937
                        group_indices
    938
                            .iter()
    939
                            .zip(values.iter())
    940
                            .zip(filter.iter())
    941
                            .for_each(|((&group_index, value), is_included)| {
    942
                                // if value passed filter
    943
                                if let Some(true) = is_included {
    944
                                    if let Some(value) = value {
    945
                                        mock.saw_value(group_index);
    946
                                        expected_values.push((group_index, value));
    947
                                    }
    948
                                }
    949
                            });
    950
                    }
    951
                }
    952
    953
                assert_eq!(accumulated_values, expected_values,
    954
                           "\n\naccumulated_values:{accumulated_values:#?}\n\nexpected_values:{expected_values:#?}");
    955
    956
                let seen_values = null_state.seen_values.finish_cloned();
    957
                mock.validate_seen_values(&seen_values);
    958
    959
                // Validate the final buffer (one value per group)
    960
                let expected_null_buffer = mock.expected_null_buffer(total_num_groups);
    961
    962
                let null_buffer = null_state.build(EmitTo::All);
    963
    964
                assert_eq!(null_buffer, expected_null_buffer);
    965
            }
    966
        }
    967
    968
        /// Parallel implementation of NullState to check expected values
    969
        #[derive(Debug, Default)]
    970
        struct MockNullState {
    971
            /// group indices that had values that passed the filter
    972
            seen_values: HashSet<usize>,
    973
        }
    974
    975
        impl MockNullState {
    976
            fn new() -> Self {
    977
                Default::default()
    978
            }
    979
    980
            fn saw_value(&mut self, group_index: usize) {
    981
                self.seen_values.insert(group_index);
    982
            }
    983
    984
            /// did this group index see any input?
    985
            fn expected_seen(&self, group_index: usize) -> bool {
    986
                self.seen_values.contains(&group_index)
    987
            }
    988
    989
            /// Validate that the seen_values matches self.seen_values
    990
            fn validate_seen_values(&self, seen_values: &BooleanBuffer) {
    991
                for (group_index, is_seen) in seen_values.iter().enumerate() {
    992
                    let expected_seen = self.expected_seen(group_index);
    993
                    assert_eq!(
    994
                        expected_seen, is_seen,
    995
                        "mismatch at for group {group_index}"
    996
                    );
    997
                }
    998
            }
    999
    1000
            /// Create the expected null buffer based on if the input had nulls and a filter
    1001
            fn expected_null_buffer(&self, total_num_groups: usize) -> NullBuffer {
    1002
                (0..total_num_groups)
    1003
                    .map(|group_index| self.expected_seen(group_index))
    1004
                    .collect()
    1005
            }
    1006
        }
    1007
    1008
        #[test]
    1009
        fn test_accumulate_multiple_no_nulls_no_filter() {
    1010
            let group_indices = vec![0, 1, 0, 1];
    1011
            let values1 = Int32Array::from(vec![1, 2, 3, 4]);
    1012
            let values2 = Int32Array::from(vec![10, 20, 30, 40]);
    1013
            let value_columns = [values1, values2];
    1014
    1015
            let mut accumulated = vec![];
    1016
            accumulate_multiple(
    1017
                &group_indices,
    1018
                &value_columns.iter().collect::<Vec<_>>(),
    1019
                None,
    1020
                |group_idx, batch_idx, columns| {
    1021
                    let values = columns.iter().map(|col| col.value(batch_idx)).collect();
    1022
                    accumulated.push((group_idx, values));
    1023
                },
    1024
            );
    1025
    1026
            let expected = vec![
    1027
                (0, vec![1, 10]),
    1028
                (1, vec![2, 20]),
    1029
                (0, vec![3, 30]),
    1030
                (1, vec![4, 40]),
    1031
            ];
    1032
            assert_eq!(accumulated, expected);
    1033
        }
    1034
    1035
        #[test]
    1036
        fn test_accumulate_multiple_with_nulls() {
    1037
            let group_indices = vec![0, 1, 0, 1];
    1038
            let values1 = Int32Array::from(vec![Some(1), None, Some(3), Some(4)]);
    1039
            let values2 = Int32Array::from(vec![Some(10), Some(20), None, Some(40)]);
    1040
            let value_columns = [values1, values2];
    1041
    1042
            let mut accumulated = vec![];
    1043
            accumulate_multiple(
    1044
                &group_indices,
    1045
                &value_columns.iter().collect::<Vec<_>>(),
    1046
                None,
    1047
                |group_idx, batch_idx, columns| {
    1048
                    let values = columns.iter().map(|col| col.value(batch_idx)).collect();
    1049
                    accumulated.push((group_idx, values));
    1050
                },
    1051
            );
    1052
    1053
            // Only rows where both columns are non-null should be accumulated
    1054
            let expected = vec![(0, vec![1, 10]), (1, vec![4, 40])];
    1055
            assert_eq!(accumulated, expected);
    1056
        }
    1057
    1058
        #[test]
    1059
        fn test_accumulate_multiple_with_filter() {
    1060
            let group_indices = vec![0, 1, 0, 1];
    1061
            let values1 = Int32Array::from(vec![1, 2, 3, 4]);
    1062
            let values2 = Int32Array::from(vec![10, 20, 30, 40]);
    1063
            let value_columns = [values1, values2];
    1064
    1065
            let filter = BooleanArray::from(vec![true, false, true, false]);
    1066
    1067
            let mut accumulated = vec![];
    1068
            accumulate_multiple(
    1069
                &group_indices,
    1070
                &value_columns.iter().collect::<Vec<_>>(),
    1071
                Some(&filter),
    1072
                |group_idx, batch_idx, columns| {
    1073
                    let values = columns.iter().map(|col| col.value(batch_idx)).collect();
    1074
                    accumulated.push((group_idx, values));
    1075
                },
    1076
            );
    1077
    1078
            // Only rows where filter is true should be accumulated
    1079
            let expected = vec![(0, vec![1, 10]), (0, vec![3, 30])];
    1080
            assert_eq!(accumulated, expected);
    1081
        }
    1082
    1083
        #[test]
    1084
        fn test_accumulate_multiple_with_nulls_and_filter() {
    1085
            let group_indices = vec![0, 1, 0, 1];
    1086
            let values1 = Int32Array::from(vec![Some(1), None, Some(3), Some(4)]);
    1087
            let values2 = Int32Array::from(vec![Some(10), Some(20), None, Some(40)]);
    1088
            let value_columns = [values1, values2];
    1089
    1090
            let filter = BooleanArray::from(vec![true, true, true, false]);
    1091
    1092
            let mut accumulated = vec![];
    1093
            accumulate_multiple(
    1094
                &group_indices,
    1095
                &value_columns.iter().collect::<Vec<_>>(),
    1096
                Some(&filter),
    1097
                |group_idx, batch_idx, columns| {
    1098
                    let values = columns.iter().map(|col| col.value(batch_idx)).collect();
    1099
                    accumulated.push((group_idx, values));
    1100
                },
    1101
            );
    1102
    1103
            // Only rows where both:
    1104
            // 1. Filter is true
    1105
            // 2. Both columns are non-null
    1106
            // should be accumulated
    1107
            let expected = [(0, vec![1, 10])];
    1108
            assert_eq!(accumulated, expected);
    1109
        }
    1110
    }