diff --git a/datafusion/functions-aggregate/benches/count_distinct.rs b/datafusion/functions-aggregate/benches/count_distinct.rs index 4d9e8c5b67b31..91ab3692eaa99 100644 --- a/datafusion/functions-aggregate/benches/count_distinct.rs +++ b/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::() + .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::() + .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::() + .with_precision_and_scale(38, 2) + .unwrap() +} + fn prepare_args(data_type: DataType) -> (Arc, 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 diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index 7c3bfdceefdf7..df523f2f0f423 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/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 { 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)), diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 61a6f4578f358..95de5169dbde8 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/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)