feat: add GroupColumn support for Decimal32/Decimal64

datafusion/physical-plan/benches/multi_group_by.rs

@@ -29,8 +29,9 @@
 //! `(List<Int32>, 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<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: Vec<usize> = (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::<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,
@@ -1026,5 +1132,6 @@ criterion_group!(
     bench_decimal256,
     bench_list_int,
     bench_list_utf8,
+    bench_narrow_decimals,
 );
 criterion_main!(benches);

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<Box<dyn GroupColumn>> {
                 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<T>(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::<false>::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::<PrimitiveArray<T>>()
+                .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::<PrimitiveArray<T>>()
+                .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::<Decimal32Type>(DataType::Decimal32(9, 2), [1000, 2000, 999_999_999]);
+        check::<Decimal64Type>(
+            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]

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;