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..dd9cde6db596c 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, @@ -1259,7 +1263,10 @@ enum Nulls { mod tests { use std::{collections::HashMap, sync::Arc}; - use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray, StringViewArray}; + use arrow::array::{ + Array, ArrayRef, Decimal256Array, 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; @@ -1295,6 +1302,7 @@ mod tests { DataType::Float32, DataType::Float64, DataType::Decimal128(38, 10), + DataType::Decimal256(76, 10), DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View, @@ -1326,7 +1334,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 @@ -1352,6 +1359,51 @@ mod tests { } } + /// Decimal256 GROUP BY on the column-wise fast path (mirroring Decimal128): + /// dedups (incl. nulls) and preserves the Decimal256 precision/scale output + /// type. Uses precision > 38 so it's genuinely Decimal256, not Decimal128. + #[test] + fn test_group_values_column_decimal256() { + use arrow::datatypes::i256; + + let schema = Arc::new(Schema::new(vec![Field::new( + "d", + DataType::Decimal256(50, 0), + true, + )])); + assert!(supported_schema(&schema)); + let mut group_values = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + + // Two distinct values and a null with repeats: row 3 repeats row 0, + // row 4 repeats the null of row 1. + let a = i256::from_i128(100); + let b = i256::from_i128(200); + let input: ArrayRef = Arc::new( + Decimal256Array::from(vec![Some(a), None, Some(b), Some(a), None]) + .with_precision_and_scale(50, 0) + .unwrap(), + ); + 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 Decimal256 type (precision + scale), not + // the bare i256 native. + assert_eq!(emitted[0].data_type(), &DataType::Decimal256(50, 0)); + let actual = emitted[0] + .as_any() + .downcast_ref::() + .expect("emitted column should be a Decimal256Array"); + // Three groups in first-seen order: 100, null, 200. + assert_eq!(actual.len(), 3); + assert_eq!(actual.value(0), a); + assert!(actual.is_null(1)); + assert_eq!(actual.value(2), b); + } + #[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 9400a09a5d4bf..2552fdda42755 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -2327,6 +2327,39 @@ statement ok DROP TABLE approx_distinct_interval_test; +# GROUP BY over Decimal256 group keys. Decimal256 columns route through the +# GroupValuesColumn column-wise fast path (like Decimal128), so a Decimal256 key +# no longer forces the row-encoded fallback. Equal decimals collapse and a +# Decimal256 key groups alongside a primitive on the multi-column path. Queries +# project only the INT counts / keys so the expected output does not depend on +# decimal Display formatting. +statement ok +CREATE TABLE decimal256_group_test AS VALUES + (1, arrow_cast(100, 'Decimal256(20, 2)')), + (1, arrow_cast(100, 'Decimal256(20, 2)')), + (1, arrow_cast(250, 'Decimal256(20, 2)')), + (2, arrow_cast(100, 'Decimal256(20, 2)')); + +# Single Decimal256 group key: {100.00, 250.00}, with 100.00 appearing three times. +query I +SELECT count(*) FROM decimal256_group_test GROUP BY column2 ORDER BY count(*); +---- +1 +3 + +# Multi-column GROUP BY: a primitive key and a Decimal256 key on the same path. +query II +SELECT column1, count(*) +FROM decimal256_group_test GROUP BY column1, column2 ORDER BY column1, count(*); +---- +1 1 +1 2 +2 1 + +statement ok +DROP TABLE decimal256_group_test; + + ## This test executes the APPROX_PERCENTILE_CONT aggregation against the test ## data, asserting the estimated quantiles are ±5% their actual values.