diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 11c2800864316..2056d96a9b4bc 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, Decimal256Array, Int32Array, UInt32Array}; use arrow::compute::take; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef, i256}; 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,92 @@ fn bench_fixed_size_binary(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)`; 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> { + 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 8: 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::::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 +539,6 @@ criterion_group!( bench_high_cardinality_scaling, bench_group_count_sweep, bench_fixed_size_binary, + 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 f275d777c3279..b070b6f1874a7 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,12 @@ 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, Decimal256Type, + 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; @@ -936,6 +936,7 @@ fn group_column_supported_type(data_type: &DataType) -> bool { | DataType::Float32 | DataType::Float64 | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) | DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary @@ -1034,6 +1035,9 @@ fn make_group_column(field: &Field) -> Result> { 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::::new( OutputType::Utf8, @@ -1295,6 +1299,7 @@ mod tests { DataType::Float32, DataType::Float64, DataType::Decimal128(38, 10), + DataType::Decimal256(76, 10), DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View, @@ -1326,7 +1331,6 @@ mod tests { let unsupported_cases: Vec = vec![ DataType::Float16, - 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/decimal.slt b/datafusion/sqllogictest/test_files/decimal.slt index 4335ec06685f2..5cd04bf5c2686 100644 --- a/datafusion/sqllogictest/test_files/decimal.slt +++ b/datafusion/sqllogictest/test_files/decimal.slt @@ -723,6 +723,21 @@ true 2 statement ok drop table decimal256_simple; +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; + # https://github.com/apache/datafusion/issues/12870 query R