From 9f0bb3eba1e07626b6fc448575c70ef1c5c97cab Mon Sep 17 00:00:00 2001 From: tohuya6 <201355151+tohuya6@users.noreply.github.com> Date: Sun, 19 Jul 2026 14:56:51 -0400 Subject: [PATCH] feat(physical-plan): add GroupColumn support for Duration in multi-column GROUP BY `multi_group_by::group_column_supported_type` gates which GROUP BY columns can use the column-wise `GroupValuesColumn` fast path. Any unsupported column forces the entire grouping onto the byte-encoded `GroupValuesRows` fallback, so a single `Duration` key dragged an otherwise-qualifying multi-column GROUP BY onto the slow path. `Duration` shares the `i64` native representation of `Timestamp`, so it reuses the existing `PrimitiveGroupValueBuilder` with no new builder type: - dispatch the four `Duration*Type` units in `make_group_column` - accept `Duration(_)` in `group_column_supported_type` (all four units are valid Arrow types, unlike Time32/Time64) - extend the `group_column_supported_type` <-> `make_group_column` consistency fuzz with the four Duration units - add an end-to-end unit test (Duration GROUP BY dedups including nulls and preserves the Duration output type) and a Duration GROUP BY block (single- and multi-column keys) in aggregate.slt - add a `(Duration, Int32)` group-count benchmark to `benches/multi_group_by.rs` Part of #22715 --- .../physical-plan/benches/multi_group_by.rs | 87 +++++++++++++++++- .../group_values/multi_group_by/mod.rs | 89 +++++++++++++++++-- .../sqllogictest/test_files/aggregate.slt | 31 +++++++ 3 files changed, 198 insertions(+), 9 deletions(-) diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 11c2800864316..63dc399072ef9 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/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,88 @@ 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> { + 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 8: 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::::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 +535,6 @@ criterion_group!( bench_high_cardinality_scaling, bench_group_count_sweep, bench_fixed_size_binary, + bench_duration, ); 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 f275d777c3279..88f56f5eeff03 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 @@ -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,9 @@ fn group_column_supported_type(data_type: &DataType) -> bool { | DataType::Time64(TimeUnit::Microsecond) | DataType::Time64(TimeUnit::Nanosecond) | DataType::Timestamp(_, _) + // All four Duration units are valid Arrow types (unlike Time32 / + // Time64), so every unit is handled by the dispatcher below. + | DataType::Duration(_) | DataType::Utf8View | DataType::BinaryView | DataType::Boolean @@ -1031,6 +1035,20 @@ fn make_group_column(field: &Field) -> Result> { 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 +1277,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 +1330,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 +1377,56 @@ mod tests { } } + /// End-to-end coverage for `Duration` group keys: a `Duration` column stays + /// on the `GroupValuesColumn` fast path, deduplicates equal durations + /// (including nulls), and round-trips with its `Duration` output type + /// preserved (not the bare `i64` native). The four units share one builder, + /// so a single unit exercises the dispatcher; the fuzz above covers routing + /// for every unit. + #[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, + )])); + assert!(supported_schema(&schema)); + let mut group_values = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + + // Two distinct durations and a null with repeats: row 3 repeats row 0, + // row 4 repeats the null of row 1. + let input: ArrayRef = Arc::new(DurationMicrosecondArray::from(vec![ + Some(10), + None, + Some(20), + Some(10), + None, + ])); + let mut groups = Vec::new(); + group_values.intern(&[input], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 2, 0, 1]); + + let emitted = group_values.emit(EmitTo::All).unwrap(); + assert_eq!(emitted.len(), 1); + // The emitted key keeps its Duration type, not the bare i64 native. + assert_eq!( + emitted[0].data_type(), + &DataType::Duration(TimeUnit::Microsecond) + ); + let actual = emitted[0] + .as_any() + .downcast_ref::() + .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 diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index cbb9c5d0317dc..2111567e4f758 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -2268,6 +2268,37 @@ statement ok DROP TABLE approx_distinct_duration_test; +# GROUP BY over Duration group keys. Duration columns route through the +# GroupValuesColumn column-wise fast path, so a Duration key no longer forces +# the row-encoded fallback. Assert equal durations collapse and a Duration key +# groups alongside a primitive on the multi-column path. +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}, with 5s appearing three times. +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; + + # This test runs approx_distinct over the intervals YearMonth, # DayTime, MonthDayNano for the scalar and the grouped path. statement ok