perf: use native distinct accumulator for COUNT(DISTINCT) on Decimal32/Decimal64

datafusion/functions-aggregate/benches/count_distinct.rs

@@ -18,8 +18,8 @@
 use std::sync::Arc;
 
 use arrow::array::{
-    Array, ArrayRef, Int8Array, Int16Array, Int32Array, Int64Array, UInt8Array,
-    UInt16Array, UInt32Array,
+    Array, ArrayRef, Decimal32Array, Decimal64Array, Decimal128Array, Int8Array,
+    Int16Array, Int32Array, Int64Array, UInt8Array, UInt16Array, UInt32Array,
 };
 use arrow::datatypes::{DataType, Field, Schema};
 use criterion::{Criterion, criterion_group, criterion_main};
@@ -102,6 +102,33 @@ fn create_i32_array(n_distinct: usize) -> Int32Array {
         .collect()
 }
 
+fn create_decimal32_array(n_distinct: usize) -> Decimal32Array {
+    let mut rng = StdRng::seed_from_u64(42);
+    (0..BATCH_SIZE)
+        .map(|_| Some(rng.random_range(0..n_distinct as i32)))
+        .collect::<Decimal32Array>()
+        .with_precision_and_scale(9, 2)
+        .unwrap()
+}
+
+fn create_decimal64_array(n_distinct: usize) -> Decimal64Array {
+    let mut rng = StdRng::seed_from_u64(42);
+    (0..BATCH_SIZE)
+        .map(|_| Some(rng.random_range(0..n_distinct as i64)))
+        .collect::<Decimal64Array>()
+        .with_precision_and_scale(18, 2)
+        .unwrap()
+}
+
+fn create_decimal128_array(n_distinct: usize) -> Decimal128Array {
+    let mut rng = StdRng::seed_from_u64(42);
+    (0..BATCH_SIZE)
+        .map(|_| Some(rng.random_range(0..n_distinct as i128)))
+        .collect::<Decimal128Array>()
+        .with_precision_and_scale(38, 2)
+        .unwrap()
+}
+
 fn prepare_args(data_type: DataType) -> (Arc<Schema>, AccumulatorArgs<'static>) {
     let schema = Arc::new(Schema::new(vec![Field::new("f", data_type, true)]));
     let schema_leaked: &'static Schema = Box::leak(Box::new((*schema).clone()));
@@ -214,6 +241,45 @@ fn count_distinct_benchmark(c: &mut Criterion) {
             })
         });
     }
+
+    // Decimal types. Decimal128 already uses the native distinct accumulator
+    // and serves as a reference for Decimal32 and Decimal64.
+    for pct in [80, 99] {
+        let n_distinct = BATCH_SIZE * pct / 100;
+
+        // Decimal32
+        let values = Arc::new(create_decimal32_array(n_distinct)) as ArrayRef;
+        c.bench_function(&format!("count_distinct decimal32 {pct}% distinct"), |b| {
+            b.iter(|| {
+                let mut accumulator = prepare_accumulator(DataType::Decimal32(9, 2));
+                accumulator
+                    .update_batch(std::slice::from_ref(&values))
+                    .unwrap()
+            })
+        });
+
+        // Decimal64
+        let values = Arc::new(create_decimal64_array(n_distinct)) as ArrayRef;
+        c.bench_function(&format!("count_distinct decimal64 {pct}% distinct"), |b| {
+            b.iter(|| {
+                let mut accumulator = prepare_accumulator(DataType::Decimal64(18, 2));
+                accumulator
+                    .update_batch(std::slice::from_ref(&values))
+                    .unwrap()
+            })
+        });
+
+        // Decimal128
+        let values = Arc::new(create_decimal128_array(n_distinct)) as ArrayRef;
+        c.bench_function(&format!("count_distinct decimal128 {pct}% distinct"), |b| {
+            b.iter(|| {
+                let mut accumulator = prepare_accumulator(DataType::Decimal128(38, 2));
+                accumulator
+                    .update_batch(std::slice::from_ref(&values))
+                    .unwrap()
+            })
+        });
+    }
 }
 
 /// Create group indices with uniform distribution

datafusion/functions-aggregate/src/count.rs

@@ -20,10 +20,10 @@ use arrow::{
     buffer::BooleanBuffer,
     compute,
     datatypes::{
-        DataType, Date32Type, Date64Type, Decimal128Type, Decimal256Type, Field,
-        FieldRef, Float16Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type,
-        Int64Type, Time32MillisecondType, Time32SecondType, Time64MicrosecondType,
-        Time64NanosecondType, TimeUnit, TimestampMicrosecondType,
+        DataType, Date32Type, Date64Type, Decimal32Type, Decimal64Type, Decimal128Type,
+        Decimal256Type, Field, FieldRef, Float16Type, Float32Type, Float64Type, Int8Type,
+        Int16Type, Int32Type, Int64Type, Time32MillisecondType, Time32SecondType,
+        Time64MicrosecondType, Time64NanosecondType, TimeUnit, TimestampMicrosecondType,
         TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType,
         UInt8Type, UInt16Type, UInt32Type, UInt64Type,
     },
@@ -193,6 +193,13 @@ fn get_count_accumulator(data_type: &DataType) -> Box<dyn Accumulator> {
         DataType::UInt8 | DataType::Int8 | DataType::UInt16 | DataType::Int16 => {
             get_small_int_accumulator(data_type).unwrap()
         }
+
+        DataType::Decimal32(_, _) => Box::new(PrimitiveDistinctCountAccumulator::<
+            Decimal32Type,
+        >::new(data_type)),
+        DataType::Decimal64(_, _) => Box::new(PrimitiveDistinctCountAccumulator::<
+            Decimal64Type,
+        >::new(data_type)),
         DataType::Decimal128(_, _) => Box::new(PrimitiveDistinctCountAccumulator::<
             Decimal128Type,
         >::new(data_type)),

datafusion/sqllogictest/test_files/aggregate.slt

@@ -7112,6 +7112,33 @@ select c2, count(DISTINCT cast(c1 AS DECIMAL(10, 2))) from d_table GROUP BY c2 O
 A 2
 B 2
 
+# COUNT(DISTINCT) on Decimal32 / Decimal64. Counting two different columns
+# keeps single_distinct_to_groupby from rewriting the query, so the distinct
+# count accumulator is used.
+statement ok
+CREATE TABLE decimal_count_distinct AS SELECT
+  g,
+  arrow_cast(v, 'Decimal32(9, 2)') AS d32,
+  arrow_cast(v, 'Decimal64(18, 2)') AS d64
+FROM (VALUES
+  (1, 1.50), (1, 1.50), (1, 2.25), (1, NULL),
+  (2, 3.00), (2, 3.00), (2, 3.00)
+) AS t(g, v);
+
+query II
+SELECT COUNT(DISTINCT d32), COUNT(DISTINCT d64) FROM decimal_count_distinct;
+----
+3 3
+
+query III rowsort
+SELECT g, COUNT(DISTINCT d32), COUNT(DISTINCT d64) FROM decimal_count_distinct GROUP BY g;
+----
+1 2 2
+2 1 1
+
+statement ok
+DROP TABLE decimal_count_distinct;
+
 # Creating the table
 statement ok
 CREATE TABLE test_table (c1 INT, c2 INT, c3 INT)