Implement specialized group values for single Uft8/LargeUtf8/Binary/LargeBinary column

datafusion/physical-expr/src/aggregate/count_distinct/bytes.rs

@@ -0,0 +1,90 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`BytesDistinctCountAccumulator`] for Utf8/LargeUtf8/Binary/LargeBinary values
+
+use crate::binary_map::{ArrowBytesSet, OutputType};
+use arrow_array::{ArrayRef, OffsetSizeTrait};
+use datafusion_common::cast::as_list_array;
+use datafusion_common::utils::array_into_list_array;
+use datafusion_common::ScalarValue;
+use datafusion_expr::Accumulator;
+use std::fmt::Debug;
+use std::sync::Arc;
+
+/// Specialized implementation of
+/// `COUNT DISTINCT` for [`StringArray`] [`LargeStringArray`],
+/// [`BinaryArray`] and [`LargeBinaryArray`].
+///
+/// [`StringArray`]: arrow::array::StringArray
+/// [`LargeStringArray`]: arrow::array::LargeStringArray
+/// [`BinaryArray`]: arrow::array::BinaryArray
+/// [`LargeBinaryArray`]: arrow::array::LargeBinaryArray
+#[derive(Debug)]
+pub(super) struct BytesDistinctCountAccumulator<O: OffsetSizeTrait>(ArrowBytesSet<O>);
+
+impl<O: OffsetSizeTrait> BytesDistinctCountAccumulator<O> {
+    pub(super) fn new(output_type: OutputType) -> Self {
+        Self(ArrowBytesSet::new(output_type))
+    }
+}
+
+impl<O: OffsetSizeTrait> Accumulator for BytesDistinctCountAccumulator<O> {
+    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
+        let set = self.0.take();
+        let arr = set.into_state();
+        let list = Arc::new(array_into_list_array(arr));
+        Ok(vec![ScalarValue::List(list)])
+    }
+
+    fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
+        if values.is_empty() {
+            return Ok(());
+        }
+
+        self.0.insert(&values[0]);
+
+        Ok(())
+    }
+
+    fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
+        if states.is_empty() {
+            return Ok(());
+        }
+        assert_eq!(
+            states.len(),
+            1,
+            "count_distinct states must be single array"
+        );
+
+        let arr = as_list_array(&states[0])?;
+        arr.iter().try_for_each(|maybe_list| {
+            if let Some(list) = maybe_list {
+                self.0.insert(&list);
+            };
+            Ok(())
+        })
+    }
+
+    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
+        Ok(ScalarValue::Int64(Some(self.0.non_null_len() as i64)))
+    }
+
+    fn size(&self) -> usize {
+        std::mem::size_of_val(self) + self.0.size()
+    }
+}

datafusion/physical-expr/src/aggregate/count_distinct/mod.rs

@@ -15,8 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
+mod bytes;
 mod native;
-mod strings;
 
 use std::any::Any;
 use std::collections::HashSet;
@@ -37,11 +37,12 @@ use arrow_array::types::{
 use datafusion_common::{Result, ScalarValue};
 use datafusion_expr::Accumulator;
 
+use crate::aggregate::count_distinct::bytes::BytesDistinctCountAccumulator;
 use crate::aggregate::count_distinct::native::{
     FloatDistinctCountAccumulator, PrimitiveDistinctCountAccumulator,
 };
-use crate::aggregate::count_distinct::strings::StringDistinctCountAccumulator;
 use crate::aggregate::utils::down_cast_any_ref;
+use crate::binary_map::OutputType;
 use crate::expressions::format_state_name;
 use crate::{AggregateExpr, PhysicalExpr};
 
@@ -144,8 +145,16 @@ impl AggregateExpr for DistinctCount {
             Float32 => Box::new(FloatDistinctCountAccumulator::<Float32Type>::new()),
             Float64 => Box::new(FloatDistinctCountAccumulator::<Float64Type>::new()),
 
-            Utf8 => Box::new(StringDistinctCountAccumulator::<i32>::new()),
-            LargeUtf8 => Box::new(StringDistinctCountAccumulator::<i64>::new()),
+            Utf8 => Box::new(BytesDistinctCountAccumulator::<i32>::new(OutputType::Utf8)),
+            LargeUtf8 => {
+                Box::new(BytesDistinctCountAccumulator::<i64>::new(OutputType::Utf8))
+            }
+            Binary => Box::new(BytesDistinctCountAccumulator::<i32>::new(
+                OutputType::Binary,
+            )),
+            LargeBinary => Box::new(BytesDistinctCountAccumulator::<i64>::new(
+                OutputType::Binary,
+            )),
 
             _ => Box::new(DistinctCountAccumulator {
                 values: HashSet::default(),
@@ -175,7 +184,7 @@ impl PartialEq<dyn Any> for DistinctCount {
 /// General purpose distinct accumulator that works for any DataType by using
 /// [`ScalarValue`]. Some types have specialized accumulators that are (much)
 /// more efficient such as [`PrimitiveDistinctCountAccumulator`] and
-/// [`StringDistinctCountAccumulator`]
+/// [`BytesDistinctCountAccumulator`]
 #[derive(Debug)]
 struct DistinctCountAccumulator {
     values: HashSet<ScalarValue, RandomState>,