## 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/count_distinct/bytes.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
    //! [`BytesDistinctCountAccumulator`] for Utf8/LargeUtf8/Binary/LargeBinary values
    19
    20
    use arrow::array::{ArrayRef, OffsetSizeTrait};
    21
    use datafusion_common::cast::as_list_array;
    22
    use datafusion_common::utils::SingleRowListArrayBuilder;
    23
    use datafusion_common::ScalarValue;
    24
    use datafusion_expr_common::accumulator::Accumulator;
    25
    use datafusion_physical_expr_common::binary_map::{ArrowBytesSet, OutputType};
    26
    use datafusion_physical_expr_common::binary_view_map::ArrowBytesViewSet;
    27
    use std::fmt::Debug;
    28
    use std::mem::size_of_val;
    29
    30
    /// Specialized implementation of
    31
    /// `COUNT DISTINCT` for [`StringArray`] [`LargeStringArray`],
    32
    /// [`BinaryArray`] and [`LargeBinaryArray`].
    33
    ///
    34
    /// [`StringArray`]: arrow::array::StringArray
    35
    /// [`LargeStringArray`]: arrow::array::LargeStringArray
    36
    /// [`BinaryArray`]: arrow::array::BinaryArray
    37
    /// [`LargeBinaryArray`]: arrow::array::LargeBinaryArray
    38
    #[derive(Debug)]
    39
    pub struct BytesDistinctCountAccumulator<O: OffsetSizeTrait>(ArrowBytesSet<O>);
    40
    41
    impl<O: OffsetSizeTrait> BytesDistinctCountAccumulator<O> {
    42
    0
        pub fn new(output_type: OutputType) -> Self {
    43
    0
            Self(ArrowBytesSet::new(output_type))
    44
    0
        }
    45
    }
    46
    47
    impl<O: OffsetSizeTrait> Accumulator for BytesDistinctCountAccumulator<O> {
    48
    0
        fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
    49
    0
            let set = self.0.take();
    50
    0
            let arr = set.into_state();
    51
    0
            Ok(vec![SingleRowListArrayBuilder::new(arr).build_list_scalar()])
    52
    0
        }
    53
    54
    0
        fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
    55
    0
            if values.is_empty() {
    56
    0
                return Ok(());
    57
    0
            }
    58
    59
    0
            self.0.insert(&values[0]);
    60
    61
    0
            Ok(())
    62
    0
        }
    63
    64
    0
        fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
    65
    0
            if states.is_empty() {
    66
    0
                return Ok(());
    67
    0
            }
    68
    0
            assert_eq!(
    69
    0
                states.len(),
    70
                1,
    71
    0
                "count_distinct states must be single array"
    72
            );
    73
    74
    0
            let arr = as_list_array(&states[0])?;
    75
    0
            arr.iter().try_for_each(|maybe_list| {
    76
    0
                if let Some(list) = maybe_list {
    77
    0
                    self.0.insert(&list);
    78
    0
                };
    79
    0
                Ok(())
    80
    0
            })
    81
    0
        }
    82
    83
    0
        fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
    84
    0
            Ok(ScalarValue::Int64(Some(self.0.non_null_len() as i64)))
    85
    0
        }
    86
    87
    0
        fn size(&self) -> usize {
    88
    0
            size_of_val(self) + self.0.size()
    89
    0
        }
    90
    }
    91
    92
    /// Specialized implementation of
    93
    /// `COUNT DISTINCT` for [`StringViewArray`] and [`BinaryViewArray`].
    94
    ///
    95
    /// [`StringViewArray`]: arrow::array::StringViewArray
    96
    /// [`BinaryViewArray`]: arrow::array::BinaryViewArray
    97
    #[derive(Debug)]
    98
    pub struct BytesViewDistinctCountAccumulator(ArrowBytesViewSet);
    99
    100
    impl BytesViewDistinctCountAccumulator {
    101
        pub fn new(output_type: OutputType) -> Self {
    102
            Self(ArrowBytesViewSet::new(output_type))
    103
        }
    104
    }
    105
    106
    impl Accumulator for BytesViewDistinctCountAccumulator {
    107
        fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
    108
            let set = self.0.take();
    109
            let arr = set.into_state();
    110
            Ok(vec![SingleRowListArrayBuilder::new(arr).build_list_scalar()])
    111
        }
    112
    113
        fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
    114
            if values.is_empty() {
    115
                return Ok(());
    116
            }
    117
    118
            self.0.insert(&values[0]);
    119
    120
            Ok(())
    121
        }
    122
    123
        fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
    124
            if states.is_empty() {
    125
                return Ok(());
    126
            }
    127
            assert_eq!(
    128
                states.len(),
    129
                1,
    130
                "count_distinct states must be single array"
    131
            );
    132
    133
            let arr = as_list_array(&states[0])?;
    134
    0
            arr.iter().try_for_each(|maybe_list| {
    135
    0
                if let Some(list) = maybe_list {
    136
    0
                    self.0.insert(&list);
    137
    0
                };
    138
    0
                Ok(())
    139
    0
            })
    140
        }
    141
    142
        fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
    143
            Ok(ScalarValue::Int64(Some(self.0.non_null_len() as i64)))
    144
        }
    145
    146
        fn size(&self) -> usize {
    147
            size_of_val(self) + self.0.size()
    148
        }
    149
    }