diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 11c2800864316..3871ac7997607 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -27,9 +27,11 @@ //! covers a `(FixedSizeBinary, Int32)` key to exercise the //! `FixedSizeBinaryGroupValueBuilder`. -use arrow::array::{ArrayRef, Int32Array, UInt32Array}; +use arrow::array::{ArrayRef, Int32Array, IntervalMonthDayNanoArray, UInt32Array}; use arrow::compute::take; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::datatypes::{ + DataType, Field, IntervalMonthDayNano, IntervalUnit, Schema, SchemaRef, +}; 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 +446,91 @@ fn bench_fixed_size_binary(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> { + 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 8: 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::::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 +540,6 @@ criterion_group!( bench_high_cardinality_scaling, bench_group_count_sweep, bench_fixed_size_binary, + 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 f275d777c3279..4953ac10abd25 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 @@ -33,7 +33,8 @@ 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, + Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, IntervalDayTimeType, + IntervalMonthDayNanoType, IntervalUnit, IntervalYearMonthType, Schema, SchemaRef, StringViewType, Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, @@ -952,6 +953,7 @@ fn group_column_supported_type(data_type: &DataType) -> bool { | DataType::Time64(TimeUnit::Microsecond) | DataType::Time64(TimeUnit::Nanosecond) | DataType::Timestamp(_, _) + | DataType::Interval(_) | DataType::Utf8View | DataType::BinaryView | DataType::Boolean @@ -1031,6 +1033,19 @@ fn make_group_column(field: &Field) -> Result> { instantiate_primitive!(v, nullable, TimestampNanosecondType, 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) } @@ -1259,7 +1274,10 @@ enum Nulls { mod tests { use std::{collections::HashMap, sync::Arc}; - use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray, StringViewArray}; + use arrow::array::{ + Array, ArrayRef, Int64Array, IntervalMonthDayNanoArray, 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 +1327,9 @@ mod tests { DataType::Time64(arrow::datatypes::TimeUnit::Microsecond), DataType::Time64(arrow::datatypes::TimeUnit::Nanosecond), DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None), + DataType::Interval(arrow::datatypes::IntervalUnit::YearMonth), + DataType::Interval(arrow::datatypes::IntervalUnit::DayTime), + DataType::Interval(arrow::datatypes::IntervalUnit::MonthDayNano), ]; for dt in &supported_cases { @@ -1352,6 +1373,56 @@ mod tests { } } + /// Interval GROUP BY on the column-wise fast path: dedups (incl. nulls), + /// keeps "1 month" and "30 days" distinct (field-wise compare, no cross-unit + /// folding), and preserves the Interval output type. One unit exercises the + /// shared builder; the fuzz above covers routing for all three. + #[test] + fn test_group_values_column_interval() { + use arrow::datatypes::{IntervalMonthDayNano, IntervalUnit}; + + let schema = Arc::new(Schema::new(vec![Field::new( + "i", + DataType::Interval(IntervalUnit::MonthDayNano), + true, + )])); + assert!(supported_schema(&schema)); + let mut group_values = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + + // "1 month" and "30 days" are distinct Interval values, not folded + // together. Row 3 repeats row 0, row 4 repeats the null of row 1. + let one_month = IntervalMonthDayNano::new(1, 0, 0); + let thirty_days = IntervalMonthDayNano::new(0, 30, 0); + let input: ArrayRef = Arc::new(IntervalMonthDayNanoArray::from(vec![ + Some(one_month), + None, + Some(thirty_days), + Some(one_month), + 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 Interval type, not the bare native. + assert_eq!( + emitted[0].data_type(), + &DataType::Interval(IntervalUnit::MonthDayNano) + ); + let actual = emitted[0] + .as_any() + .downcast_ref::() + .expect("emitted column should be an IntervalMonthDayNanoArray"); + // Three groups in first-seen order: 1 month, null, 30 days. + assert_eq!(actual.len(), 3); + assert_eq!(actual.value(0), one_month); + assert!(actual.is_null(1)); + assert_eq!(actual.value(2), thirty_days); + } + #[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..9e149059ac25f 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -2297,6 +2297,39 @@ statement ok DROP TABLE approx_distinct_interval_test; +# GROUP BY over Interval group keys. Interval columns route through the +# GroupValuesColumn column-wise fast path, so an Interval key no longer forces +# the row-encoded fallback. Equal intervals collapse and "1 month" / "30 days" +# stay distinct (no cross-unit folding); a primitive and an Interval key also +# group together on the multi-column path. Queries project only the INT counts +# / keys so the expected output does not depend on interval Display formatting. +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}, with 1 month appearing 3 times. +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; + + ## This test executes the APPROX_PERCENTILE_CONT aggregation against the test ## data, asserting the estimated quantiles are ±5% their actual values.