From 98a25ce5b03498bf78ecd02750ab12251f48e675 Mon Sep 17 00:00:00 2001 From: Michael Kleen Date: Tue, 30 Jun 2026 23:50:07 +0200 Subject: [PATCH 1/4] feat: Support Duration type in approx_distinct --- .../benches/approx_distinct.rs | 34 ++++++++++-- .../src/approx_distinct.rs | 54 +++++++++++++++++-- .../sqllogictest/test_files/aggregate.slt | 27 ++++++++++ 3 files changed, 109 insertions(+), 6 deletions(-) diff --git a/datafusion/functions-aggregate/benches/approx_distinct.rs b/datafusion/functions-aggregate/benches/approx_distinct.rs index 2ab783d9acf05..4f102a860c5ea 100644 --- a/datafusion/functions-aggregate/benches/approx_distinct.rs +++ b/datafusion/functions-aggregate/benches/approx_distinct.rs @@ -20,11 +20,13 @@ use std::sync::Arc; use arrow::array::{ ArrayRef, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, - Int8Array, Int16Array, Int64Array, IntervalDayTimeArray, IntervalMonthDayNanoArray, - IntervalYearMonthArray, StringArray, StringViewArray, UInt8Array, UInt16Array, + DurationNanosecondArray, Int8Array, Int16Array, Int64Array, IntervalDayTimeArray, + IntervalMonthDayNanoArray, IntervalYearMonthArray, StringArray, StringViewArray, + UInt8Array, UInt16Array, }; +use arrow::datatypes::{DataType, Field, TimeUnit}; use arrow::datatypes::{ - DataType, Field, IntervalDayTime, IntervalMonthDayNano, IntervalUnit, Schema, i256, + IntervalDayTime, IntervalMonthDayNano, IntervalUnit, Schema, i256, }; use criterion::{Criterion, criterion_group, criterion_main}; use datafusion_expr::function::AccumulatorArgs; @@ -124,6 +126,14 @@ fn create_i64_array(n_distinct: usize) -> Int64Array { .collect() } +/// Creates a `DurationNanosecondArray` where values are drawn from `0..n_distinct`. +fn create_duration_array(n_distinct: usize) -> DurationNanosecondArray { + let mut rng = StdRng::seed_from_u64(42); + (0..BATCH_SIZE) + .map(|_| Some(rng.random_range(0..n_distinct as i64))) + .collect() +} + fn create_u8_array(n_distinct: usize) -> UInt8Array { let mut rng = StdRng::seed_from_u64(42); let max_val = n_distinct.min(256) as u8; @@ -231,6 +241,18 @@ fn approx_distinct_benchmark(c: &mut Criterion) { }) }); + // --- Duration benchmarks --- + let values = Arc::new(create_duration_array(n_distinct)) as ArrayRef; + c.bench_function(&format!("approx_distinct duration {pct}% distinct"), |b| { + b.iter(|| { + let mut accumulator = + prepare_accumulator(DataType::Duration(TimeUnit::Nanosecond)); + accumulator + .update_batch(std::slice::from_ref(&values)) + .unwrap() + }) + }); + for (label, str_len) in [("short", SHORT_STRING_LENGTH), ("long", LONG_STRING_LENGTH)] { @@ -472,6 +494,11 @@ fn build_grouped_batches(data_type: &DataType) -> Vec<(ArrayRef, Vec)> { .map(|_| Some(rng.random::())) .collect::(), ), + DataType::Duration(TimeUnit::Nanosecond) => Arc::new( + (0..BATCH_SIZE) + .map(|_| Some(rng.random::())) + .collect::(), + ), DataType::Utf8 => Arc::new( (0..BATCH_SIZE) .map(|_| Some(pool[rng.random_range(0..pool.len())].as_str())) @@ -551,6 +578,7 @@ fn approx_distinct_grouped_benchmark(c: &mut Criterion) { for data_type in [ DataType::Int64, + DataType::Duration(TimeUnit::Nanosecond), DataType::Utf8, DataType::Utf8View, DataType::Decimal32(DECIMAL32_PRECISION, DECIMAL_SCALE), diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 5fe3f350d73fb..b15894d547e34 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -25,7 +25,8 @@ use arrow::array::{ use arrow::buffer::NullBuffer; use arrow::datatypes::{ ArrowPrimitiveType, DataType, Date32Type, Date64Type, Decimal32Type, Decimal64Type, - Decimal128Type, Decimal256Type, Field, FieldRef, Int32Type, Int64Type, + Decimal128Type, Decimal256Type, DurationMicrosecondType, DurationMillisecondType, + DurationNanosecondType, DurationSecondType, Field, FieldRef, Int32Type, Int64Type, IntervalDayTimeType, IntervalMonthDayNanoType, IntervalUnit, IntervalYearMonthType, Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, @@ -781,6 +782,18 @@ impl AggregateUDFImpl for ApproxDistinct { DataType::Decimal256(_, _) => { Box::new(NumericHLLAccumulator::::new()) } + DataType::Duration(TimeUnit::Second) => { + Box::new(NumericHLLAccumulator::::new()) + } + DataType::Duration(TimeUnit::Millisecond) => { + Box::new(NumericHLLAccumulator::::new()) + } + DataType::Duration(TimeUnit::Microsecond) => { + Box::new(NumericHLLAccumulator::::new()) + } + DataType::Duration(TimeUnit::Nanosecond) => { + Box::new(NumericHLLAccumulator::::new()) + } DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View @@ -848,6 +861,7 @@ fn is_hll_groups_type(data_type: &DataType) -> bool { | DataType::Decimal64(_, _) | DataType::Decimal128(_, _) | DataType::Decimal256(_, _) + | DataType::Duration(_) | DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View @@ -866,8 +880,9 @@ mod tests { use super::*; use arrow::array::{ AsArray, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, - Int64Array, IntervalDayTimeArray, IntervalMonthDayNanoArray, - IntervalYearMonthArray, StringViewArray, + DurationMicrosecondArray, DurationMillisecondArray, DurationNanosecondArray, + DurationSecondArray, Int64Array, IntervalDayTimeArray, + IntervalMonthDayNanoArray, IntervalYearMonthArray, StringViewArray, }; use arrow::datatypes::{IntervalDayTime, IntervalMonthDayNano, i256}; use std::sync::Arc; @@ -1039,6 +1054,39 @@ mod tests { ); } + #[test] + fn duration_support_numerical_acc_and_group_acc() { + let values = vec![1i64, 2, 2, 3, 3, 3, 0, 0, i64::MAX, i64::MIN, i64::MIN]; + + let duration_second: ArrayRef = + Arc::new(DurationSecondArray::from(values.clone())); + assert_count_numerical_acc_and_group_acc::( + duration_second, + 6, + ); + + let duration_millisecond: ArrayRef = + Arc::new(DurationMillisecondArray::from(values.clone())); + assert_count_numerical_acc_and_group_acc::( + duration_millisecond, + 6, + ); + + let duration_microsecond: ArrayRef = + Arc::new(DurationMicrosecondArray::from(values.clone())); + assert_count_numerical_acc_and_group_acc::( + duration_microsecond, + 6, + ); + + let duration_nanosecond: ArrayRef = + Arc::new(DurationNanosecondArray::from(values)); + assert_count_numerical_acc_and_group_acc::( + duration_nanosecond, + 6, + ); + } + /// `approx_distinct(v) FILTER (WHERE nullable_bool)` — a NULL filter row /// must not be counted (null filter is treated the same as false). #[test] diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index dbf5063a30188..bdcc2135aedb3 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -2005,6 +2005,33 @@ FROM approx_distinct_decimal_test GROUP BY g ORDER BY g; statement ok DROP TABLE approx_distinct_decimal_test; +# This test runs approx_distinct over all four Duration units for the scalar and the grouped path. +statement ok +CREATE TABLE approx_distinct_duration_test AS VALUES + (1, arrow_cast(1, 'Duration(Second)'), arrow_cast(1, 'Duration(Millisecond)'), arrow_cast(1, 'Duration(Microsecond)'), arrow_cast(1, 'Duration(Nanosecond)')), + (1, arrow_cast(2, 'Duration(Second)'), arrow_cast(2, 'Duration(Millisecond)'), arrow_cast(2, 'Duration(Microsecond)'), arrow_cast(2, 'Duration(Nanosecond)')), + (1, arrow_cast(2, 'Duration(Second)'), arrow_cast(2, 'Duration(Millisecond)'), arrow_cast(2, 'Duration(Microsecond)'), arrow_cast(2, 'Duration(Nanosecond)')), + (2, arrow_cast(3, 'Duration(Second)'), arrow_cast(3, 'Duration(Millisecond)'), arrow_cast(3, 'Duration(Microsecond)'), arrow_cast(3, 'Duration(Nanosecond)')), + (2, arrow_cast(0, 'Duration(Second)'), arrow_cast(0, 'Duration(Millisecond)'), arrow_cast(0, 'Duration(Microsecond)'), arrow_cast(0, 'Duration(Nanosecond)')), + (2, arrow_cast(0, 'Duration(Second)'), arrow_cast(0, 'Duration(Millisecond)'), arrow_cast(0, 'Duration(Microsecond)'), arrow_cast(0, 'Duration(Nanosecond)')); + +# Scalar path +query IIII +SELECT approx_distinct(column2), approx_distinct(column3), approx_distinct(column4), approx_distinct(column5) FROM approx_distinct_duration_test; +---- +4 4 4 4 + +# Grouped path +query IIIII +SELECT column1, approx_distinct(column2), approx_distinct(column3), approx_distinct(column4), approx_distinct(column5) +FROM approx_distinct_duration_test GROUP BY column1 ORDER BY column1; +---- +1 2 2 2 2 +2 2 2 2 2 + +statement ok +DROP TABLE approx_distinct_duration_test; + # This test runs approx_distinct over the intervals YearMonth, # DayTime, MonthDayNano for the scalar and the grouped path. From 4d9b72b6b6ed1897c69d1a9e578cd111121fd3d0 Mon Sep 17 00:00:00 2001 From: Michael Kleen Date: Mon, 6 Jul 2026 09:18:02 +0200 Subject: [PATCH 2/4] Remove tests in approx_distinct.rs --- .../src/approx_distinct.rs | 38 +------------------ 1 file changed, 2 insertions(+), 36 deletions(-) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index b15894d547e34..7ae7d26e89919 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -880,9 +880,8 @@ mod tests { use super::*; use arrow::array::{ AsArray, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, - DurationMicrosecondArray, DurationMillisecondArray, DurationNanosecondArray, - DurationSecondArray, Int64Array, IntervalDayTimeArray, - IntervalMonthDayNanoArray, IntervalYearMonthArray, StringViewArray, + Int64Array, IntervalDayTimeArray, IntervalMonthDayNanoArray, IntervalYearMonthArray, + StringViewArray, }; use arrow::datatypes::{IntervalDayTime, IntervalMonthDayNano, i256}; use std::sync::Arc; @@ -1054,39 +1053,6 @@ mod tests { ); } - #[test] - fn duration_support_numerical_acc_and_group_acc() { - let values = vec![1i64, 2, 2, 3, 3, 3, 0, 0, i64::MAX, i64::MIN, i64::MIN]; - - let duration_second: ArrayRef = - Arc::new(DurationSecondArray::from(values.clone())); - assert_count_numerical_acc_and_group_acc::( - duration_second, - 6, - ); - - let duration_millisecond: ArrayRef = - Arc::new(DurationMillisecondArray::from(values.clone())); - assert_count_numerical_acc_and_group_acc::( - duration_millisecond, - 6, - ); - - let duration_microsecond: ArrayRef = - Arc::new(DurationMicrosecondArray::from(values.clone())); - assert_count_numerical_acc_and_group_acc::( - duration_microsecond, - 6, - ); - - let duration_nanosecond: ArrayRef = - Arc::new(DurationNanosecondArray::from(values)); - assert_count_numerical_acc_and_group_acc::( - duration_nanosecond, - 6, - ); - } - /// `approx_distinct(v) FILTER (WHERE nullable_bool)` — a NULL filter row /// must not be counted (null filter is treated the same as false). #[test] From fdc9b33cdae0399b00cd073f22094493e3590690 Mon Sep 17 00:00:00 2001 From: Michael Kleen Date: Mon, 6 Jul 2026 09:18:35 +0200 Subject: [PATCH 3/4] Remove benchmarks --- .../benches/approx_distinct.rs | 34 ++----------------- 1 file changed, 3 insertions(+), 31 deletions(-) diff --git a/datafusion/functions-aggregate/benches/approx_distinct.rs b/datafusion/functions-aggregate/benches/approx_distinct.rs index 4f102a860c5ea..2ab783d9acf05 100644 --- a/datafusion/functions-aggregate/benches/approx_distinct.rs +++ b/datafusion/functions-aggregate/benches/approx_distinct.rs @@ -20,13 +20,11 @@ use std::sync::Arc; use arrow::array::{ ArrayRef, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, - DurationNanosecondArray, Int8Array, Int16Array, Int64Array, IntervalDayTimeArray, - IntervalMonthDayNanoArray, IntervalYearMonthArray, StringArray, StringViewArray, - UInt8Array, UInt16Array, + Int8Array, Int16Array, Int64Array, IntervalDayTimeArray, IntervalMonthDayNanoArray, + IntervalYearMonthArray, StringArray, StringViewArray, UInt8Array, UInt16Array, }; -use arrow::datatypes::{DataType, Field, TimeUnit}; use arrow::datatypes::{ - IntervalDayTime, IntervalMonthDayNano, IntervalUnit, Schema, i256, + DataType, Field, IntervalDayTime, IntervalMonthDayNano, IntervalUnit, Schema, i256, }; use criterion::{Criterion, criterion_group, criterion_main}; use datafusion_expr::function::AccumulatorArgs; @@ -126,14 +124,6 @@ fn create_i64_array(n_distinct: usize) -> Int64Array { .collect() } -/// Creates a `DurationNanosecondArray` where values are drawn from `0..n_distinct`. -fn create_duration_array(n_distinct: usize) -> DurationNanosecondArray { - let mut rng = StdRng::seed_from_u64(42); - (0..BATCH_SIZE) - .map(|_| Some(rng.random_range(0..n_distinct as i64))) - .collect() -} - fn create_u8_array(n_distinct: usize) -> UInt8Array { let mut rng = StdRng::seed_from_u64(42); let max_val = n_distinct.min(256) as u8; @@ -241,18 +231,6 @@ fn approx_distinct_benchmark(c: &mut Criterion) { }) }); - // --- Duration benchmarks --- - let values = Arc::new(create_duration_array(n_distinct)) as ArrayRef; - c.bench_function(&format!("approx_distinct duration {pct}% distinct"), |b| { - b.iter(|| { - let mut accumulator = - prepare_accumulator(DataType::Duration(TimeUnit::Nanosecond)); - accumulator - .update_batch(std::slice::from_ref(&values)) - .unwrap() - }) - }); - for (label, str_len) in [("short", SHORT_STRING_LENGTH), ("long", LONG_STRING_LENGTH)] { @@ -494,11 +472,6 @@ fn build_grouped_batches(data_type: &DataType) -> Vec<(ArrayRef, Vec)> { .map(|_| Some(rng.random::())) .collect::(), ), - DataType::Duration(TimeUnit::Nanosecond) => Arc::new( - (0..BATCH_SIZE) - .map(|_| Some(rng.random::())) - .collect::(), - ), DataType::Utf8 => Arc::new( (0..BATCH_SIZE) .map(|_| Some(pool[rng.random_range(0..pool.len())].as_str())) @@ -578,7 +551,6 @@ fn approx_distinct_grouped_benchmark(c: &mut Criterion) { for data_type in [ DataType::Int64, - DataType::Duration(TimeUnit::Nanosecond), DataType::Utf8, DataType::Utf8View, DataType::Decimal32(DECIMAL32_PRECISION, DECIMAL_SCALE), From 6174140e2cb0618b44239b7eabf21e4155873c80 Mon Sep 17 00:00:00 2001 From: Michael Kleen Date: Mon, 6 Jul 2026 09:23:03 +0200 Subject: [PATCH 4/4] fix fmt --- datafusion/functions-aggregate/src/approx_distinct.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 7ae7d26e89919..3b74d3b7148c3 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -880,8 +880,8 @@ mod tests { use super::*; use arrow::array::{ AsArray, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, - Int64Array, IntervalDayTimeArray, IntervalMonthDayNanoArray, IntervalYearMonthArray, - StringViewArray, + Int64Array, IntervalDayTimeArray, IntervalMonthDayNanoArray, + IntervalYearMonthArray, StringViewArray, }; use arrow::datatypes::{IntervalDayTime, IntervalMonthDayNano, i256}; use std::sync::Arc;