This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23786-9051efdb87cd417b632e69759ed65e119bd06360 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 62650eff81847928fadaf1d5e8ba2303cb13b819 Author: Varun <[email protected]> AuthorDate: Mon Aug 3 03:39:38 2026 -0400 feat: add GroupColumn support for Interval in multi-column GROUP BY (#23786) ## Which issue does this PR close? - Part of #22715. ## Rationale for this change `multi_group_by::group_column_supported_type` gates which GROUP BY columns may use the column-wise `GroupValuesColumn` fast path, and the gate is all-or-nothing: a single unsupported column forces the **entire** grouping onto the byte-encoded `GroupValuesRows` fallback, even when every other key column would have qualified. An `Interval` key triggers exactly that today. All three `Interval` units reuse the existing `PrimitiveGroupValueBuilder` with no new builder type and no new `HashValue` impl — the `i32` / `IntervalDayTime` / `IntervalMonthDayNano` natives already implement it. ## What changes are included in this PR? - Dispatch the three `Interval*Type` units in `make_group_column`. `IntervalUnit` has exactly three variants, so the match is exhaustive with no fallback arm (unlike `Time32`/`Time64`, which must reject invalid unit combinations). - Accept `Interval(_)` in `group_column_supported_type`. - Extend the `group_column_supported_type` ↔ `make_group_column` consistency fuzz with the three Interval units. - Add an `(Interval, Int32)` group-count benchmark to `benches/multi_group_by.rs`. ## Are these changes tested? Yes. - New unit test: an `Interval` key stays on the `GroupValuesColumn` path, dedups including nulls, keeps "1 month" and "30 days" as **distinct** groups (no cross-unit folding), and round-trips with the `Interval` output type preserved. - The consistency fuzz now asserts all three Interval units route through the dispatcher. - New single- and multi-column `Interval` `GROUP BY` coverage in `group_by.slt`. ## Are there any user-facing changes? No API changes. `GROUP BY` queries with an `Interval` key now use the column-wise fast path instead of the row-encoded fallback; results are unchanged. --------- Co-authored-by: tohuya6 <[email protected]> --- datafusion/physical-plan/benches/multi_group_by.rs | 94 +++++++++++++++++++++- .../aggregates/group_values/multi_group_by/mod.rs | 81 ++++++++++++++++++- datafusion/sqllogictest/test_files/group_by.slt | 27 +++++++ 3 files changed, 198 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 12d2fa680a..11481d4f91 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/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); diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 9dced09ed0..8b68152c47 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/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 diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index cf18123fec..afde5b9331 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/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; --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
