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

datafusion/physical-plan/benches/multi_group_by.rs

@@ -28,10 +28,13 @@
 //! `FixedSizeBinaryGroupValueBuilder`.
 
 use arrow::array::{
-    ArrayRef, DurationMicrosecondArray, Float16Array, Int32Array, UInt32Array,
+    ArrayRef, DurationMicrosecondArray, Float16Array, Int32Array,
+    IntervalMonthDayNanoArray, UInt32Array,
 };
 use arrow::compute::take;
-use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
+use arrow::datatypes::{
+    DataType, Field, IntervalMonthDayNano, IntervalUnit, 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;
@@ -620,6 +623,92 @@ fn bench_duration(c: &mut Criterion) {
     group.finish();
 }
 
+fn make_interval_schema() -> SchemaRef {
+    Arc::new(Schema::new(vec![
+        Field::new("iv", DataType::Interval(IntervalUnit::MonthDayNano), false),
+        Field::new("id", DataType::Int32, false),
+    ]))
+}
+
+/// Generate `(Interval(MonthDayNano), Int32)` batches with `num_distinct_groups`
+/// distinct keys.
+///
+/// Each distinct interval is `MonthDayNano(g, 0, 0)`. The `Int32` column is
+/// keyed identically so the combined cardinality equals `num_distinct_groups`.
+fn generate_interval_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 = IntervalMonthDayNanoArray::from_iter_values(
+                group_ids
+                    .clone()
+                    .map(|g| IntervalMonthDayNano::new(g as i32, 0, 0)),
+            );
+            let id: Int32Array = group_ids.map(|g| g as i32).collect();
+
+            vec![Arc::new(keys) as ArrayRef, Arc::new(id) as ArrayRef]
+        })
+        .collect()
+}
+
+/// Experiment 10: Group count sweep for an `(Interval, Int32)` key.
+///
+/// Exercises the primitive `GroupColumn` builder for `Interval` on the
+/// multi-column path (previously such a schema fell back to `GroupValuesRows`).
+fn bench_interval(c: &mut Criterion) {
+    let mut group = c.benchmark_group("interval");
+    group.sample_size(15);
+
+    let schema = make_interval_schema();
+
+    for num_groups in [1_000, 1_000_000] {
+        let batches =
+            generate_interval_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,
@@ -631,5 +720,6 @@ criterion_group!(
     bench_fixed_size_binary,
     bench_float16,
     bench_duration,
+    bench_interval,
 );
 criterion_main!(benches);

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

@@ -37,7 +37,8 @@ use arrow::datatypes::{
     BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type,
     DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType,
     DurationSecondType, Field, Float16Type, Float32Type, Float64Type, Int8Type,
-    Int16Type, Int32Type, Int64Type, Schema, SchemaRef, StringViewType,
+    Int16Type, Int32Type, Int64Type, IntervalDayTimeType, IntervalMonthDayNanoType,
+    IntervalUnit, IntervalYearMonthType, Schema, SchemaRef, StringViewType,
     Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType,
     TimeUnit, TimestampMicrosecondType, TimestampMillisecondType,
     TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, UInt32Type,
@@ -967,6 +968,7 @@ fn group_column_supported_type(data_type: &DataType) -> bool {
             | DataType::Time64(TimeUnit::Nanosecond)
             | DataType::Timestamp(_, _)
             | DataType::Duration(_)
+            | DataType::Interval(_)
             | DataType::Utf8View
             | DataType::BinaryView
             | DataType::Boolean
@@ -1063,6 +1065,19 @@ fn make_group_column(field: &Field) -> Result<Box<dyn GroupColumn>> {
                 instantiate_primitive!(v, nullable, DurationNanosecondType, data_type)
             }
         },
+        // `IntervalUnit` has exactly three variants, so this match is exhaustive
+        // with no fallback arm (unlike Time32 / Time64).
+        DataType::Interval(u) => match u {
+            IntervalUnit::YearMonth => {
+                instantiate_primitive!(v, nullable, IntervalYearMonthType, data_type)
+            }
+            IntervalUnit::DayTime => {
+                instantiate_primitive!(v, nullable, IntervalDayTimeType, data_type)
+            }
+            IntervalUnit::MonthDayNano => {
+                instantiate_primitive!(v, nullable, IntervalMonthDayNanoType, data_type)
+            }
+        },
         DataType::Decimal128(_, _) => {
             instantiate_primitive!(v, nullable, Decimal128Type, data_type)
         }
@@ -1301,7 +1316,7 @@ mod tests {
 
     use arrow::array::{
         Array, ArrayRef, DurationMicrosecondArray, Float16Array, Int32Array, Int64Array,
-        RecordBatch, StringArray, StringViewArray,
+        PrimitiveArray, RecordBatch, StringArray, StringViewArray,
     };
     use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
     use arrow::{compute::concat_batches, util::pretty::pretty_format_batches};
@@ -1606,6 +1621,9 @@ mod tests {
             DataType::Duration(arrow::datatypes::TimeUnit::Millisecond),
             DataType::Duration(arrow::datatypes::TimeUnit::Microsecond),
             DataType::Duration(arrow::datatypes::TimeUnit::Nanosecond),
+            DataType::Interval(arrow::datatypes::IntervalUnit::YearMonth),
+            DataType::Interval(arrow::datatypes::IntervalUnit::DayTime),
+            DataType::Interval(arrow::datatypes::IntervalUnit::MonthDayNano),
         ];
 
         for dt in &supported_cases {
@@ -1758,6 +1776,65 @@ mod tests {
         assert_eq!(ids.values().to_vec(), vec![3, 3, 4, 3, 3]);
     }
 
+    // `(Interval, Int32)` keys for each of the three interval units: null keys
+    // dedup, the Int32 key splits equal intervals, and emit gives back Interval.
+    #[test]
+    fn test_group_values_column_interval() {
+        use arrow::datatypes::{
+            ArrowPrimitiveType, IntervalDayTime, IntervalDayTimeType,
+            IntervalMonthDayNano, IntervalMonthDayNanoType, IntervalUnit,
+            IntervalYearMonthType,
+        };
+
+        fn check<T: ArrowPrimitiveType>(unit: IntervalUnit, value: T::Native) {
+            let schema = Arc::new(Schema::new(vec![
+                Field::new("i", DataType::Interval(unit), true),
+                Field::new("n", DataType::Int32, true),
+            ]));
+            assert!(supported_schema(&schema), "{unit:?} schema not supported");
+            let mut group_values =
+                GroupValuesColumn::<false>::try_new(Arc::clone(&schema)).unwrap();
+
+            let i: ArrayRef = Arc::new(PrimitiveArray::<T>::from_iter([
+                Some(value),
+                None,
+                Some(value),
+                None,
+                Some(value),
+            ]));
+            let n: ArrayRef = Arc::new(Int32Array::from(vec![3, 3, 3, 3, 4]));
+            let mut groups = Vec::new();
+            group_values.intern(&[i, n], &mut groups).unwrap();
+            assert_eq!(groups, vec![0, 1, 0, 1, 2], "{unit:?}");
+
+            let emitted = group_values.emit(EmitTo::All).unwrap();
+            assert_eq!(emitted.len(), 2);
+            // The emitted key keeps its Interval type, not the bare native.
+            assert_eq!(emitted[0].data_type(), &DataType::Interval(unit));
+            let actual = emitted[0]
+                .as_any()
+                .downcast_ref::<PrimitiveArray<T>>()
+                .unwrap_or_else(|| panic!("emitted column should be a {unit:?} array"));
+            // Three groups in first-seen order: value, null, value (n=4).
+            assert_eq!(actual.len(), 3, "{unit:?}");
+            assert_eq!(actual.value(0), value, "{unit:?}");
+            assert!(actual.is_null(1), "{unit:?}");
+            assert_eq!(actual.value(2), value, "{unit:?}");
+            let ids = emitted[1]
+                .as_any()
+                .downcast_ref::<Int32Array>()
+                .expect("emitted column should be an Int32Array");
+            assert_eq!(ids.values().to_vec(), vec![3, 3, 4], "{unit:?}");
+        }
+
+        check::<IntervalYearMonthType>(IntervalUnit::YearMonth, 13);
+        check::<IntervalDayTimeType>(IntervalUnit::DayTime, IntervalDayTime::new(1, 500));
+        check::<IntervalMonthDayNanoType>(
+            IntervalUnit::MonthDayNano,
+            IntervalMonthDayNano::new(1, 0, 0),
+        );
+    }
+
     #[test]
     fn supported_schema_rejects_mix_of_supported_and_unsupported() {
         // One unsupported column flips the whole schema to the GroupValuesRows

datafusion/sqllogictest/test_files/group_by.slt

@@ -5745,3 +5745,30 @@ SELECT count(*) FROM (SELECT column1, column2 FROM float16_group_test GROUP BY c
 
 statement ok
 DROP TABLE float16_group_test;
+
+# Test multi group by int + Interval
+statement ok
+CREATE TABLE interval_group_test AS VALUES
+  (1, INTERVAL '1' MONTH),
+  (1, INTERVAL '1' MONTH),
+  (1, INTERVAL '30' DAY),
+  (2, INTERVAL '1' MONTH);
+
+# Single Interval group key ({1 month, 30 days}) via the GroupValuesPrimitive path.
+query I
+SELECT count(*) FROM interval_group_test GROUP BY column2 ORDER BY count(*);
+----
+1
+3
+
+# Multi-column GROUP BY: a primitive key and an Interval key on the same path.
+query II
+SELECT column1, count(*)
+FROM interval_group_test GROUP BY column1, column2 ORDER BY column1, count(*);
+----
+1 1
+1 2
+2 1
+
+statement ok
+DROP TABLE interval_group_test;