diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs index 7310cf262dbf..12d2fa680a55 100644 --- a/datafusion/physical-plan/benches/multi_group_by.rs +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -27,7 +27,9 @@ //! covers a `(FixedSizeBinary, Int32)` key to exercise the //! `FixedSizeBinaryGroupValueBuilder`. -use arrow::array::{ArrayRef, DurationMicrosecondArray, Int32Array, UInt32Array}; +use arrow::array::{ + ArrayRef, DurationMicrosecondArray, Float16Array, Int32Array, UInt32Array, +}; use arrow::compute::take; use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; use arrow::util::bench_util::create_fsb_array; @@ -35,6 +37,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 +447,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, skipping NaN and inf. 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(); +} + fn make_duration_schema() -> SchemaRef { Arc::new(Schema::new(vec![ Field::new("dur", DataType::Duration(TimeUnit::Microsecond), false), @@ -537,6 +629,7 @@ criterion_group!( bench_high_cardinality_scaling, bench_group_count_sweep, bench_fixed_size_binary, + bench_float16, bench_duration, ); 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 abbbee4277aa..9dced09ed015 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 @@ -36,11 +36,12 @@ use arrow::compute::cast; use arrow::datatypes::{ BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type, DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType, - DurationSecondType, Field, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, - Int64Type, Schema, SchemaRef, StringViewType, Time32MillisecondType, - Time32SecondType, Time64MicrosecondType, Time64NanosecondType, TimeUnit, - TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, - TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, + DurationSecondType, 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; @@ -945,6 +946,7 @@ fn group_column_supported_type(data_type: &DataType) -> bool { | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 + | DataType::Float16 | DataType::Float32 | DataType::Float64 | DataType::Decimal128(_, _) @@ -999,6 +1001,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) } @@ -1295,8 +1300,8 @@ mod tests { use std::{collections::HashMap, sync::Arc}; use arrow::array::{ - Array, ArrayRef, DurationMicrosecondArray, Int64Array, RecordBatch, StringArray, - StringViewArray, + Array, ArrayRef, DurationMicrosecondArray, Float16Array, Int32Array, Int64Array, + RecordBatch, StringArray, StringViewArray, }; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::{compute::concat_batches, util::pretty::pretty_format_batches}; @@ -1581,6 +1586,7 @@ mod tests { DataType::UInt64, DataType::Float32, DataType::Float64, + DataType::Float16, DataType::Decimal128(38, 10), DataType::Utf8, DataType::LargeUtf8, @@ -1616,7 +1622,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 / @@ -1694,14 +1699,77 @@ mod tests { assert_eq!(actual.value(2), 20); } + // `(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] + fn test_group_values_column_float16() { + use half::f16; + + let schema = Arc::new(Schema::new(vec![ + Field::new("f", DataType::Float16, true), + Field::new("i", DataType::Int32, true), + ])); + assert!(supported_schema(&schema)); + let mut group_values = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + + let f: 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::from_f32(0.0)), + Some(f16::NAN), + Some(f16::NAN), + None, + None, + ])); + let i: ArrayRef = Arc::new(Int32Array::from(vec![ + Some(3), + Some(3), + Some(3), + Some(4), + Some(3), + Some(3), + Some(3), + Some(3), + ])); + let mut groups = Vec::new(); + group_values.intern(&[f, i], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 1, 2, 3, 3, 4, 4]); + + let emitted = group_values.emit(EmitTo::All).unwrap(); + assert_eq!(emitted.len(), 2); + assert_eq!(emitted[0].data_type(), &DataType::Float16); + let keys = emitted[0] + .as_any() + .downcast_ref::() + .expect("emitted column should be a Float16Array"); + assert_eq!(keys.len(), 5); + assert_eq!(keys.value(0), f16::from_f32(1.0)); + // The ±0.0 group is stored canonically as +0.0 (not -0.0). + assert_eq!(keys.value(1).to_bits(), f16::from_f32(0.0).to_bits()); + assert_eq!(keys.value(2).to_bits(), f16::from_f32(0.0).to_bits()); + assert!(keys.value(3).is_nan()); + assert!(keys.is_null(4)); + let ids = emitted[1] + .as_any() + .downcast_ref::() + .expect("emitted column should be an Int32Array"); + assert_eq!(ids.values().to_vec(), vec![3, 3, 4, 3, 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 unsupported column flips the whole schema to the GroupValuesRows + // fallback. Time64(Second) stays invalid 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)); @@ -1720,8 +1788,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) => { @@ -1808,7 +1879,7 @@ mod tests { // `emit(EmitTo::First(4))` calls can `take_n` without panicking. // The hashmap entries below reference group indices 0..=11, so the // single column builder needs at least 12 rows to back them. - let seed: ArrayRef = Arc::new(arrow::array::Int32Array::from(vec![0_i32; 12])); + let seed: ArrayRef = Arc::new(Int32Array::from(vec![0_i32; 12])); for row in 0..12 { group_values.group_values[0] .append_val(&seed, row) diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index b718d3df2074..cf18123fecfa 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -5721,3 +5721,27 @@ FROM duration_group_test GROUP BY column1, column2 ORDER BY column1, column2; statement ok DROP TABLE duration_group_test; + +# Test multi group by int + Float16 +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 Float16 group key ({1.5, 2.5, ±0.0}) via the GroupValuesPrimitive path. +query I +SELECT count(*) FROM (SELECT column1 FROM float16_group_test GROUP BY column1); +---- +3 + +# Multi-column GROUP BY: a primitive key and a Float16 key on the same path. +query I +SELECT count(*) FROM (SELECT column1, column2 FROM float16_group_test GROUP BY column1, column2); +---- +3 + +statement ok +DROP TABLE float16_group_test;