perf: Optimize approx_distinct() for string, binary inputs (apache#21037)

## Which issue does this PR close?

- Closes apache#21035.

## Rationale for this change

We were making a defensive copy of every string and binary value before
calling HyperLogLog, but that was unnecessary: HyperLogLog doesn't need
owned types, and indeed never stores the input value.

This improves the performance of `approx_distinct` on string and binary
values by 6x-7x.

## What changes are included in this PR?

* Add benchmark for `approx_distinct`
* Optimize `approx_distinct` to avoid unnecessary copies
* Cleanup: remove spurious type parameter from
`StringViewHLLAccumulator` (unused)

## Are these changes tested?

Yes.

## Are there any user-facing changes?

No.

Author: neilconway

datafusion/functions-aggregate/Cargo.toml

@@ -75,3 +75,7 @@ harness = false
 [[bench]]
 harness = false
 name = "min_max_bytes"
+
+[[bench]]
+name = "approx_distinct"
+harness = false

datafusion/functions-aggregate/benches/approx_distinct.rs

@@ -0,0 +1,128 @@
+// 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.
+
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int64Array, StringArray, StringViewArray};
+use arrow::datatypes::{DataType, Field, Schema};
+use criterion::{Criterion, criterion_group, criterion_main};
+use datafusion_expr::function::AccumulatorArgs;
+use datafusion_expr::{Accumulator, AggregateUDFImpl};
+use datafusion_functions_aggregate::approx_distinct::ApproxDistinct;
+use datafusion_physical_expr::expressions::col;
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+
+const BATCH_SIZE: usize = 8192;
+const STRING_LENGTH: usize = 20;
+
+fn prepare_accumulator(data_type: DataType) -> Box<dyn Accumulator> {
+    let schema = Arc::new(Schema::new(vec![Field::new("f", data_type, true)]));
+    let expr = col("f", &schema).unwrap();
+    let accumulator_args = AccumulatorArgs {
+        return_field: Field::new("f", DataType::UInt64, true).into(),
+        schema: &schema,
+        expr_fields: &[expr.return_field(&schema).unwrap()],
+        ignore_nulls: false,
+        order_bys: &[],
+        is_reversed: false,
+        name: "approx_distinct(f)",
+        is_distinct: false,
+        exprs: &[expr],
+    };
+    ApproxDistinct::new().accumulator(accumulator_args).unwrap()
+}
+
+/// Creates an Int64Array where values are drawn from `0..n_distinct`.
+fn create_i64_array(n_distinct: usize) -> Int64Array {
+    let mut rng = StdRng::seed_from_u64(42);
+    (0..BATCH_SIZE)
+        .map(|_| Some(rng.random_range(0..n_distinct as i64)))
+        .collect()
+}
+
+/// Creates a pool of `n_distinct` random strings.
+fn create_string_pool(n_distinct: usize) -> Vec<String> {
+    let mut rng = StdRng::seed_from_u64(42);
+    (0..n_distinct)
+        .map(|_| {
+            (0..STRING_LENGTH)
+                .map(|_| rng.random_range(b'a'..=b'z') as char)
+                .collect()
+        })
+        .collect()
+}
+
+/// Creates a StringArray where values are drawn from the given pool.
+fn create_string_array(pool: &[String]) -> StringArray {
+    let mut rng = StdRng::seed_from_u64(99);
+    (0..BATCH_SIZE)
+        .map(|_| Some(pool[rng.random_range(0..pool.len())].as_str()))
+        .collect()
+}
+
+/// Creates a StringViewArray where values are drawn from the given pool.
+fn create_string_view_array(pool: &[String]) -> StringViewArray {
+    let mut rng = StdRng::seed_from_u64(99);
+    (0..BATCH_SIZE)
+        .map(|_| Some(pool[rng.random_range(0..pool.len())].as_str()))
+        .collect()
+}
+
+fn approx_distinct_benchmark(c: &mut Criterion) {
+    for pct in [80, 99] {
+        let n_distinct = BATCH_SIZE * pct / 100;
+
+        // --- Int64 benchmarks ---
+        let values = Arc::new(create_i64_array(n_distinct)) as ArrayRef;
+        c.bench_function(&format!("approx_distinct i64 {pct}% distinct"), |b| {
+            b.iter(|| {
+                let mut accumulator = prepare_accumulator(DataType::Int64);
+                accumulator
+                    .update_batch(std::slice::from_ref(&values))
+                    .unwrap()
+            })
+        });
+
+        let string_pool = create_string_pool(n_distinct);
+
+        // --- Utf8 benchmarks ---
+        let values = Arc::new(create_string_array(&string_pool)) as ArrayRef;
+        c.bench_function(&format!("approx_distinct utf8 {pct}% distinct"), |b| {
+            b.iter(|| {
+                let mut accumulator = prepare_accumulator(DataType::Utf8);
+                accumulator
+                    .update_batch(std::slice::from_ref(&values))
+                    .unwrap()
+            })
+        });
+
+        // --- Utf8View benchmarks ---
+        let values = Arc::new(create_string_view_array(&string_pool)) as ArrayRef;
+        c.bench_function(&format!("approx_distinct utf8view {pct}% distinct"), |b| {
+            b.iter(|| {
+                let mut accumulator = prepare_accumulator(DataType::Utf8View);
+                accumulator
+                    .update_batch(std::slice::from_ref(&values))
+                    .unwrap()
+            })
+        });
+    }
+}
+
+criterion_group!(benches, approx_distinct_benchmark);
+criterion_main!(benches);

datafusion/functions-aggregate/src/approx_distinct.rs

@@ -55,14 +55,14 @@ make_udaf_expr_and_func!(
     approx_distinct_udaf
 );
 
-impl<T: Hash> From<&HyperLogLog<T>> for ScalarValue {
+impl<T: Hash + ?Sized> From<&HyperLogLog<T>> for ScalarValue {
     fn from(v: &HyperLogLog<T>) -> ScalarValue {
         let values = v.as_ref().to_vec();
         ScalarValue::Binary(Some(values))
     }
 }
 
-impl<T: Hash> TryFrom<&[u8]> for HyperLogLog<T> {
+impl<T: Hash + ?Sized> TryFrom<&[u8]> for HyperLogLog<T> {
     type Error = DataFusionError;
     fn try_from(v: &[u8]) -> Result<HyperLogLog<T>> {
         let arr: [u8; 16384] = v.try_into().map_err(|_| {
@@ -72,7 +72,7 @@ impl<T: Hash> TryFrom<&[u8]> for HyperLogLog<T> {
     }
 }
 
-impl<T: Hash> TryFrom<&ScalarValue> for HyperLogLog<T> {
+impl<T: Hash + ?Sized> TryFrom<&ScalarValue> for HyperLogLog<T> {
     type Error = DataFusionError;
     fn try_from(v: &ScalarValue) -> Result<HyperLogLog<T>> {
         if let ScalarValue::Binary(Some(slice)) = v {
@@ -99,7 +99,6 @@ where
     T: ArrowPrimitiveType,
     T::Native: Hash,
 {
-    /// new approx_distinct accumulator
     pub fn new() -> Self {
         Self {
             hll: HyperLogLog::new(),
@@ -112,15 +111,14 @@ struct StringHLLAccumulator<T>
 where
     T: OffsetSizeTrait,
 {
-    hll: HyperLogLog<String>,
+    hll: HyperLogLog<str>,
     phantom_data: PhantomData<T>,
 }
 
 impl<T> StringHLLAccumulator<T>
 where
     T: OffsetSizeTrait,
 {
-    /// new approx_distinct accumulator
     pub fn new() -> Self {
         Self {
             hll: HyperLogLog::new(),
@@ -130,22 +128,14 @@ where
 }
 
 #[derive(Debug)]
-struct StringViewHLLAccumulator<T>
-where
-    T: OffsetSizeTrait,
-{
-    hll: HyperLogLog<String>,
-    phantom_data: PhantomData<T>,
+struct StringViewHLLAccumulator {
+    hll: HyperLogLog<str>,
 }
 
-impl<T> StringViewHLLAccumulator<T>
-where
-    T: OffsetSizeTrait,
-{
+impl StringViewHLLAccumulator {
     pub fn new() -> Self {
         Self {
             hll: HyperLogLog::new(),
-            phantom_data: PhantomData,
         }
     }
 }
@@ -155,15 +145,14 @@ struct BinaryHLLAccumulator<T>
 where
     T: OffsetSizeTrait,
 {
-    hll: HyperLogLog<Vec<u8>>,
+    hll: HyperLogLog<[u8]>,
     phantom_data: PhantomData<T>,
 }
 
 impl<T> BinaryHLLAccumulator<T>
 where
     T: OffsetSizeTrait,
 {
-    /// new approx_distinct accumulator
     pub fn new() -> Self {
         Self {
             hll: HyperLogLog::new(),
@@ -213,23 +202,18 @@ where
         let array: &GenericBinaryArray<T> =
             downcast_value!(values[0], GenericBinaryArray, T);
         // flatten because we would skip nulls
-        self.hll
-            .extend(array.into_iter().flatten().map(|v| v.to_vec()));
+        self.hll.extend(array.into_iter().flatten());
         Ok(())
     }
 
     default_accumulator_impl!();
 }
 
-impl<T> Accumulator for StringViewHLLAccumulator<T>
-where
-    T: OffsetSizeTrait,
-{
+impl Accumulator for StringViewHLLAccumulator {
     fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
         let array: &StringViewArray = downcast_value!(values[0], StringViewArray);
         // flatten because we would skip nulls
-        self.hll
-            .extend(array.iter().flatten().map(|s| s.to_string()));
+        self.hll.extend(array.iter().flatten());
         Ok(())
     }
 
@@ -244,8 +228,7 @@ where
         let array: &GenericStringArray<T> =
             downcast_value!(values[0], GenericStringArray, T);
         // flatten because we would skip nulls
-        self.hll
-            .extend(array.into_iter().flatten().map(|i| i.to_string()));
+        self.hll.extend(array.into_iter().flatten());
         Ok(())
     }
 
@@ -391,7 +374,7 @@ impl AggregateUDFImpl for ApproxDistinct {
             }
             DataType::Utf8 => Box::new(StringHLLAccumulator::<i32>::new()),
             DataType::LargeUtf8 => Box::new(StringHLLAccumulator::<i64>::new()),
-            DataType::Utf8View => Box::new(StringViewHLLAccumulator::<i32>::new()),
+            DataType::Utf8View => Box::new(StringViewHLLAccumulator::new()),
             DataType::Binary => Box::new(BinaryHLLAccumulator::<i32>::new()),
             DataType::LargeBinary => Box::new(BinaryHLLAccumulator::<i64>::new()),
             DataType::Null => {