feat: Support duration type in approx_distinct

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<usize>)> {
                         .map(|_| Some(rng.random::<i64>()))
                         .collect::<Int64Array>(),
                 ),
+                DataType::Duration(TimeUnit::Nanosecond) => Arc::new(
+                    (0..BATCH_SIZE)
+                        .map(|_| Some(rng.random::<i64>()))
+                        .collect::<DurationNanosecondArray>(),
+                ),
                 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),

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::<Decimal256Type>::new())
             }
+            DataType::Duration(TimeUnit::Second) => {
+                Box::new(NumericHLLAccumulator::<DurationSecondType>::new())
+            }
+            DataType::Duration(TimeUnit::Millisecond) => {
+                Box::new(NumericHLLAccumulator::<DurationMillisecondType>::new())
+            }
+            DataType::Duration(TimeUnit::Microsecond) => {
+                Box::new(NumericHLLAccumulator::<DurationMicrosecondType>::new())
+            }
+            DataType::Duration(TimeUnit::Nanosecond) => {
+                Box::new(NumericHLLAccumulator::<DurationNanosecondType>::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::<DurationSecondType>(
+                duration_second,
+                6,
+            );
+
+            let duration_millisecond: ArrayRef =
+                Arc::new(DurationMillisecondArray::from(values.clone()));
+            assert_count_numerical_acc_and_group_acc::<DurationMillisecondType>(
+                duration_millisecond,
+                6,
+            );
+
+            let duration_microsecond: ArrayRef =
+                Arc::new(DurationMicrosecondArray::from(values.clone()));
+            assert_count_numerical_acc_and_group_acc::<DurationMicrosecondType>(
+                duration_microsecond,
+                6,
+            );
+
+            let duration_nanosecond: ArrayRef =
+                Arc::new(DurationNanosecondArray::from(values));
+            assert_count_numerical_acc_and_group_acc::<DurationNanosecondType>(
+                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]

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.