feat: add GroupColumn support for Duration in multi-column GROUP BY

datafusion/physical-plan/benches/multi_group_by.rs

@@ -27,9 +27,9 @@
 //! covers a `(FixedSizeBinary, Int32)` key to exercise the
 //! `FixedSizeBinaryGroupValueBuilder`.
 
-use arrow::array::{ArrayRef, Int32Array, UInt32Array};
+use arrow::array::{ArrayRef, DurationMicrosecondArray, Int32Array, UInt32Array};
 use arrow::compute::take;
-use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
 use arrow::util::bench_util::create_fsb_array;
 use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
 use datafusion_physical_plan::aggregates::group_values::GroupValues;
@@ -444,6 +444,90 @@ fn bench_fixed_size_binary(c: &mut Criterion) {
     group.finish();
 }
 
+fn make_duration_schema() -> SchemaRef {
+    Arc::new(Schema::new(vec![
+        Field::new("dur", DataType::Duration(TimeUnit::Microsecond), false),
+        Field::new("id", DataType::Int32, false),
+    ]))
+}
+
+/// Generate `(Duration(Microsecond), Int32)` batches with `num_distinct_groups`
+/// distinct keys.
+///
+/// Each distinct duration is `g` microseconds. The `Int32` column is keyed
+/// identically so the combined cardinality equals `num_distinct_groups`.
+fn generate_duration_batches(
+    num_distinct_groups: usize,
+    num_rows: usize,
+    batch_size: usize,
+) -> Vec<Vec<ArrayRef>> {
+    let num_full_batches = num_rows / batch_size;
+    let remainder = num_rows % batch_size;
+    let num_batches = num_full_batches + if remainder > 0 { 1 } else { 0 };
+
+    (0..num_batches)
+        .map(|batch_idx| {
+            let batch_start = batch_idx * batch_size;
+            let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 {
+                remainder
+            } else {
+                batch_size
+            };
+
+            let group_ids = (0..current_batch_size)
+                .map(|row| (batch_start + row) % num_distinct_groups);
+
+            let keys = DurationMicrosecondArray::from_iter_values(
+                group_ids.clone().map(|g| g as i64),
+            );
+            let id: Int32Array = group_ids.map(|g| g as i32).collect();
+
+            vec![Arc::new(keys) as ArrayRef, Arc::new(id) as ArrayRef]
+        })
+        .collect()
+}
+
+/// Experiment 9: Group count sweep for a `(Duration, Int32)` key.
+///
+/// Exercises the primitive `GroupColumn` builder for `Duration` on the
+/// multi-column path (previously such a schema fell back to `GroupValuesRows`).
+fn bench_duration(c: &mut Criterion) {
+    let mut group = c.benchmark_group("duration");
+    group.sample_size(15);
+
+    let schema = make_duration_schema();
+
+    for num_groups in [1_000, 1_000_000] {
+        let batches =
+            generate_duration_batches(num_groups, 1_000_000, DEFAULT_BATCH_SIZE);
+
+        for vectorized in [true, false] {
+            let label = if vectorized {
+                "vectorized"
+            } else {
+                "row_based"
+            };
+            group.bench_with_input(
+                BenchmarkId::new(label, format!("grp_{num_groups}")),
+                &batches,
+                |b, batches| {
+                    b.iter_batched_ref(
+                        || {
+                            (
+                                create_group_values(&schema, vectorized),
+                                Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE),
+                            )
+                        },
+                        |(gv, groups)| bench_intern(gv, batches, groups),
+                        criterion::BatchSize::LargeInput,
+                    );
+                },
+            );
+        }
+    }
+    group.finish();
+}
+
 criterion_group!(
     benches,
     bench_issue_17850_regression,
@@ -453,5 +537,6 @@ criterion_group!(
     bench_high_cardinality_scaling,
     bench_group_count_sweep,
     bench_fixed_size_binary,
+    bench_duration,
 );
 criterion_main!(benches);

datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs

@@ -32,12 +32,13 @@ use crate::aggregates::group_values::multi_group_by::{
 use arrow::array::{Array, ArrayRef, BooleanBufferBuilder};
 use arrow::compute::cast;
 use arrow::datatypes::{
-    BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type, Field, Float32Type,
-    Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, Schema, SchemaRef,
-    StringViewType, Time32MillisecondType, Time32SecondType, Time64MicrosecondType,
-    Time64NanosecondType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType,
-    TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, UInt32Type,
-    UInt64Type,
+    BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type,
+    DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType,
+    DurationSecondType, Field, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type,
+    Int64Type, Schema, SchemaRef, StringViewType, Time32MillisecondType,
+    Time32SecondType, Time64MicrosecondType, Time64NanosecondType, TimeUnit,
+    TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
+    TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
 };
 use datafusion_common::hash_utils::RandomState;
 use datafusion_common::hash_utils::create_hashes;
@@ -952,6 +953,7 @@ fn group_column_supported_type(data_type: &DataType) -> bool {
             | DataType::Time64(TimeUnit::Microsecond)
             | DataType::Time64(TimeUnit::Nanosecond)
             | DataType::Timestamp(_, _)
+            | DataType::Duration(_)
             | DataType::Utf8View
             | DataType::BinaryView
             | DataType::Boolean
@@ -1031,6 +1033,20 @@ fn make_group_column(field: &Field) -> Result<Box<dyn GroupColumn>> {
                 instantiate_primitive!(v, nullable, TimestampNanosecondType, data_type)
             }
         },
+        DataType::Duration(t) => match t {
+            TimeUnit::Second => {
+                instantiate_primitive!(v, nullable, DurationSecondType, data_type)
+            }
+            TimeUnit::Millisecond => {
+                instantiate_primitive!(v, nullable, DurationMillisecondType, data_type)
+            }
+            TimeUnit::Microsecond => {
+                instantiate_primitive!(v, nullable, DurationMicrosecondType, data_type)
+            }
+            TimeUnit::Nanosecond => {
+                instantiate_primitive!(v, nullable, DurationNanosecondType, data_type)
+            }
+        },
         DataType::Decimal128(_, _) => {
             instantiate_primitive!(v, nullable, Decimal128Type, data_type)
         }
@@ -1259,7 +1275,10 @@ enum Nulls {
 mod tests {
     use std::{collections::HashMap, sync::Arc};
 
-    use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray, StringViewArray};
+    use arrow::array::{
+        Array, ArrayRef, DurationMicrosecondArray, Int64Array, RecordBatch, StringArray,
+        StringViewArray,
+    };
     use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
     use arrow::{compute::concat_batches, util::pretty::pretty_format_batches};
     use datafusion_common::utils::proxy::HashTableAllocExt;
@@ -1309,6 +1328,10 @@ mod tests {
             DataType::Time64(arrow::datatypes::TimeUnit::Microsecond),
             DataType::Time64(arrow::datatypes::TimeUnit::Nanosecond),
             DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None),
+            DataType::Duration(arrow::datatypes::TimeUnit::Second),
+            DataType::Duration(arrow::datatypes::TimeUnit::Millisecond),
+            DataType::Duration(arrow::datatypes::TimeUnit::Microsecond),
+            DataType::Duration(arrow::datatypes::TimeUnit::Nanosecond),
         ];
 
         for dt in &supported_cases {
@@ -1352,6 +1375,57 @@ mod tests {
         }
     }
 
+    // `Duration` group keys stay on the `GroupValuesColumn` fast path, dedup
+    // (including nulls), and round-trip with the `Duration` type preserved.
+    #[test]
+    fn test_group_values_column_duration() {
+        use arrow::datatypes::TimeUnit;
+
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("d", DataType::Duration(TimeUnit::Microsecond), true),
+            Field::new("i", DataType::Int64, true),
+        ]));
+        assert!(supported_schema(&schema));
+        let mut group_values =
+            GroupValuesColumn::<false>::try_new(Arc::clone(&schema)).unwrap();
+
+        // (d, i) rows, where row 3 repeats row 0 and row 4 repeats the null pair.
+        let d: ArrayRef = Arc::new(DurationMicrosecondArray::from(vec![
+            Some(10),
+            None,
+            Some(20),
+            Some(10),
+            None,
+        ]));
+        let i: ArrayRef = Arc::new(Int64Array::from(vec![
+            Some(1),
+            None,
+            Some(2),
+            Some(1),
+            None,
+        ]));
+        let mut groups = Vec::new();
+        group_values.intern(&[d, i], &mut groups).unwrap();
+        assert_eq!(groups, vec![0, 1, 2, 0, 1]);
+
+        let emitted = group_values.emit(EmitTo::All).unwrap();
+        assert_eq!(emitted.len(), 2);
+        // The Duration column round-trips as Duration on emit, not bare i64.
+        assert_eq!(
+            emitted[0].data_type(),
+            &DataType::Duration(TimeUnit::Microsecond)
+        );
+        let actual = emitted[0]
+            .as_any()
+            .downcast_ref::<DurationMicrosecondArray>()
+            .expect("emitted column should be a DurationMicrosecondArray");
+        // Three groups in first-seen order: 10, null, 20.
+        assert_eq!(actual.len(), 3);
+        assert_eq!(actual.value(0), 10);
+        assert!(actual.is_null(1));
+        assert_eq!(actual.value(2), 20);
+    }
+
     #[test]
     fn supported_schema_rejects_mix_of_supported_and_unsupported() {
         // One Float16 column among supported columns flips the whole

datafusion/sqllogictest/test_files/group_by.slt

@@ -5694,3 +5694,30 @@ drop table users_with_pk;
 
 statement ok
 drop table user_orders;
+
+# Test multi group by int + Duration
+statement ok
+CREATE TABLE duration_group_test AS VALUES
+  (1, arrow_cast(5, 'Duration(Second)')),
+  (1, arrow_cast(5, 'Duration(Second)')),
+  (1, arrow_cast(7, 'Duration(Second)')),
+  (2, arrow_cast(5, 'Duration(Second)'));
+
+# Single Duration group key ({5s, 7s}, 5s x3) via the GroupValuesPrimitive path.
+query ?I
+SELECT column2, count(*) FROM duration_group_test GROUP BY column2 ORDER BY column2;
+----
+0 days 0 hours 0 mins 5 secs 3
+0 days 0 hours 0 mins 7 secs 1
+
+# Multi-column GROUP BY: a primitive key and a Duration key on the same path.
+query I?I
+SELECT column1, column2, count(*)
+FROM duration_group_test GROUP BY column1, column2 ORDER BY column1, column2;
+----
+1 0 days 0 hours 0 mins 5 secs 2
+1 0 days 0 hours 0 mins 7 secs 1
+2 0 days 0 hours 0 mins 5 secs 1
+
+statement ok
+DROP TABLE duration_group_test;