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-23849-e0a05a693da6c56f97f001e0d5b3df1329b11ac2 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 038bfa22a5cb32fb9dd68c335d4053e6d377d1e8 Author: Varun <[email protected]> AuthorDate: Sat Aug 8 02:10:44 2026 -0400 feat: add GroupColumn support for Decimal256 in multi-column GROUP BY (#23849) ## Which issue does this PR close? - Part of #22715. ## Rationale for this change A `Decimal256` GROUP BY key currently forces the whole grouping onto the row-encoded `GroupValuesRows` fallback. `i256` is `Copy` and already implements `ArrowNativeTypeOp` + `HashValue`, so it reuses the existing `PrimitiveGroupValueBuilder`, exactly like `Decimal128`. ## What changes are included in this PR? - Support `Decimal256` in `make_group_column` and `group_column_supported_type`. - Update the supported-type consistency test and add `Decimal256` benchmark. ## Are these changes tested? Yes, the consistency test + multi-column `Decimal256` GROUP BY (with NULL key) in `group_by.slt`. ## Are there any user-facing changes? No. --------- Co-authored-by: tohuya6 <[email protected]> --- datafusion/physical-plan/benches/multi_group_by.rs | 92 +++++++++++++++++++++- .../aggregates/group_values/multi_group_by/mod.rs | 8 +- datafusion/sqllogictest/test_files/group_by.slt | 16 ++++ 3 files changed, 113 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 11481d4f91..0c689f9fcb 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -28,12 +28,13 @@ //! `FixedSizeBinaryGroupValueBuilder`. use arrow::array::{ - ArrayRef, DurationMicrosecondArray, Float16Array, Int32Array, + ArrayRef, Decimal256Array, DurationMicrosecondArray, Float16Array, Int32Array, IntervalMonthDayNanoArray, UInt32Array, }; use arrow::compute::take; use arrow::datatypes::{ DataType, Field, IntervalMonthDayNano, IntervalUnit, Schema, SchemaRef, TimeUnit, + i256, }; use arrow::util::bench_util::create_fsb_array; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; @@ -709,6 +710,94 @@ fn bench_interval(c: &mut Criterion) { group.finish(); } +fn make_decimal256_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("dec", DataType::Decimal256(50, 0), false), + Field::new("id", DataType::Int32, false), + ])) +} + +/// Generate `(Decimal256(50, 0), Int32)` batches with `num_distinct_groups` +/// distinct keys. +/// +/// Each distinct value is `i256::from_i128(g)`, and precision > 38 keeps it a +/// genuine `Decimal256`. The `Int32` column is keyed identically so the combined +/// cardinality equals `num_distinct_groups`. +fn generate_decimal256_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 = Decimal256Array::from_iter_values( + group_ids.clone().map(|g| i256::from_i128(g as i128)), + ) + .with_precision_and_scale(50, 0) + .unwrap(); + let id: Int32Array = group_ids.map(|g| g as i32).collect(); + + vec![Arc::new(keys) as ArrayRef, Arc::new(id) as ArrayRef] + }) + .collect() +} + +/// Experiment 11: Group count sweep for a `(Decimal256, Int32)` key. +/// +/// Exercises the primitive `GroupColumn` builder for `Decimal256` (32-byte +/// `i256` native) on the multi-column path (previously such a schema fell back +/// to `GroupValuesRows`). +fn bench_decimal256(c: &mut Criterion) { + let mut group = c.benchmark_group("decimal256"); + group.sample_size(15); + + let schema = make_decimal256_schema(); + + for num_groups in [1_000, 1_000_000] { + let batches = + generate_decimal256_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, @@ -721,5 +810,6 @@ criterion_group!( bench_float16, bench_duration, bench_interval, + bench_decimal256, ); 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 8b68152c47..a3ac23c15e 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 @@ -34,7 +34,7 @@ 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, + BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type, Decimal256Type, DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType, DurationSecondType, Field, Float16Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, IntervalDayTimeType, IntervalMonthDayNanoType, @@ -951,6 +951,7 @@ fn group_column_supported_type(data_type: &DataType) -> bool { | DataType::Float32 | DataType::Float64 | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) | DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary @@ -1081,6 +1082,9 @@ fn make_group_column(field: &Field) -> Result<Box<dyn GroupColumn>> { DataType::Decimal128(_, _) => { instantiate_primitive!(v, nullable, Decimal128Type, data_type) } + DataType::Decimal256(_, _) => { + instantiate_primitive!(v, nullable, Decimal256Type, data_type) + } DataType::Utf8 => { v.push(Box::new(ByteGroupValueBuilder::<i32>::new( OutputType::Utf8, @@ -1603,6 +1607,7 @@ mod tests { DataType::Float64, DataType::Float16, DataType::Decimal128(38, 10), + DataType::Decimal256(76, 10), DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View, @@ -1640,7 +1645,6 @@ mod tests { } let unsupported_cases: Vec<DataType> = vec![ - DataType::Decimal256(76, 10), // Invalid Time-unit combinations: Time32 is defined only for // Second / Millisecond and Time64 only for Microsecond / // Nanosecond. The TimeUnit enum allows constructing the other diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 637b7c1882..38d1b78214 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -5765,3 +5765,19 @@ FROM interval_group_test GROUP BY column1, column2 ORDER BY column1, count(*); statement ok DROP TABLE interval_group_test; + +# Test multi group by int + Decimal256 +statement ok +create table decimal256_multi_group (k int, d decimal(50, 2)) as values + (1, 100.00), (1, 100.00), (1, 250.00), (2, 100.00), (2, NULL); + +query IRI +select k, d, count(*) from decimal256_multi_group group by k, d order by k, d; +---- +1 100 2 +1 250 1 +2 100 1 +2 NULL 1 + +statement ok +drop table decimal256_multi_group; --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
