diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 5fe3f350d73fb..3b74d3b7148c3 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 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.