## 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/utils.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
    use ahash::RandomState;
    19
    use arrow::array::{
    20
        Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, PrimitiveArray,
    21
    };
    22
    use arrow::compute::SortOptions;
    23
    use arrow::datatypes::{
    24
        ArrowNativeType, DataType, DecimalType, Field, FieldRef, ToByteSlice,
    25
    };
    26
    use datafusion_common::cast::{as_list_array, as_primitive_array};
    27
    use datafusion_common::utils::memory::estimate_memory_size;
    28
    use datafusion_common::utils::SingleRowListArrayBuilder;
    29
    use datafusion_common::{
    30
        exec_err, internal_datafusion_err, HashSet, Result, ScalarValue,
    31
    };
    32
    use datafusion_expr_common::accumulator::Accumulator;
    33
    use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
    34
    use std::sync::Arc;
    35
    36
    /// Convert scalar values from an accumulator into arrays.
    37
    pub fn get_accum_scalar_values_as_arrays(
    38
        accum: &mut dyn Accumulator,
    39
    ) -> Result<Vec<ArrayRef>> {
    40
        accum
    41
            .state()?
    42
            .iter()
    43
    0
            .map(|s| s.to_array_of_size(1))
    44
            .collect()
    45
    }
    46
    47
    /// Construct corresponding fields for the expressions in an ORDER BY clause.
    48
    pub fn ordering_fields(
    49
        order_bys: &[PhysicalSortExpr],
    50
        // Data type of each expression in the ordering requirement
    51
        data_types: &[DataType],
    52
    ) -> Vec<FieldRef> {
    53
        order_bys
    54
            .iter()
    55
            .zip(data_types.iter())
    56
    0
            .map(|(sort_expr, dtype)| {
    57
    0
                Field::new(
    58
    0
                    sort_expr.expr.to_string().as_str(),
    59
    0
                    dtype.clone(),
    60
                    // Multi partitions may be empty hence field should be nullable.
    61
                    true,
    62
                )
    63
    0
            })
    64
            .map(Arc::new)
    65
            .collect()
    66
    }
    67
    68
    /// Selects the sort option attribute from all the given `PhysicalSortExpr`s.
    69
    pub fn get_sort_options(ordering_req: &LexOrdering) -> Vec<SortOptions> {
    70
        ordering_req.iter().map(|item| item.options).collect()
    71
    }
    72
    73
    /// A wrapper around a type to provide hash for floats
    74
    #[derive(Copy, Clone, Debug)]
    75
    pub struct Hashable<T>(pub T);
    76
    77
    impl<T: ToByteSlice> std::hash::Hash for Hashable<T> {
    78
    0
        fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
    79
    0
            self.0.to_byte_slice().hash(state)
    80
    0
        }
    81
    }
    82
    83
    impl<T: ArrowNativeTypeOp> PartialEq for Hashable<T> {
    84
    0
        fn eq(&self, other: &Self) -> bool {
    85
    0
            self.0.is_eq(other.0)
    86
    0
        }
    87
    }
    88
    89
    impl<T: ArrowNativeTypeOp> Eq for Hashable<T> {}
    90
    91
    /// Computes averages for `Decimal128`/`Decimal256` values, checking for overflow
    92
    ///
    93
    /// This is needed because different precisions for Decimal128/Decimal256 can
    94
    /// store different ranges of values and thus sum/count may not fit in
    95
    /// the target type.
    96
    ///
    97
    /// For example, the precision is 3, the max of value is `999` and the min
    98
    /// value is `-999`
    99
    pub struct DecimalAverager<T: DecimalType> {
    100
        /// scale factor for sum values (10^sum_scale)
    101
        sum_mul: T::Native,
    102
        /// scale factor for target (10^target_scale)
    103
        target_mul: T::Native,
    104
        /// the output precision
    105
        target_precision: u8,
    106
        /// the output scale
    107
        target_scale: i8,
    108
    }
    109
    110
    impl<T: DecimalType> DecimalAverager<T> {
    111
        /// Create a new `DecimalAverager`:
    112
        ///
    113
        /// * sum_scale: the scale of `sum` values passed to [`Self::avg`]
    114
        /// * target_precision: the output precision
    115
        /// * target_scale: the output scale
    116
        ///
    117
        /// Errors if the resulting data can not be stored
    118
    0
        pub fn try_new(
    119
    0
            sum_scale: i8,
    120
    0
            target_precision: u8,
    121
    0
            target_scale: i8,
    122
    0
        ) -> Result<Self> {
    123
    0
            let sum_mul = T::Native::from_usize(10_usize)
    124
    0
                .map(|b| b.pow_wrapping(sum_scale as u32))
    125
    0
                .ok_or_else(|| {
    126
    0
                    internal_datafusion_err!("Failed to compute sum_mul in DecimalAverager")
    127
    0
                })?;
    128
    129
    0
            let target_mul = T::Native::from_usize(10_usize)
    130
    0
                .map(|b| b.pow_wrapping(target_scale as u32))
    131
    0
                .ok_or_else(|| {
    132
    0
                    internal_datafusion_err!(
    133
    0
                        "Failed to compute target_mul in DecimalAverager"
    134
                    )
    135
    0
                })?;
    136
    137
    0
            if target_mul >= sum_mul {
    138
    0
                Ok(Self {
    139
    0
                    sum_mul,
    140
    0
                    target_mul,
    141
    0
                    target_precision,
    142
    0
                    target_scale,
    143
    0
                })
    144
            } else {
    145
                // can't convert the lit decimal to the returned data type
    146
    0
                exec_err!("Arithmetic Overflow in AvgAccumulator")
    147
            }
    148
    0
        }
    149
    150
        /// Returns the `sum`/`count` as a i128/i256 Decimal128/Decimal256 with
    151
        /// target_scale and target_precision and reporting overflow.
    152
        ///
    153
        /// * sum: The total sum value stored as Decimal128 with sum_scale
    154
        ///   (passed to `Self::try_new`)
    155
        /// * count: total count, stored as a i128/i256 (*NOT* a Decimal128/Decimal256 value)
    156
        #[inline(always)]
    157
    0
        pub fn avg(&self, sum: T::Native, count: T::Native) -> Result<T::Native> {
    158
    0
            if let Ok(value) = sum.mul_checked(self.target_mul.div_wrapping(self.sum_mul)) {
    159
    0
                let new_value = value.div_wrapping(count);
    160
    161
    0
                let validate = T::validate_decimal_precision(
    162
    0
                    new_value,
    163
    0
                    self.target_precision,
    164
    0
                    self.target_scale,
    165
                );
    166
    167
    0
                if validate.is_ok() {
    168
    0
                    Ok(new_value)
    169
                } else {
    170
    0
                    exec_err!("Arithmetic Overflow in AvgAccumulator")
    171
                }
    172
            } else {
    173
                // can't convert the lit decimal to the returned data type
    174
    0
                exec_err!("Arithmetic Overflow in AvgAccumulator")
    175
            }
    176
    0
        }
    177
    }
    178
    179
    /// Generic way to collect distinct values for accumulators.
    180
    ///
    181
    /// The intermediate state is represented as a List of scalar values updated by
    182
    /// `merge_batch` and a `Vec` of `ArrayRef` that are converted to scalar values
    183
    /// in the final evaluation step so that we avoid expensive conversions and
    184
    /// allocations during `update_batch`.
    185
    pub struct GenericDistinctBuffer<T: ArrowPrimitiveType> {
    186
        pub values: HashSet<Hashable<T::Native>, RandomState>,
    187
        data_type: DataType,
    188
    }
    189
    190
    impl<T: ArrowPrimitiveType> std::fmt::Debug for GenericDistinctBuffer<T> {
    191
    0
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    192
    0
            write!(
    193
    0
                f,
    194
    0
                "GenericDistinctBuffer({}, values={})",
    195
                self.data_type,
    196
    0
                self.values.len()
    197
            )
    198
    0
        }
    199
    }
    200
    201
    impl<T: ArrowPrimitiveType> GenericDistinctBuffer<T> {
    202
    0
        pub fn new(data_type: DataType) -> Self {
    203
    0
            Self {
    204
    0
                values: HashSet::default(),
    205
    0
                data_type,
    206
    0
            }
    207
    0
        }
    208
    209
        /// Mirrors [`Accumulator::state`].
    210
    0
        pub fn state(&self) -> Result<Vec<ScalarValue>> {
    211
    0
            let arr = Arc::new(
    212
    0
                PrimitiveArray::<T>::from_iter_values(self.values.iter().map(|v| v.0))
    213
                    // Ideally we'd just use T::DATA_TYPE but this misses things like
    214
                    // decimal scale/precision and timestamp timezones, which need to
    215
                    // match up with Accumulator::state_fields
    216
    0
                    .with_data_type(self.data_type.clone()),
    217
            );
    218
    0
            Ok(vec![SingleRowListArrayBuilder::new(arr).build_list_scalar()])
    219
    0
        }
    220
    221
        /// Mirrors [`Accumulator::update_batch`].
    222
    0
        pub fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
    223
    0
            if values.is_empty() {
    224
    0
                return Ok(());
    225
    0
            }
    226
    227
    0
            debug_assert_eq!(
    228
    0
                values.len(),
    229
                1,
    230
    0
                "DistinctValuesBuffer::update_batch expects only a single input array"
    231
            );
    232
    233
    0
            let arr = as_primitive_array::<T>(&values[0])?;
    234
    0
            if arr.null_count() > 0 {
    235
    0
                self.values.extend(arr.iter().flatten().map(Hashable));
    236
    0
            } else {
    237
    0
                self.values
    238
    0
                    .extend(arr.values().iter().cloned().map(Hashable));
    239
    0
            }
    240
    241
    0
            Ok(())
    242
    0
        }
    243
    244
        /// Mirrors [`Accumulator::merge_batch`].
    245
    0
        pub fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
    246
    0
            if states.is_empty() {
    247
    0
                return Ok(());
    248
    0
            }
    249
    250
    0
            let array = as_list_array(&states[0])?;
    251
    0
            for list in array.iter().flatten() {
    252
    0
                self.update_batch(&[list])?;
    253
            }
    254
    255
    0
            Ok(())
    256
    0
        }
    257
    258
        /// Mirrors [`Accumulator::size`].
    259
    0
        pub fn size(&self) -> usize {
    260
    0
            let num_elements = self.values.len();
    261
    0
            let fixed_size = size_of_val(self) + size_of_val(&self.values);
    262
    0
            estimate_memory_size::<T::Native>(num_elements, fixed_size).unwrap()
    263
    0
        }
    264
    }