This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 62650eff81 feat: add GroupColumn support for Interval in multi-column
GROUP BY (#23786)
62650eff81 is described below
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]
