From b93c342a9e42c9b7b962776160121a96dd247289 Mon Sep 17 00:00:00 2001 From: tohuya6 <201355151+tohuya6@users.noreply.github.com> Date: Sun, 19 Jul 2026 16:08:19 -0400 Subject: [PATCH] feat(physical-plan): add GroupColumn support for Float16 in multi-column GROUP BY `multi_group_by::group_column_supported_type` gates which GROUP BY columns can use the column-wise `GroupValuesColumn` fast path. Any unsupported column forces the entire grouping onto the byte-encoded `GroupValuesRows` fallback, so a single `Float16` key dragged an otherwise-qualifying multi-column GROUP BY onto the slow path. `Float16` reuses the existing `PrimitiveGroupValueBuilder` with no new builder type: its native `half::f16` already implements the `HashValue` canonicalization (`hash_float!(f16, f32, f64)`), so `-0.0` / `+0.0` folding and NaN grouping match Float32 / Float64 with no extra handling. - accept `Float16` in `group_column_supported_type` and dispatch it in `make_group_column` - move `Float16` from the rejected to the accepted set in the `group_column_supported_type` <-> `make_group_column` consistency fuzz, and repoint the two tests that used `Float16` as their stock "unsupported" example to a permanently-invalid unit combo (`Time64(Second)`), so they stay stable as sibling primitive builders (Decimal256, Interval, ...) land - add an end-to-end unit test (Float16 GROUP BY dedups including nulls, folds -0.0/+0.0 into one group stored as +0.0, groups equal NaNs, and preserves the Float16 output type) and a Float16 GROUP BY block in aggregate.slt - add a `(Float16, Int32)` group-count benchmark to `benches/multi_group_by.rs` (capped below f16's ~63.5k distinct finite values) Part of #22715 --- .../physical-plan/benches/multi_group_by.rs | 93 ++++++++++++++++++- .../group_values/multi_group_by/mod.rs | 89 +++++++++++++++--- .../sqllogictest/test_files/aggregate.slt | 33 +++++++ 3 files changed, 200 insertions(+), 15 deletions(-) diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 11c2800864316..99039d60b3463 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -27,7 +27,7 @@ //! covers a `(FixedSizeBinary, Int32)` key to exercise the //! `FixedSizeBinaryGroupValueBuilder`. -use arrow::array::{ArrayRef, Int32Array, UInt32Array}; +use arrow::array::{ArrayRef, Float16Array, Int32Array, UInt32Array}; use arrow::compute::take; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::util::bench_util::create_fsb_array; @@ -35,6 +35,7 @@ use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; use datafusion_physical_plan::aggregates::group_values::GroupValues; use datafusion_physical_plan::aggregates::group_values::GroupValuesRows; use datafusion_physical_plan::aggregates::group_values::multi_group_by::GroupValuesColumn; +use half::f16; use std::hint::black_box; use std::sync::Arc; @@ -444,6 +445,95 @@ fn bench_fixed_size_binary(c: &mut Criterion) { group.finish(); } +fn make_f16_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("f16", DataType::Float16, false), + Field::new("id", DataType::Int32, false), + ])) +} + +/// Generate `(Float16, Int32)` batches with `num_distinct_groups` distinct keys. +/// +/// `f16` has only ~63.5k finite values, so `num_distinct_groups` must stay well +/// under that (see `bench_float16`). Distinct keys are the low finite `f16` bit +/// patterns (NaN/inf skipped); the `Int32` column is keyed identically so the +/// combined cardinality equals `num_distinct_groups`. +fn generate_f16_batches( + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec> { + let pool: Vec = (0u16..) + .map(f16::from_bits) + .filter(|v| v.is_finite()) + .take(num_distinct_groups) + .collect(); + + 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 = Float16Array::from_iter_values(group_ids.clone().map(|g| pool[g])); + 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 `(Float16, Int32)` key. +/// +/// Exercises the primitive `GroupColumn` builder for `Float16` on the +/// multi-column path (previously such a schema fell back to `GroupValuesRows`). +/// Group counts are capped below `f16`'s ~63.5k distinct finite values. +fn bench_float16(c: &mut Criterion) { + let mut group = c.benchmark_group("float16"); + group.sample_size(15); + + let schema = make_f16_schema(); + + for num_groups in [1_000, 60_000] { + let batches = generate_f16_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 +543,6 @@ criterion_group!( bench_high_cardinality_scaling, bench_group_count_sweep, bench_fixed_size_binary, + bench_float16, ); 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..a3f9bd2e83208 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, Field, Float16Type, + 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; @@ -933,6 +933,7 @@ fn group_column_supported_type(data_type: &DataType) -> bool { | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 + | DataType::Float16 | DataType::Float32 | DataType::Float64 | DataType::Decimal128(_, _) @@ -986,6 +987,9 @@ fn make_group_column(field: &Field) -> Result> { DataType::UInt16 => instantiate_primitive!(v, nullable, UInt16Type, data_type), DataType::UInt32 => instantiate_primitive!(v, nullable, UInt32Type, data_type), DataType::UInt64 => instantiate_primitive!(v, nullable, UInt64Type, data_type), + DataType::Float16 => { + instantiate_primitive!(v, nullable, Float16Type, data_type) + } DataType::Float32 => { instantiate_primitive!(v, nullable, Float32Type, data_type) } @@ -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, Float16Array, 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; @@ -1283,7 +1290,7 @@ mod tests { /// /// This test fuzzes a representative cross-section of types and asserts /// both directions of the biconditional. When a new specialization is - /// added (`Float16`, `FixedSizeList`, `Struct`, ...) it should be added + /// added (`FixedSizeList`, `Struct`, `Decimal256`, ...) it should be added /// to the supported_cases vector; when a type is intentionally rejected /// it should be added to unsupported_cases. #[test] @@ -1294,6 +1301,7 @@ mod tests { DataType::UInt64, DataType::Float32, DataType::Float64, + DataType::Float16, DataType::Decimal128(38, 10), DataType::Utf8, DataType::LargeUtf8, @@ -1325,7 +1333,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 / @@ -1352,14 +1359,65 @@ mod tests { } } + /// End-to-end coverage for `Float16` group keys. Beyond routing through the + /// column-wise `GroupValuesColumn` fast path, this exercises the float + /// canonicalization the builder relies on (shared with Float32 / Float64 via + /// `HashValue`): `-0.0` and `+0.0` collapse into one group stored as `+0.0`, + /// and equal `NaN`s collapse into a single group. + #[test] + fn test_group_values_column_float16() { + use half::f16; + + let schema = + Arc::new(Schema::new(vec![Field::new("f", DataType::Float16, true)])); + assert!(supported_schema(&schema)); + let mut group_values = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + + // Rows: 1.0, -0.0, +0.0, NaN, 1.0, NaN, null, null. + // First-seen groups: 1.0 -> 0, -0.0/+0.0 -> 1, NaN -> 2, null -> 3. + let input: ArrayRef = Arc::new(Float16Array::from(vec![ + Some(f16::from_f32(1.0)), + Some(f16::from_f32(-0.0)), + Some(f16::from_f32(0.0)), + Some(f16::NAN), + Some(f16::from_f32(1.0)), + Some(f16::NAN), + None, + None, + ])); + let mut groups = Vec::new(); + group_values.intern(&[input], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 1, 2, 0, 2, 3, 3]); + + let emitted = group_values.emit(EmitTo::All).unwrap(); + assert_eq!(emitted.len(), 1); + assert_eq!(emitted[0].data_type(), &DataType::Float16); + let actual = emitted[0] + .as_any() + .downcast_ref::() + .expect("emitted column should be a Float16Array"); + assert_eq!(actual.len(), 4); + assert_eq!(actual.value(0), f16::from_f32(1.0)); + // The ±0.0 group is stored canonically as +0.0 (not -0.0). + assert_eq!(actual.value(1).to_bits(), f16::from_f32(0.0).to_bits()); + assert!(actual.value(2).is_nan()); + assert!(actual.is_null(3)); + } + #[test] fn supported_schema_rejects_mix_of_supported_and_unsupported() { - // One Float16 column among supported columns flips the whole - // schema to GroupValuesRows fallback. + // One permanently-unsupported column flips the whole schema to the + // GroupValuesRows fallback. Use an invalid Arrow unit combo + // (Time64(Second)) so this stays stable as new primitive builders land. let schema = Schema::new(vec![ Field::new("a", DataType::Int32, true), Field::new("b", DataType::Utf8, true), - Field::new("c", DataType::Float16, true), + Field::new( + "c", + DataType::Time64(arrow::datatypes::TimeUnit::Second), + true, + ), ]); assert!(!supported_schema(&schema)); @@ -1378,8 +1436,11 @@ mod tests { // rejected at construction time rather than at first `intern`. // `GroupValuesColumn` doesn't implement `Debug`, so explicit match // instead of `unwrap_err`. - let schema = - Arc::new(Schema::new(vec![Field::new("x", DataType::Float16, true)])); + let schema = Arc::new(Schema::new(vec![Field::new( + "x", + DataType::Time64(arrow::datatypes::TimeUnit::Second), + true, + )])); match GroupValuesColumn::::try_new(schema) { Ok(_) => panic!("expected NotImpl error, but try_new succeeded"), Err(e) => { diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index cbb9c5d0317dc..a29221bc5152d 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -2268,6 +2268,39 @@ statement ok DROP TABLE approx_distinct_duration_test; +# GROUP BY over Float16 group keys. In a multi-column GROUP BY a Float16 key +# routes through the GroupValuesColumn column-wise fast path (added here); a +# single-column GROUP BY uses the pre-existing single-column primitive path. +# Verify both, including that -0.0 / +0.0 fold into one group via the float +# canonicalization shared with Float32 / Float64. The count-of-groups form +# avoids depending on Float16's textual display. +statement ok +CREATE TABLE float16_group_test AS VALUES + (arrow_cast(1.5, 'Float16'), 1), + (arrow_cast(1.5, 'Float16'), 1), + (arrow_cast(2.5, 'Float16'), 2), + (arrow_cast(-0.0, 'Float16'), 3), + (arrow_cast(0.0, 'Float16'), 3); + +# Single-column GROUP BY (single-column primitive path): five rows collapse to +# three groups: 1.5, 2.5, and the merged -0.0 / +0.0. +query I +SELECT count(*) FROM (SELECT column1 FROM float16_group_test GROUP BY column1); +---- +3 + +# Multi-column GROUP BY: a Float16 key alongside a primitive key, exercising the +# GroupValuesColumn path this PR adds. column2 is aligned 1:1 with the Float16 +# value (the -0.0 / +0.0 rows share id 3), so the result is still three groups. +query I +SELECT count(*) FROM (SELECT column1, column2 FROM float16_group_test GROUP BY column1, column2); +---- +3 + +statement ok +DROP TABLE float16_group_test; + + # This test runs approx_distinct over the intervals YearMonth, # DayTime, MonthDayNano for the scalar and the grouped path. statement ok