diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 56570927623b0..4753fe516fe5c 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -29,8 +29,9 @@ //! `(List, Int32)` key to exercise the `ListGroupValueBuilder`. use arrow::array::{ - ArrayRef, Decimal256Array, DurationMicrosecondArray, Float16Array, Int32Array, - IntervalMonthDayNanoArray, ListArray, StringArray, UInt32Array, + ArrayRef, Decimal32Array, Decimal64Array, Decimal256Array, DurationMicrosecondArray, + Float16Array, Int32Array, IntervalMonthDayNanoArray, ListArray, StringArray, + UInt32Array, }; use arrow::buffer::OffsetBuffer; use arrow::compute::take; @@ -1011,6 +1012,111 @@ fn bench_list_utf8(c: &mut Criterion) { group.finish(); } +/// Generate `(Decimal32 | Decimal64, Int32)` batches with `num_distinct_groups` +/// distinct keys. +/// +/// The decimal column is built by `make_keys` so one generator serves both +/// narrow widths. The `Int32` column is keyed identically so the combined +/// cardinality equals `num_distinct_groups`. +fn generate_narrow_decimal_batches( + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, + make_keys: fn(&[usize]) -> ArrayRef, +) -> 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: Vec = (0..current_batch_size) + .map(|row| (batch_start + row) % num_distinct_groups) + .collect(); + + let id: Int32Array = group_ids.iter().map(|&g| g as i32).collect(); + + vec![make_keys(&group_ids), Arc::new(id) as ArrayRef] + }) + .collect() +} + +/// Experiment 14: Group count sweep for `(Decimal32, Int32)` and +/// `(Decimal64, Int32)` keys. +/// +/// Exercises the primitive `GroupColumn` builder for the two narrow decimal +/// widths (4-byte `i32` and 8-byte `i64` native) on the multi-column path +/// (previously such a schema fell back to `GroupValuesRows`). +fn bench_narrow_decimals(c: &mut Criterion) { + let cases: [(&str, DataType, fn(&[usize]) -> ArrayRef); 2] = [ + ("decimal32", DataType::Decimal32(9, 0), |group_ids| { + Arc::new( + Decimal32Array::from_iter_values(group_ids.iter().map(|&g| g as i32)) + .with_precision_and_scale(9, 0) + .unwrap(), + ) as ArrayRef + }), + ("decimal64", DataType::Decimal64(18, 0), |group_ids| { + Arc::new( + Decimal64Array::from_iter_values(group_ids.iter().map(|&g| g as i64)) + .with_precision_and_scale(18, 0) + .unwrap(), + ) as ArrayRef + }), + ]; + + for (name, data_type, make_keys) in cases { + let mut group = c.benchmark_group(name); + group.sample_size(15); + + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("dec", data_type, false), + Field::new("id", DataType::Int32, false), + ])); + + for num_groups in [1_000, 1_000_000] { + let batches = generate_narrow_decimal_batches( + num_groups, + 1_000_000, + DEFAULT_BATCH_SIZE, + make_keys, + ); + + 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, @@ -1026,5 +1132,6 @@ criterion_group!( bench_decimal256, bench_list_int, bench_list_utf8, + bench_narrow_decimals, ); 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 f10990a73af50..098b49dd4fa6b 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 @@ -38,13 +38,13 @@ use crate::aggregates::group_values::multi_group_by::{ }; use arrow::array::{Array, ArrayRef, BooleanBufferBuilder, new_empty_array}; use arrow::datatypes::{ - BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type, Decimal256Type, - DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType, - DurationSecondType, Field, Float16Type, Float32Type, Float64Type, Int8Type, - Int16Type, Int32Type, Int64Type, IntervalDayTimeType, IntervalMonthDayNanoType, - IntervalUnit, IntervalYearMonthType, Schema, SchemaRef, StringViewType, - Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType, - TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, + BinaryViewType, DataType, Date32Type, Date64Type, Decimal32Type, Decimal64Type, + Decimal128Type, Decimal256Type, DurationMicrosecondType, DurationMillisecondType, + DurationNanosecondType, DurationSecondType, Field, Float16Type, Float32Type, + Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, IntervalDayTimeType, + IntervalMonthDayNanoType, IntervalUnit, IntervalYearMonthType, Schema, SchemaRef, + StringViewType, Time32MillisecondType, Time32SecondType, Time64MicrosecondType, + Time64NanosecondType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; @@ -1036,6 +1036,12 @@ fn make_group_column(field: &Field) -> Result> { instantiate_primitive!(nullable, IntervalMonthDayNanoType, data_type) } }, + DataType::Decimal32(_, _) => { + instantiate_primitive!(nullable, Decimal32Type, data_type) + } + DataType::Decimal64(_, _) => { + instantiate_primitive!(nullable, Decimal64Type, data_type) + } DataType::Decimal128(_, _) => { instantiate_primitive!(nullable, Decimal128Type, data_type) } @@ -1333,7 +1339,9 @@ mod tests { Int32Array, Int64Array, PrimitiveArray, RecordBatch, StringArray, StringViewArray, UInt32Array, }; - use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use arrow::datatypes::{ + DataType, Decimal32Type, Decimal64Type, Field, Schema, SchemaRef, + }; use arrow::{ compute::{concat_batches, take}, util::pretty::pretty_format_batches, @@ -1633,6 +1641,8 @@ mod tests { DataType::Float32, DataType::Float64, DataType::Float16, + DataType::Decimal32(9, 2), + DataType::Decimal64(18, 2), DataType::Decimal128(38, 10), DataType::Decimal256(76, 10), DataType::Utf8, @@ -1684,6 +1694,16 @@ mod tests { )), ), DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Float16)), + // The narrow decimals are reachable through the dictionary + // recursion too, not just as top-level keys. + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Decimal32(9, 2)), + ), + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Decimal64(18, 2)), + ), ]; for dt in &supported_cases { @@ -1792,6 +1812,82 @@ mod tests { assert_eq!(actual.value(2), 20); } + // `Decimal32` / `Decimal64` group keys stay on the `GroupValuesColumn` fast + // path, dedup (including nulls), and round-trip with their declared + // precision and scale preserved -- not widened to `Decimal128`, and not + // emitted as the bare i32 / i64 storage type. + #[test] + fn test_group_values_column_narrow_decimals() { + // Both narrow widths go through the same primitive builder, so drive + // them from one body rather than testing only the first. + fn check(data_type: DataType, values: [T::Native; 3]) + where + T: arrow::datatypes::DecimalType + arrow::datatypes::ArrowPrimitiveType, + { + let schema = Arc::new(Schema::new(vec![ + Field::new("d", data_type.clone(), true), + Field::new("i", DataType::Int64, true), + ])); + assert!(supported_schema(&schema), "{data_type} not on column path"); + let mut group_values = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + + let (DataType::Decimal32(p, s) | DataType::Decimal64(p, s)) = data_type + else { + unreachable!("this test only covers the narrow decimals") + }; + // Row 3 repeats row 0 and row 4 repeats the null pair. Row 5 carries + // a value at the storage type's full width. + let d: ArrayRef = Arc::new( + [ + Some(values[0]), + None, + Some(values[1]), + Some(values[0]), + None, + Some(values[2]), + ] + .into_iter() + .collect::>() + .with_precision_and_scale(p, s) + .unwrap(), + ); + let i: ArrayRef = Arc::new(Int64Array::from(vec![ + Some(1), + None, + Some(2), + Some(1), + None, + Some(3), + ])); + let mut groups = Vec::new(); + group_values.intern(&[d, i], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 2, 0, 1, 3]); + + let emitted = group_values.emit(EmitTo::All).unwrap(); + assert_eq!(emitted.len(), 2); + // Precision and scale survive the round trip. + assert_eq!(emitted[0].data_type(), &data_type); + let actual = emitted[0] + .as_any() + .downcast_ref::>() + .expect("emitted column should keep its decimal type"); + // Four groups in first-seen order. + assert_eq!(actual.len(), 4); + assert_eq!(actual.value(0), values[0]); + assert!(actual.is_null(1)); + assert_eq!(actual.value(2), values[1]); + // The full-width value is not truncated by the narrower storage. + assert_eq!(actual.value(3), values[2]); + } + + check::(DataType::Decimal32(9, 2), [1000, 2000, 999_999_999]); + check::( + DataType::Decimal64(18, 2), + [1000, 2000, 999_999_999_999_999_999], + ); + } + // `(Float16, Int32)` keys: ±0.0 collapse (stored as +0.0), NaNs collapse, and // the Int32 key keeps `(0.0, 4)` distinct from `(±0.0, 3)`. #[test] diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 97539ea479d40..9eb1865ecf107 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -5880,3 +5880,45 @@ select k, d, count(*) from decimal256_multi_group group by k, d order by k, d; statement ok drop table decimal256_multi_group; + +# Test multi group by int + Decimal32 / Decimal64 +# +# Unlike the Decimal128 / Decimal256 cases above, these two widths are not +# reachable from a SQL `DECIMAL(p, s)` declaration -- that maps to Decimal128 or +# Decimal256 by precision -- so the key has to be built with arrow_cast. The +# arrow_typeof queries pin that the group key really is the narrow type, so +# these cases cannot silently degrade into more Decimal128 coverage. +statement ok +create table narrow_decimal_multi_group (k int, d decimal(9, 2)) as values + (1, 100.00), (1, 100.00), (1, 250.00), (2, 100.00), (2, NULL); + +query T +select distinct arrow_typeof(arrow_cast(d, 'Decimal32(9, 2)')) from narrow_decimal_multi_group; +---- +Decimal32(9, 2) + +query IRI +select k, arrow_cast(d, 'Decimal32(9, 2)'), count(*) +from narrow_decimal_multi_group group by 1, 2 order by 1, 2; +---- +1 100.00 2 +1 250.00 1 +2 100.00 1 +2 NULL 1 + +query T +select distinct arrow_typeof(arrow_cast(d, 'Decimal64(18, 2)')) from narrow_decimal_multi_group; +---- +Decimal64(18, 2) + +query IRI +select k, arrow_cast(d, 'Decimal64(18, 2)'), count(*) +from narrow_decimal_multi_group group by 1, 2 order by 1, 2; +---- +1 100.00 2 +1 250.00 1 +2 100.00 1 +2 NULL 1 + +statement ok +drop table narrow_decimal_multi_group;