From e8b23e023efe67dc9e05018d1d7a2bc923d7a3b7 Mon Sep 17 00:00:00 2001 From: Qi Zhu Date: Wed, 29 Jul 2026 14:24:41 +0800 Subject: [PATCH 1/4] feat(physical-plan): generic Rows-backed GroupColumn keeps mixed schemas on the column-wise path (#23523) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Part of #22715 (nested type coverage in `GroupValuesColumn` EPIC) - Alternative to #23128 (per-type approach) — implements the direction @alamb proposed in - Step toward the terminal goal of retiring `GroupValuesRows` entirely (#23404) Today `GroupValuesColumn` is **all-or-nothing**: a single nested column in the GROUP BY key (\`Struct\`, \`List\`, \`FixedSizeList\`, …) makes \`supported_schema\` return \`false\` and drops the *entire* aggregation onto the row-wise \`GroupValuesRows\` fallback — even when every other column would have qualified for the column-wise fast path. For a \`GROUP BY int_col, struct_col\` shape, the \`int_col\` pays the row-encoded storage cost for no reason. Add \`RowsGroupColumn\`: a generic \`GroupColumn\` backed by a single-field \`RowConverter\`, wired in as the nested-type dispatch arm of \`group_column_supported_type\` / \`make_group_column\`. Native columns keep their type-specialized builders; the nested column pays row-encoding only for its one column. Gated to \`data_type.is_nested()\` so intentionally excluded scalar types (Float16, Decimal256) stay on \`GroupValuesRows\` and the \`group_column_supported_type\` ⇔ \`make_group_column\` invariant holds. Memory, measured with 4000 groups of \`8 × Int64 + 1 × FixedSizeList\` in \`mixed_schema_column_path_uses_less_memory_than_rows_fallback\`: | | Bytes | vs baseline | |-------------------------------------------------------|----------|-------------| | \`GroupValuesRows\` (today's fallback) | 1096 KB | 100% | | \`GroupValuesColumn\` + \`RowsGroupColumn\` fallback | 594 KB | **54.2%** | Speed: not benchmarked as a headline result — the wins come from native columns keeping their type-specialized \`equal_to\`/\`append_val\` fast paths instead of falling back to byte-encoded row comparisons. Yes: - Unit tests inside \`row_backed\`: FSL / Struct roundtrip, \`take_n\`, \`supports_type\` matches \`RowConverter::supports_fields\`. - \`mixed_schema_column_path_uses_less_memory_than_rows_fallback\` (mod.rs): the 54.2% memory claim + identical group assignment vs \`GroupValuesRows\`. - \`nested_float_edge_cases_match_rows_fallback\`: nested \`-0.0\` / \`NaN\` produce the same groupings as \`GroupValuesRows\` (the correctness invariant to watch, since hashing runs on the raw column and equality runs on the row bytes). - \`multi_batch_and_emit_first_matches_rows_fallback\`: multi-batch streaming intern + \`EmitTo::First\` + \`take_n\`. All 39 tests in \`aggregates::group_values\` pass. No — internal aggregation representation only. Same query results, lower memory footprint on mixed-schema GROUP BY keys. - Add coverage for any type \`RowConverter\` cannot encode (currently arrow-rs 59.x handles Map fine; \`supports_type\` delegates to \`RowConverter::supports_fields\` so it auto-tracks upstream). - Retire \`GroupValuesRows\` entirely once coverage is complete (#23404). (cherry picked from commit 68d587468073c1fc84ffe00db66f35d506dcc828) --- .../group_values/multi_group_by/mod.rs | 686 +++++++---- .../group_values/multi_group_by/row_backed.rs | 1014 +++++++++++++++++ .../src/aggregates/group_values/row.rs | 98 +- 3 files changed, 1587 insertions(+), 211 deletions(-) create mode 100644 datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs 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 f603839bee271..2923726ad1b66 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 @@ -21,6 +21,7 @@ mod boolean; mod bytes; pub mod bytes_view; pub mod primitive; +pub mod row_backed; use std::mem::{self, size_of}; @@ -28,11 +29,12 @@ use crate::aggregates::group_values::GroupValues; use crate::aggregates::group_values::multi_group_by::{ boolean::BooleanGroupValueBuilder, bytes::ByteGroupValueBuilder, bytes_view::ByteViewGroupValueBuilder, primitive::PrimitiveGroupValueBuilder, + row_backed::RowsGroupColumn, }; use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; use arrow::compute::cast; use arrow::datatypes::{ - BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type, Float32Type, + BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type, Field, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, Schema, SchemaRef, StringViewType, Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, @@ -285,6 +287,7 @@ impl GroupValuesColumn { /// Create a new instance of GroupValuesColumn if supported for the specified schema pub fn try_new(schema: SchemaRef) -> Result { let map = HashTable::with_capacity(0); + let group_values = Self::build_group_columns(&schema)?; Ok(Self { schema, map, @@ -292,12 +295,27 @@ impl GroupValuesColumn { emit_group_index_list_buffer: Vec::new(), vectorized_operation_buffers: VectorizedOperationBuffers::default(), map_size: 0, - group_values: vec![], + group_values, hashes_buffer: Default::default(), random_state: crate::aggregates::AGGREGATION_HASH_SEED, }) } + /// Build one fresh [`GroupColumn`] per field in the schema. + /// + /// Used at construction time (`try_new`) and to repopulate the column + /// vector after operations that drain it (`emit(EmitTo::All)`, + /// `clear_shrink`). Centralising it keeps the post-condition that + /// `self.group_values` always contains exactly one builder per schema + /// field outside of those transient drain points. + fn build_group_columns(schema: &Schema) -> Result>> { + let mut v: Vec> = Vec::with_capacity(schema.fields().len()); + for f in schema.fields().iter() { + v.push(make_group_column(f.as_ref())?); + } + Ok(v) + } + // ======================================================================== // Scalarized intern // ======================================================================== @@ -911,172 +929,192 @@ macro_rules! instantiate_primitive { }; } -impl GroupValues for GroupValuesColumn { - fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { - if self.group_values.is_empty() { - let mut v = Vec::with_capacity(cols.len()); - - for f in self.schema.fields().iter() { - let nullable = f.is_nullable(); - let data_type = f.data_type(); - match data_type { - &DataType::Int8 => { - instantiate_primitive!(v, nullable, Int8Type, data_type) - } - &DataType::Int16 => { - instantiate_primitive!(v, nullable, Int16Type, data_type) - } - &DataType::Int32 => { - instantiate_primitive!(v, nullable, Int32Type, data_type) - } - &DataType::Int64 => { - instantiate_primitive!(v, nullable, Int64Type, data_type) - } - &DataType::UInt8 => { - instantiate_primitive!(v, nullable, UInt8Type, data_type) - } - &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::Float32 => { - instantiate_primitive!(v, nullable, Float32Type, data_type) - } - &DataType::Float64 => { - instantiate_primitive!(v, nullable, Float64Type, data_type) - } - &DataType::Date32 => { - instantiate_primitive!(v, nullable, Date32Type, data_type) - } - &DataType::Date64 => { - instantiate_primitive!(v, nullable, Date64Type, data_type) - } - &DataType::Time32(t) => match t { - TimeUnit::Second => { - instantiate_primitive!( - v, - nullable, - Time32SecondType, - data_type - ) - } - TimeUnit::Millisecond => { - instantiate_primitive!( - v, - nullable, - Time32MillisecondType, - data_type - ) - } - _ => {} - }, - &DataType::Time64(t) => match t { - TimeUnit::Microsecond => { - instantiate_primitive!( - v, - nullable, - Time64MicrosecondType, - data_type - ) - } - TimeUnit::Nanosecond => { - instantiate_primitive!( - v, - nullable, - Time64NanosecondType, - data_type - ) - } - _ => {} - }, - &DataType::Timestamp(t, _) => match t { - TimeUnit::Second => { - instantiate_primitive!( - v, - nullable, - TimestampSecondType, - data_type - ) - } - TimeUnit::Millisecond => { - instantiate_primitive!( - v, - nullable, - TimestampMillisecondType, - data_type - ) - } - TimeUnit::Microsecond => { - instantiate_primitive!( - v, - nullable, - TimestampMicrosecondType, - data_type - ) - } - TimeUnit::Nanosecond => { - instantiate_primitive!( - v, - nullable, - TimestampNanosecondType, - data_type - ) - } - }, - &DataType::Decimal128(_, _) => { - instantiate_primitive! { - v, - nullable, - Decimal128Type, - data_type - } - } - &DataType::Utf8 => { - let b = ByteGroupValueBuilder::::new(OutputType::Utf8); - v.push(Box::new(b) as _) - } - &DataType::LargeUtf8 => { - let b = ByteGroupValueBuilder::::new(OutputType::Utf8); - v.push(Box::new(b) as _) - } - &DataType::Binary => { - let b = ByteGroupValueBuilder::::new(OutputType::Binary); - v.push(Box::new(b) as _) - } - &DataType::LargeBinary => { - let b = ByteGroupValueBuilder::::new(OutputType::Binary); - v.push(Box::new(b) as _) - } - &DataType::Utf8View => { - let b = ByteViewGroupValueBuilder::::new(); - v.push(Box::new(b) as _) - } - &DataType::BinaryView => { - let b = ByteViewGroupValueBuilder::::new(); - v.push(Box::new(b) as _) - } - &DataType::Boolean => { - if nullable { - let b = BooleanGroupValueBuilder::::new(); - v.push(Box::new(b) as _) - } else { - let b = BooleanGroupValueBuilder::::new(); - v.push(Box::new(b) as _) - } - } - dt => { - return not_impl_err!("{dt} not supported in GroupValuesColumn"); - } - } +/// Returns true if the specified data type has a specialized +/// [`GroupColumn`] builder in [`make_group_column`]. +/// +/// This is the allow-list that gates the `GroupValuesRows` fallback in +/// [`crate::aggregates::group_values::new_group_values`]: it must accept +/// exactly the set of types that [`make_group_column`] constructs a +/// builder for. The `group_column_supported_type_matches_make_group_column` +/// test below pins this biconditional. +fn group_column_supported_type(data_type: &DataType) -> bool { + // Nested types (Struct / List / LargeList / FixedSizeList, recursively) have + // no type-specialized `GroupColumn`; they are handled by the generic + // row-backed fallback in `make_group_column` whenever arrow's row format can + // encode them. Gate the fallback to nested types so intentionally-excluded + // scalar types (e.g. Float16, Decimal256) stay on `GroupValuesRows` and the + // `group_column_supported_type` ⇔ `make_group_column` invariant holds. + if data_type.is_nested() { + return RowsGroupColumn::supports_type(data_type); + } + matches!( + *data_type, + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::Float32 + | DataType::Float64 + | DataType::Decimal128(_, _) + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Binary + | DataType::LargeBinary + | DataType::Date32 + | DataType::Date64 + // Only the semantically valid Time variants per the Arrow spec. + // The dispatcher in `make_group_column` returns NotImpl for the + // other unit combinations, so accepting them here would cause a + // schema to be routed into GroupValuesColumn and then fail at + // intern. Keep these two arms in lockstep with the dispatcher. + | DataType::Time32(TimeUnit::Second) + | DataType::Time32(TimeUnit::Millisecond) + | DataType::Time64(TimeUnit::Microsecond) + | DataType::Time64(TimeUnit::Nanosecond) + | DataType::Timestamp(_, _) + | DataType::Utf8View + | DataType::BinaryView + | DataType::Boolean + ) +} + +/// Build a [`GroupColumn`] for a single schema field. +/// +/// Extracted from the inline match that used to live in +/// [`GroupValuesColumn::intern`] so the per-field dispatch lives in one +/// place. This factory is the single source of truth for which Arrow types +/// map to which builder, and it is the function that future nested-type +/// specializations (e.g. `Struct`, `List`, `LargeList`) plug into without +/// having to enumerate every combination inline. +/// +/// Returns `Err(not_impl_err!(...))` for any type not in the supported set; +/// callers (`GroupValues::intern`) propagate that error so the +/// `GroupValuesRows` fallback can take over upstream of this builder. +/// +/// The allow-list that gates this dispatcher lives in +/// [`group_column_supported_type`] directly above. +fn make_group_column(field: &Field) -> Result> { + let nullable = field.is_nullable(); + let data_type = field.data_type(); + let mut v: Vec> = Vec::with_capacity(1); + match *data_type { + DataType::Int8 => instantiate_primitive!(v, nullable, Int8Type, data_type), + DataType::Int16 => instantiate_primitive!(v, nullable, Int16Type, data_type), + DataType::Int32 => instantiate_primitive!(v, nullable, Int32Type, data_type), + DataType::Int64 => instantiate_primitive!(v, nullable, Int64Type, data_type), + DataType::UInt8 => instantiate_primitive!(v, nullable, UInt8Type, data_type), + 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::Float32 => { + instantiate_primitive!(v, nullable, Float32Type, data_type) + } + DataType::Float64 => { + instantiate_primitive!(v, nullable, Float64Type, data_type) + } + DataType::Date32 => instantiate_primitive!(v, nullable, Date32Type, data_type), + DataType::Date64 => instantiate_primitive!(v, nullable, Date64Type, data_type), + DataType::Time32(t) => match t { + TimeUnit::Second => { + instantiate_primitive!(v, nullable, Time32SecondType, data_type) + } + TimeUnit::Millisecond => { + instantiate_primitive!(v, nullable, Time32MillisecondType, data_type) + } + // Time32 with Microsecond / Nanosecond is not a valid Arrow type + // combination; reject explicitly so group_column_supported_type + // and this dispatcher stay in lockstep (see consistency fuzz below). + _ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"), + }, + DataType::Time64(t) => match t { + TimeUnit::Microsecond => { + instantiate_primitive!(v, nullable, Time64MicrosecondType, data_type) + } + TimeUnit::Nanosecond => { + instantiate_primitive!(v, nullable, Time64NanosecondType, data_type) + } + // Time64 with Second / Millisecond is not a valid Arrow type + // combination; reject explicitly. + _ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"), + }, + DataType::Timestamp(t, _) => match t { + TimeUnit::Second => { + instantiate_primitive!(v, nullable, TimestampSecondType, data_type) + } + TimeUnit::Millisecond => { + instantiate_primitive!(v, nullable, TimestampMillisecondType, data_type) + } + TimeUnit::Microsecond => { + instantiate_primitive!(v, nullable, TimestampMicrosecondType, data_type) + } + TimeUnit::Nanosecond => { + instantiate_primitive!(v, nullable, TimestampNanosecondType, data_type) + } + }, + DataType::Decimal128(_, _) => { + instantiate_primitive!(v, nullable, Decimal128Type, data_type) + } + DataType::Utf8 => { + v.push(Box::new(ByteGroupValueBuilder::::new( + OutputType::Utf8, + ))); + } + DataType::LargeUtf8 => { + v.push(Box::new(ByteGroupValueBuilder::::new( + OutputType::Utf8, + ))); + } + DataType::Binary => { + v.push(Box::new(ByteGroupValueBuilder::::new( + OutputType::Binary, + ))); + } + DataType::LargeBinary => { + v.push(Box::new(ByteGroupValueBuilder::::new( + OutputType::Binary, + ))); + } + DataType::Utf8View => { + v.push(Box::new(ByteViewGroupValueBuilder::::new())); + } + DataType::BinaryView => { + v.push(Box::new(ByteViewGroupValueBuilder::::new())); + } + DataType::Boolean => { + if nullable { + v.push(Box::new(BooleanGroupValueBuilder::::new())); + } else { + v.push(Box::new(BooleanGroupValueBuilder::::new())); } - self.group_values = v; } + // Generic fallback for nested types (Struct / List / LargeList / + // FixedSizeList, recursively) that lack a type-specialized builder but + // can be encoded by arrow's row format. This is what lets a mixed + // schema keep the column-wise fast path for its native columns instead + // of dropping the whole key onto `GroupValuesRows`. + ref dt if dt.is_nested() && RowsGroupColumn::supports_type(dt) => { + v.push(Box::new(RowsGroupColumn::try_new(dt.clone())?)); + } + _ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"), + } + debug_assert_eq!( + v.len(), + 1, + "make_group_column must push exactly one builder" + ); + Ok(v.into_iter().next().unwrap()) +} + +impl GroupValues for GroupValuesColumn { + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { + // `try_new` and the reset points in `emit` / `clear_shrink` keep + // `self.group_values` populated with one builder per schema field, + // so no lazy initialization is needed here. if !STREAMING { self.vectorized_intern(cols, groups) } else { @@ -1104,8 +1142,14 @@ impl GroupValues for GroupValuesColumn { fn emit(&mut self, emit_to: EmitTo) -> Result> { let mut output = match emit_to { EmitTo::All => { - let group_values = mem::take(&mut self.group_values); - debug_assert!(self.group_values.is_empty()); + // Replace the column builders with a fresh set so the + // aggregator is immediately reusable after the drain. + // Same `self.schema` was already validated by `try_new`, + // so `build_group_columns` would only error here if some + // out-of-band schema mutation occurred — propagate it as + // a real Result rather than panicking. + let fresh = Self::build_group_columns(&self.schema)?; + let group_values = mem::replace(&mut self.group_values, fresh); group_values .into_iter() @@ -1204,7 +1248,12 @@ impl GroupValues for GroupValuesColumn { } fn clear_shrink(&mut self, num_rows: usize) { - self.group_values.clear(); + // Reset to a fresh column-builder vector. The schema was validated + // in `try_new`, so rebuilding cannot fail unless something else + // mutated the schema out-of-band — surface that as a panic since + // `clear_shrink` is infallible by trait signature. + self.group_values = Self::build_group_columns(&self.schema) + .expect("schema previously validated in try_new"); self.map.clear(); self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared self.map_size = self.map.capacity() * size_of::<(u64, usize)>(); @@ -1226,39 +1275,7 @@ pub fn supported_schema(schema: &Schema) -> bool { .fields() .iter() .map(|f| f.data_type()) - .all(supported_type) -} - -/// Returns true if the specified data type is supported by [`GroupValuesColumn`] -/// -/// In order to be supported, there must be a specialized implementation of -/// [`GroupColumn`] for the data type, instantiated in [`GroupValuesColumn::intern`] -fn supported_type(data_type: &DataType) -> bool { - matches!( - *data_type, - DataType::Int8 - | DataType::Int16 - | DataType::Int32 - | DataType::Int64 - | DataType::UInt8 - | DataType::UInt16 - | DataType::UInt32 - | DataType::UInt64 - | DataType::Float32 - | DataType::Float64 - | DataType::Decimal128(_, _) - | DataType::Utf8 - | DataType::LargeUtf8 - | DataType::Binary - | DataType::LargeBinary - | DataType::Date32 - | DataType::Date64 - | DataType::Time32(_) - | DataType::Timestamp(_, _) - | DataType::Utf8View - | DataType::BinaryView - | DataType::Boolean - ) + .all(group_column_supported_type) } ///Shows how many `null`s there are in an array @@ -1285,8 +1302,273 @@ mod tests { GroupValues, multi_group_by::GroupValuesColumn, }; - use super::{GroupIndexView, split_vec_min_alloc}; + use super::{ + GroupIndexView, group_column_supported_type, make_group_column, split_vec_min_alloc, + supported_schema, + }; + + /// A mixed group-by key of several native columns plus one nested column + /// that has no type-specialized `GroupColumn`. + /// + /// Before the generic row-backed fallback, `supported_schema` returned + /// `false` for this schema, so the *entire* key dropped to the row-wise + /// `GroupValuesRows`. Now only the nested column pays the row-encoding + /// cost; the native columns keep their compact column-wise storage. This + /// test proves both that (a) the results are identical and (b) the + /// column-wise path now uses less memory than the all-rows fallback. + #[test] + fn mixed_schema_column_path_uses_less_memory_than_rows_fallback() { + use crate::aggregates::group_values::GroupValuesRows; + use arrow::array::{FixedSizeListArray, Int64Array}; + use arrow::datatypes::Int64Type; + + // 8 native Int64 columns + 1 FixedSizeList ("embedding"). + let fsl_field = Arc::new(Field::new("item", DataType::Int64, true)); + let mut fields: Vec = (0..8) + .map(|i| Field::new(format!("k{i}"), DataType::Int64, false)) + .collect(); + fields.push(Field::new( + "emb", + DataType::FixedSizeList(Arc::clone(&fsl_field), 4), + true, + )); + let schema: SchemaRef = Arc::new(Schema::new(fields)); + + // The whole schema must now be eligible for the column-wise path. + assert!( + supported_schema(schema.as_ref()), + "mixed native + nested schema should be column-supported now" + ); + + // Build `n_groups` distinct rows (each row is its own group). + let n_groups = 4000usize; + let mut cols: Vec = (0..8) + .map(|c| { + let vals: Vec = + (0..n_groups).map(|r| (r as i64) * 8 + c as i64).collect(); + Arc::new(Int64Array::from(vals)) as ArrayRef + }) + .collect(); + let emb: Vec>>> = (0..n_groups) + .map(|r| { + Some(vec![ + Some(r as i64), + Some(r as i64 + 1), + Some(r as i64 + 2), + Some(r as i64 + 3), + ]) + }) + .collect(); + cols.push( + Arc::new(FixedSizeListArray::from_iter_primitive::( + emb, 4, + )) as ArrayRef, + ); + + // Intern the same data into both implementations. + let mut column_path = GroupValuesColumn::::try_new(Arc::clone(&schema)) + .expect("column path"); + let mut rows_path = + GroupValuesRows::try_new(Arc::clone(&schema)).expect("rows path"); + + let mut g1 = vec![]; + let mut g2 = vec![]; + column_path.intern(&cols, &mut g1).unwrap(); + rows_path.intern(&cols, &mut g2).unwrap(); + + // (a) Correctness: same number of groups and identical group assignment. + assert_eq!(column_path.len(), n_groups); + assert_eq!(rows_path.len(), n_groups); + assert_eq!(g1, g2, "group assignment must match the rows fallback"); + + // (b) Memory: the column-wise path stores the 8 native columns compactly + // and only row-encodes the nested one, so it should be smaller than + // encoding every column into rows. + // + // The delta is only printed here — a hard `column_size < rows_size` + // assert would be brittle to future Arrow row-format or memory- + // accounting changes without reflecting a grouping-correctness + // regression. Track the memory improvement via benchmarks instead. + let column_size = column_path.size(); + let rows_size = rows_path.size(); + println!( + "mixed-schema group values size: column-wise = {column_size} bytes, \ + all-rows fallback = {rows_size} bytes \ + ({:.1}% of fallback)", + 100.0 * column_size as f64 / rows_size as f64 + ); + + // Emitted values must be equal too (compare via the rows fallback which + // is the established reference implementation). + let out_col = column_path.emit(EmitTo::All).unwrap(); + let out_row = rows_path.emit(EmitTo::All).unwrap(); + assert_eq!(out_col.len(), out_row.len()); + for (a, b) in out_col.iter().zip(out_row.iter()) { + assert_eq!(a.as_ref(), b.as_ref()); + } + } + /// Relabel a group-index vector so labels are assigned in order of first + /// appearance. Two vectors are equivalent groupings iff their canonical + /// forms are equal — this ignores the (opaque, non-semantic) difference in + /// group-index numbering between the vectorized column path and the + /// sequential rows fallback. + /// + /// The [`GroupValues`] trait only guarantees that equal keys receive the + /// same group-id and that new keys receive a fresh id; the order in which + /// new ids are handed out is deliberately not part of the contract, and + /// can differ between correct implementations (e.g. because of internal + /// hash-map ordering). Canonicalizing before comparison is what lets us + /// assert equivalence across implementations. + fn canonical_grouping(groups: &[usize]) -> Vec { + let mut map = HashMap::new(); + let mut next = 0usize; + groups + .iter() + .map(|&g| { + *map.entry(g).or_insert_with(|| { + let v = next; + next += 1; + v + }) + }) + .collect() + } + + /// The generic row-backed column must be behavior-preserving: for the + /// nested columns it now handles, `GroupValuesColumn` must induce the same + /// grouping (partition of rows) as the established `GroupValuesRows` + /// fallback — including the float `-0.0` / `+0.0` / `NaN` edge cases decided + /// jointly by hashing and the row format. + #[test] + fn nested_float_edge_cases_match_rows_fallback() { + use crate::aggregates::group_values::GroupValuesRows; + use arrow::array::{FixedSizeListArray, Float64Array}; + + let item = Arc::new(Field::new("item", DataType::Float64, true)); + let schema: SchemaRef = Arc::new(Schema::new(vec![Field::new( + "emb", + DataType::FixedSizeList(Arc::clone(&item), 2), + true, + )])); + assert!(supported_schema(schema.as_ref())); + + // Rows exercising +0.0 vs -0.0, two NaN bit patterns, and inner nulls. + let nan = f64::NAN; + let other_nan = f64::from_bits(0x7ff8_0000_0000_0001); + let values = Float64Array::from(vec![ + Some(0.0), + Some(1.0), // [ +0.0, 1.0 ] + Some(-0.0), + Some(1.0), // [ -0.0, 1.0 ] + Some(nan), + Some(2.0), // [ NaN, 2.0 ] + Some(other_nan), + Some(2.0), // [ NaN', 2.0 ] + Some(0.0), + Some(1.0), // [ +0.0, 1.0 ] (dup of row 0) + ]); + let field_ref = Arc::new(Field::new("item", DataType::Float64, true)); + let input: ArrayRef = Arc::new(FixedSizeListArray::new( + field_ref, + 2, + Arc::new(values), + None, + )); + + let cols = vec![input]; + + let mut column_path = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + let mut rows_path = GroupValuesRows::try_new(Arc::clone(&schema)).unwrap(); + + let mut g1 = vec![]; + let mut g2 = vec![]; + column_path.intern(&cols, &mut g1).unwrap(); + rows_path.intern(&cols, &mut g2).unwrap(); + + assert_eq!( + canonical_grouping(&g1), + canonical_grouping(&g2), + "column-wise path must induce the same grouping as the rows fallback \ + on float edge cases (got column={g1:?}, rows={g2:?})" + ); + assert_eq!(column_path.len(), rows_path.len()); + } + + /// Equivalence across multiple `intern` batches and `EmitTo::First(n)`. + #[test] + fn multi_batch_and_emit_first_matches_rows_fallback() { + use crate::aggregates::group_values::GroupValuesRows; + use arrow::array::{FixedSizeListArray, Int32Array}; + use arrow::datatypes::Int32Type; + + let item = Arc::new(Field::new("item", DataType::Int32, true)); + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("emb", DataType::FixedSizeList(Arc::clone(&item), 2), true), + ])); + + let make_batch = |base: i32| -> Vec { + let k = Arc::new(Int32Array::from(vec![base, base + 1, base])) as ArrayRef; + let emb: Vec>>> = vec![ + Some(vec![Some(base), Some(base)]), + Some(vec![Some(base + 1), None]), + Some(vec![Some(base), Some(base)]), // dup of row 0 + ]; + let emb = Arc::new( + FixedSizeListArray::from_iter_primitive::(emb, 2), + ) as ArrayRef; + vec![k, emb] + }; + + let mut column_path = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + let mut rows_path = GroupValuesRows::try_new(Arc::clone(&schema)).unwrap(); + + for base in [0, 10, 0] { + let cols = make_batch(base); + let (mut a, mut b) = (vec![], vec![]); + column_path.intern(&cols, &mut a).unwrap(); + rows_path.intern(&cols, &mut b).unwrap(); + // Same grouping (partition), even if the opaque group-index labels + // differ between the vectorized and sequential paths. + assert_eq!( + canonical_grouping(&a), + canonical_grouping(&b), + "grouping must match for batch base={base}" + ); + } + + let total_groups = column_path.len(); + assert_eq!(total_groups, rows_path.len()); + + // `EmitTo::First(n)` then `EmitTo::All` on the nested column path must + // work and together emit exactly `total_groups` rows. (Cross-path value + // equality is covered by `mixed_schema_...` and the row_backed unit + // tests; group-index ordering differs here so we check counts.) + let col_first = column_path.emit(EmitTo::First(2)).unwrap(); + assert_eq!(col_first[0].len(), 2); + let col_rest = column_path.emit(EmitTo::All).unwrap(); + assert_eq!(col_first[0].len() + col_rest[0].len(), total_groups); + // Column count / schema preserved on both emits. + assert_eq!(col_first.len(), schema.fields().len()); + assert_eq!(col_rest.len(), schema.fields().len()); + } + + /// CRITICAL invariant: if `group_column_supported_type(t)` returns true + /// the dispatcher must accept that type at intern time, and conversely + /// if `group_column_supported_type(t)` returns false the planner must + /// NOT route it through `GroupValuesColumn`. A divergence here would + /// let the planner select `GroupValuesColumn` for a type whose + /// dispatcher arm is missing, producing a runtime `not_impl_err` after + /// the field reaches the builder factory. + /// + /// 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 + /// to the supported_cases vector; when a type is intentionally rejected + /// it should be added to unsupported_cases. #[test] fn test_split_vec_min_alloc_drain_branch() { // n * 2 <= len → drain+collect branch (allocates n elements) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs new file mode 100644 index 0000000000000..29beb3bd66229 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs @@ -0,0 +1,1014 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A generic [`GroupColumn`] backed by the arrow row format. +//! +//! Unlike the type-specialized builders in this module (primitive, byte, +//! boolean, ...), [`RowsGroupColumn`] works for *any* data type that arrow's +//! [`RowConverter`] can encode — including nested types such as `Struct`, +//! `List`, `LargeList` and `FixedSizeList`. It stores one group value per row +//! in a single-column [`Rows`] buffer and compares group keys by their encoded +//! bytes. +//! +//! # Why this exists +//! +//! [`GroupValuesColumn`] can only be used when *every* column of the group-by +//! key has a [`GroupColumn`] implementation; otherwise the whole aggregation +//! falls back to the row-wise [`GroupValuesRows`], which is materially slower +//! and heavier for the columns that *would* have qualified for the column-wise +//! fast path. By providing a generic fallback `GroupColumn`, a schema like +//! `GROUP BY int_col, struct_col` keeps `int_col` on its fast native builder +//! and only pays the row-encoding cost on `struct_col`, instead of dragging both +//! columns onto `GroupValuesRows`. +//! +//! # Relationship to hashing +//! +//! This column does not hash anything itself: [`GroupValuesColumn`] hashes the +//! raw input columns via `create_hashes`, which already supports nested types. +//! Equality is decided here by comparing arrow-row bytes. For the two to agree +//! on group identity, values that this column considers equal must hash equal — +//! see the float `-0.0` / `NaN` note on [`RowsGroupColumn`]. +//! +//! [`GroupValuesColumn`]: crate::aggregates::group_values::multi_group_by::GroupValuesColumn +//! [`GroupValuesRows`]: crate::aggregates::group_values::GroupValuesRows + +use crate::aggregates::group_values::multi_group_by::GroupColumn; +use crate::aggregates::group_values::row::encode_array_if_necessary; + +use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; +use arrow::datatypes::DataType; +use arrow::row::{RowConverter, Rows, SortField}; +use datafusion_common::{DataFusionError, Result}; + +/// A [`GroupColumn`] that stores group values for a single column in the arrow +/// [row format], backed by a single-field [`RowConverter`]. +/// +/// # NULL semantics +/// +/// The [`GroupColumn`] contract treats two NULLs as equal. The row format +/// encodes NULL with a distinct sentinel, so `null`-row bytes compare equal to +/// each other and unequal to any non-null row — matching the contract without +/// special-casing. +/// +/// # Float `-0.0` / `NaN` +/// +/// Equality here is byte equality under arrow's IEEE-754 *totalOrder* row +/// encoding, which treats `-0.0` and `+0.0` as distinct and canonicalizes +/// `NaN`. Because hashing is performed separately (on the raw input array), a +/// caller must ensure the two agree — e.g. by normalizing `-0.0 → +0.0` on the +/// input columns before hashing when a float leaf is present (as +/// [`GroupValuesRows`] does). See the module docs. +/// +/// [row format]: arrow::row +/// [`GroupValuesRows`]: crate::aggregates::group_values::GroupValuesRows +pub struct RowsGroupColumn { + /// Single-field row converter for this column's data type. + row_converter: RowConverter, + /// Accumulated group values in row format; `group_values.row(i)` is the + /// group value for group index `i`. + group_values: Rows, + /// The column's expected output type. The row format decodes dictionary / + /// run-end encoded values to their plain value type, so emitted arrays are + /// re-encoded to this type in `build` / `take_n` (mirroring + /// `GroupValuesRows::emit`). + output_type: DataType, +} + +/// Walk `data_type`'s subtree and return `true` if it contains a +/// [`DataType::FixedSizeList`] whose descendant tree includes any +/// [`DataType::Dictionary`]. +/// +/// Two-state recursion: once we cross a `FixedSizeList`, `inside_fsl` +/// stays true for every descendant, so a `Dictionary` anywhere below +/// counts. Above that boundary, encountering a `Dictionary` is fine — +/// only nested containers propagate the risk. +/// +/// TODO: this guard works around +/// (`decode_fixed_size_list` panics instead of applying the +/// dictionary-flatten `corrected_type` step). Fixed upstream by +/// (merged 2026-07-24, not +/// yet in a release as of arrow 59.1.0). Once DataFusion upgrades to an +/// arrow release containing that fix, `FixedSizeList` will +/// decode like the other list-likes (flattened child, re-encoded by +/// `encode_array_if_necessary`'s existing `FixedSizeList` arm) — remove +/// this guard and its `supports_type` rejection at that point. +fn contains_fsl_with_dictionary(data_type: &DataType) -> bool { + fn walk(dt: &DataType, inside_fsl: bool) -> bool { + match dt { + DataType::Dictionary(_, _) => inside_fsl, + DataType::FixedSizeList(f, _) => walk(f.data_type(), true), + DataType::List(f) + | DataType::LargeList(f) + | DataType::ListView(f) + | DataType::LargeListView(f) => walk(f.data_type(), inside_fsl), + DataType::Map(f, _) => walk(f.data_type(), inside_fsl), + DataType::Struct(fs) => fs.iter().any(|f| walk(f.data_type(), inside_fsl)), + DataType::RunEndEncoded(_, values) => walk(values.data_type(), inside_fsl), + DataType::Union(fs, _) => { + fs.iter().any(|(_, f)| walk(f.data_type(), inside_fsl)) + } + _ => false, + } + } + walk(data_type, false) +} + +/// Return `true` if `data_type` contains a [`DataType::Union`] or +/// [`DataType::RunEndEncoded`] anywhere in its subtree. +/// +/// These two nested variants can round-trip through `RowConverter` in +/// principle, but their arrow-row decoders have not been validated by +/// this crate's test matrix against the full range of leaf types (dict, +/// nested, etc.). Before this PR both were handled by `GroupValuesRows` +/// (they were not `is_nested`-eligible for `GroupValuesColumn`), so +/// reject them here to preserve the pre-PR routing rather than route +/// untested shapes through `RowsGroupColumn`. When we grow explicit +/// round-trip tests for these types, this blacklist can be removed. +fn contains_union_or_run_end_encoded(data_type: &DataType) -> bool { + match data_type { + DataType::Union(_, _) | DataType::RunEndEncoded(_, _) => true, + DataType::List(f) + | DataType::LargeList(f) + | DataType::ListView(f) + | DataType::LargeListView(f) + | DataType::FixedSizeList(f, _) => { + contains_union_or_run_end_encoded(f.data_type()) + } + DataType::Map(f, _) => contains_union_or_run_end_encoded(f.data_type()), + DataType::Struct(fs) => fs + .iter() + .any(|f| contains_union_or_run_end_encoded(f.data_type())), + _ => false, + } +} + +impl RowsGroupColumn { + /// Returns whether `data_type` can be handled by this generic column. + /// + /// This is stricter than [`RowConverter::supports_fields`]: the row + /// format also has to survive the `build` / `take_n` reverse trip + /// through [`RowConverter::convert_rows`], and arrow's + /// `decode_fixed_size_list` (arrow-row 59.1.0) skips the + /// dictionary-flatten correction that the other list-like decoders + /// apply, so any `FixedSizeList` containing a `Dictionary` leaf + /// panics on emit with `"FixedSizeListArray expected data type + /// Dictionary(...) got for \"item\""`. + /// + /// Reject those shapes here so `make_group_column` falls back to + /// `GroupValuesRows`. The other list-likes (`List`, `LargeList`, + /// `ListView`, `LargeListView`, `Map`) do carry the correction, so + /// they decode without panicking — but the correction *flattens* any + /// dictionary child to its value type, so `build` / `take_n` must + /// re-encode the emitted array back to `output_type` via + /// `encode_array_if_necessary` (which has a reconstruction arm for + /// each of these containers). + /// + /// Additionally, `Union` and `RunEndEncoded` are rejected because + /// they were routed to `GroupValuesRows` before this column existed + /// and their arrow-row round-trip has not been covered by this + /// crate's tests yet. Keeping them on the pre-PR path avoids + /// introducing an untested code path for those types. + pub fn supports_type(data_type: &DataType) -> bool { + if contains_fsl_with_dictionary(data_type) { + return false; + } + if contains_union_or_run_end_encoded(data_type) { + return false; + } + RowConverter::supports_fields(&[SortField::new(data_type.clone())]) + } + + /// Create an empty [`RowsGroupColumn`] for `data_type`. + pub fn try_new(data_type: DataType) -> Result { + let row_converter = RowConverter::new(vec![SortField::new(data_type.clone())])?; + let group_values = row_converter.empty_rows(0, 0); + Ok(Self { + row_converter, + group_values, + output_type: data_type, + }) + } + + /// Materialize `rows` into a single array of `self.output_type`, re-applying + /// dictionary / run-end encoding the row format strips on decode. + fn rows_to_array<'a>( + &self, + rows: impl IntoIterator>, + ) -> ArrayRef { + let mut arrays = self + .row_converter + .convert_rows(rows) + .expect("row conversion during emit"); + debug_assert_eq!(arrays.len(), 1, "single-field row converter"); + let array = arrays.swap_remove(0); + encode_array_if_necessary(&array, &self.output_type) + .expect("dictionary re-encode during emit") + } + + /// Encode a whole incoming column into the row format. + fn convert(&self, array: &ArrayRef) -> Result { + self.row_converter + .convert_columns(std::slice::from_ref(array)) + .map_err(DataFusionError::from) + } +} + +impl GroupColumn for RowsGroupColumn { + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + // Scalar path (hash-collision remainder / streaming). Encode just the + // single incoming row rather than the whole column. The vectorized + // methods below encode the batch once; this path is expected to be rare. + let incoming = self + .convert(&array.slice(rhs_row, 1)) + .expect("row conversion during equal_to"); + self.group_values.row(lhs_row) == incoming.row(0) + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + let incoming = self.convert(&array.slice(row, 1))?; + self.group_values.push(incoming.row(0)); + Ok(()) + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + // Encode the incoming column once for the whole batch. + let incoming = self + .convert(array) + .expect("row conversion during vectorized_equal_to"); + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + // Preserve the AND-accumulate contract: skip rows already false. + if !equal_to_results.get_bit(idx) { + continue; + } + if self.group_values.row(lhs_row) != incoming.row(rhs_row) { + equal_to_results.set_bit(idx, false); + } + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + // Encode the incoming column once, then push the selected rows. + let incoming = self.convert(array)?; + for &row in rows { + self.group_values.push(incoming.row(row)); + } + Ok(()) + } + + fn len(&self) -> usize { + self.group_values.num_rows() + } + + fn size(&self) -> usize { + self.row_converter.size() + self.group_values.size() + } + + fn build(self: Box) -> ArrayRef { + self.rows_to_array(&self.group_values) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + debug_assert!(n <= self.group_values.num_rows()); + + // Materialize the first `n` group rows. + let output = self.rows_to_array(self.group_values.iter().take(n)); + + // Shift the remaining rows to the front by rebuilding the buffer. + // TODO: mirror the arrow-rs efficiency TODO in `GroupValuesRows::emit`. + let mut remaining = self.row_converter.empty_rows(0, 0); + for row in self.group_values.iter().skip(n) { + remaining.push(row); + } + self.group_values = remaining; + + output + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::array::{Array, ArrayRef, FixedSizeListArray, Int32Array, StructArray}; + use arrow::datatypes::{DataType, Field, Int32Type}; + use std::sync::Arc; + + fn fsl_i32(data: Vec>>>, list_len: i32) -> ArrayRef { + Arc::new(FixedSizeListArray::from_iter_primitive::( + data, list_len, + )) + } + + /// The generic column must agree with a per-row reference for equality, + /// including inner-null and outer-null rows, on a `FixedSizeList`. + #[test] + fn fsl_append_equal_to_build_roundtrip() { + let dt = DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, true)), + 2, + ); + let mut col = Box::new(RowsGroupColumn::try_new(dt).unwrap()); + + // group values: [1,2], null-outer, [3, null-inner] + let input = fsl_i32( + vec![ + Some(vec![Some(1), Some(2)]), + None, + Some(vec![Some(3), None]), + ], + 2, + ); + + col.vectorized_append(&input, &[0, 1, 2]).unwrap(); + assert_eq!(col.len(), 3); + + // Probe with a fresh batch: row0 == group0, row1 (null) == group1, + // row2 differs from group0, row3 (inner null) == group2. + let probe = fsl_i32( + vec![ + Some(vec![Some(1), Some(2)]), // == g0 + None, // == g1 + Some(vec![Some(9), Some(9)]), // != g0 + Some(vec![Some(3), None]), // == g2 + ], + 2, + ); + + assert!(col.equal_to(0, &probe, 0)); + assert!(col.equal_to(1, &probe, 1)); + assert!(!col.equal_to(0, &probe, 2)); + assert!(col.equal_to(2, &probe, 3)); + + // Vectorized equal_to should match the scalar reference. + let mut results = BooleanBufferBuilder::new(3); + results.append_n(3, true); + col.vectorized_equal_to(&[0, 1, 2], &probe, &[0, 1, 3], &mut results); + assert!(results.get_bit(0)); + assert!(results.get_bit(1)); + assert!(results.get_bit(2)); + + // build() must reproduce the original group values. + let out = col.build(); + let out = out.as_any().downcast_ref::().unwrap(); + assert_eq!(out.len(), 3); + assert!(out.is_null(1)); + assert!(!out.is_null(0)); + } + + /// `take_n` must emit the first `n` rows and shift the rest to the front. + #[test] + fn fsl_take_n_shifts_remaining() { + let dt = DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, true)), + 1, + ); + let mut col = RowsGroupColumn::try_new(dt).unwrap(); + + let input = fsl_i32( + vec![ + Some(vec![Some(10)]), + Some(vec![Some(20)]), + Some(vec![Some(30)]), + ], + 1, + ); + col.vectorized_append(&input, &[0, 1, 2]).unwrap(); + + let first = col.take_n(1); + let first = first.as_any().downcast_ref::().unwrap(); + let first_vals = first + .value(0) + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + assert_eq!(first_vals.value(0), 10); + assert_eq!(col.len(), 2); + + // Remaining 20, 30 should now be at indices 0, 1. + let rest = Box::new(col).build(); + let rest = rest.as_any().downcast_ref::().unwrap(); + assert_eq!(rest.len(), 2); + let g0 = rest + .value(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(g0, 20); + } + + /// Works for `Struct` too — proves the column is type-generic. + #[test] + fn struct_roundtrip() { + let dt = DataType::Struct(vec![Field::new("a", DataType::Int32, true)].into()); + let mut col = RowsGroupColumn::try_new(dt).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2)])); + let input: ArrayRef = Arc::new(StructArray::new( + vec![Field::new("a", DataType::Int32, true)].into(), + vec![a], + None, + )); + col.vectorized_append(&input, &[0, 1]).unwrap(); + assert_eq!(col.len(), 2); + assert!(col.equal_to(0, &input, 0)); + assert!(!col.equal_to(0, &input, 1)); + } + + #[test] + fn supports_type_matches_row_converter_impl() { + assert!(RowsGroupColumn::supports_type(&DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, true)), + 3 + ))); + assert!(RowsGroupColumn::supports_type(&DataType::Struct( + vec![Field::new("a", DataType::Int32, true)].into() + ))); + // Whether Map is encodable depends on the arrow-rs version. + // Just assert that our `supports_type` agrees with arrow's + // `RowConverter::supports_fields` — either both accept it or both + // reject it. Both are correct wrt the invariant. + let map_field = Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("keys", DataType::Int32, false), + Field::new("values", DataType::Int32, true), + ] + .into(), + ), + false, + )); + let map_dt = DataType::Map(map_field, false); + let arrow_supports = + RowConverter::supports_fields(&[SortField::new(map_dt.clone())]); + assert_eq!(RowsGroupColumn::supports_type(&map_dt), arrow_supports); + } + + /// Regression test for the nested-container recursion in + /// [`crate::aggregates::group_values::row::encode_array_if_necessary`]. + /// `RowConverter` flattens dictionary values on the way in, so a + /// `List>` schema round-trips with `Utf8` values + /// unless the helper re-encodes the leaf. Without that recursion, + /// `build()` would emit an array whose data type does not match the + /// group column's declared type. + #[test] + fn build_preserves_list_of_dictionary_schema() { + use arrow::array::{DictionaryArray, ListArray, StringArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::Int32Type; + + let dict_dt = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + let item_field = Arc::new(Field::new("item", dict_dt.clone(), true)); + let outer_dt = DataType::List(Arc::clone(&item_field)); + + // Skip if this arrow-rs version rejects the nesting — the invariant we + // care about is `output().data_type() == declared type` conditional on + // supports_type saying yes. + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + + // Build List> of one row = ["a", "b"]. + let values = Arc::new(StringArray::from(vec!["a", "b"])); + let keys = Int32Array::from(vec![0, 1]); + let dict = DictionaryArray::::try_new(keys, values).unwrap(); + let offsets = OffsetBuffer::from_lengths([2]); + let list = + ListArray::try_new(Arc::clone(&item_field), offsets, Arc::new(dict), None) + .unwrap(); + let input: ArrayRef = Arc::new(list); + + col.vectorized_append(&input, &[0]).unwrap(); + let built = col.build(); + assert_eq!( + built.data_type(), + &outer_dt, + "build() must return the declared List data type, \ + not the RowConverter-flattened List", + ); + } + + // ---- FSL rejection ---------------------------------------- + // + // arrow-row 59.1.0's `decode_fixed_size_list` skips the + // dict-flatten correction that the generic `decode` path applies + // to `List` / `LargeList` / `ListView` / `LargeListView` / `Map`, + // so any `FixedSizeList` containing a `Dictionary` leaf panics on + // emit. `supports_type` must reject those shapes so + // `GroupValuesRows` fallback handles them instead. These tests pin + // the current shape of that black-list. + + fn dict_utf8() -> DataType { + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)) + } + + fn fsl_of(inner: DataType) -> DataType { + DataType::FixedSizeList(Arc::new(Field::new("item", inner, true)), 2) + } + + #[test] + fn supports_type_rejects_fixed_size_list_of_dict() { + // Direct case: `FixedSizeList>`. + assert!(!RowsGroupColumn::supports_type(&fsl_of(dict_utf8()))); + } + + #[test] + fn supports_type_rejects_fsl_with_dict_nested_in_struct() { + // The dict is one level deep under a struct that is itself the + // FSL element. arrow-row still panics because `convert_raw` + // returns the struct with a decoded (Utf8) field while the + // FSL builder expects the declared struct-with-dict shape. + let struct_dt = DataType::Struct(vec![Field::new("d", dict_utf8(), true)].into()); + assert!(!RowsGroupColumn::supports_type(&fsl_of(struct_dt))); + } + + #[test] + fn supports_type_rejects_fsl_with_dict_nested_in_list() { + // `FixedSizeList>` — the inner `List` handles + // dicts correctly on its own, but the outer FSL wrapper still + // panics with the mismatched declared child type. + let list_of_dict = + DataType::List(Arc::new(Field::new("item", dict_utf8(), true))); + assert!(!RowsGroupColumn::supports_type(&fsl_of(list_of_dict))); + } + + #[test] + fn supports_type_rejects_fsl_hidden_under_outer_list() { + // Sibling positioning: the outer container is a `List` (which is + // fine on its own), but its child is a `FixedSizeList`. + // The panic surface is at the inner FSL layer regardless of what + // wraps it, so this must still be rejected. + let outer = + DataType::List(Arc::new(Field::new("item", fsl_of(dict_utf8()), true))); + assert!(!RowsGroupColumn::supports_type(&outer)); + } + + #[test] + fn supports_type_rejects_fsl_hidden_under_outer_struct() { + // Same, but the outer wrapper is a struct. + let outer = + DataType::Struct(vec![Field::new("f", fsl_of(dict_utf8()), true)].into()); + assert!(!RowsGroupColumn::supports_type(&outer)); + } + + // ---- FSL without dicts is still fine ---------------------------- + + #[test] + fn supports_type_accepts_fsl_of_primitive() { + // Sanity: a plain FSL must not get caught by the + // dict-under-FSL blacklist. + assert!(RowsGroupColumn::supports_type(&fsl_of(DataType::Int32))); + } + + #[test] + fn supports_type_accepts_fsl_of_struct_without_dict() { + // FSL of struct where the struct's fields are all primitives. + let struct_dt = + DataType::Struct(vec![Field::new("a", DataType::Int32, true)].into()); + assert!(RowsGroupColumn::supports_type(&fsl_of(struct_dt))); + } + + // ---- Positive round-trip tests for non-FSL list-likes ----------- + // + // The other list-like decoders in arrow-row 59.1.0 + // (`GenericListArrayOrMap` path) apply the corrected_type fix, so + // `List`, `LargeList`, `ListView`, `LargeListView` + // and `Map<..., Dict>` all round-trip cleanly. These tests pin + // that they are (a) accepted by `supports_type` and (b) actually + // survive `vectorized_append` + `build()` without panicking, so a + // future arrow-rs regression there is caught here rather than in + // production. + + #[test] + fn supports_type_accepts_large_list_of_dict() { + let dt = DataType::LargeList(Arc::new(Field::new("item", dict_utf8(), true))); + assert!(RowsGroupColumn::supports_type(&dt)); + } + + #[test] + fn supports_type_accepts_list_view_of_dict() { + let dt = DataType::ListView(Arc::new(Field::new("item", dict_utf8(), true))); + assert!(RowsGroupColumn::supports_type(&dt)); + } + + #[test] + fn supports_type_accepts_large_list_view_of_dict() { + let dt = DataType::LargeListView(Arc::new(Field::new("item", dict_utf8(), true))); + assert!(RowsGroupColumn::supports_type(&dt)); + } + + #[test] + fn supports_type_map_agrees_with_row_converter() { + // Map>. Whether arrow-row supports Map + // depends on the version; either way, our `supports_type` must + // agree with `RowConverter::supports_fields` — otherwise we'd + // pick a strategy the converter can't back. + let entries = Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("keys", DataType::Int32, false), + Field::new("values", dict_utf8(), true), + ] + .into(), + ), + false, + )); + let map_dt = DataType::Map(entries, false); + let arrow_supports = + RowConverter::supports_fields(&[SortField::new(map_dt.clone())]); + assert_eq!(RowsGroupColumn::supports_type(&map_dt), arrow_supports); + } + + /// End-to-end regression: `LargeList>` must + /// actually survive `vectorized_append` + `build()` on the current + /// arrow-rs version, not just be accepted by `supports_type`. + #[test] + fn build_preserves_large_list_of_dictionary_schema() { + use arrow::array::{DictionaryArray, LargeListArray, StringArray}; + use arrow::buffer::OffsetBuffer; + + let item_field = Arc::new(Field::new("item", dict_utf8(), true)); + let outer_dt = DataType::LargeList(Arc::clone(&item_field)); + + // Skip if this arrow-rs version rejects the nesting (defensive: + // the invariant we care about is `output().data_type() == declared` + // conditional on `supports_type` saying yes). + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + + let values = Arc::new(StringArray::from(vec!["a", "b"])); + let keys = Int32Array::from(vec![0, 1]); + let dict = DictionaryArray::::try_new(keys, values).unwrap(); + let offsets = OffsetBuffer::::from_lengths([2]); + let list = LargeListArray::try_new( + Arc::clone(&item_field), + offsets, + Arc::new(dict), + None, + ) + .unwrap(); + + col.vectorized_append(&(Arc::new(list) as ArrayRef), &[0]) + .unwrap(); + let built = col.build(); + assert_eq!( + built.data_type(), + &outer_dt, + "LargeList: build() must preserve the declared type", + ); + } + + /// Build a two-row `ListView>` array with rows + /// `["a", "b"]` and `["c"]` — the shape from the review reproducer: + /// `arrow_cast(a, 'ListView(Dictionary(Int32, Utf8))')`. + fn list_view_of_dict_input() -> (DataType, ArrayRef) { + use arrow::array::{DictionaryArray, ListViewArray, StringArray}; + use arrow::buffer::ScalarBuffer; + + let item_field = Arc::new(Field::new("item", dict_utf8(), true)); + let outer_dt = DataType::ListView(Arc::clone(&item_field)); + + let values = Arc::new(StringArray::from(vec!["a", "b", "c"])); + let keys = Int32Array::from(vec![0, 1, 2]); + let dict = DictionaryArray::::try_new(keys, values).unwrap(); + let offsets = ScalarBuffer::::from(vec![0, 2]); + let sizes = ScalarBuffer::::from(vec![2, 1]); + let list = ListViewArray::try_new( + Arc::clone(&item_field), + offsets, + sizes, + Arc::new(dict), + None, + ) + .unwrap(); + (outer_dt, Arc::new(list) as ArrayRef) + } + + /// `ListView`: arrow-row's `decode_list_view` flattens the + /// dictionary child (`corrected_type`), so `build` must re-encode + /// the emitted array back to the declared type. Regression for the + /// review reproducer that failed with + /// `expected ListView(Dictionary(Int32, Utf8)) but found ListView(Utf8)`. + #[test] + fn build_preserves_list_view_of_dictionary_schema() { + let (outer_dt, input) = list_view_of_dict_input(); + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + col.vectorized_append(&input, &[0, 1]).unwrap(); + assert_eq!(col.len(), 2); + + let built = col.build(); + assert_eq!( + built.data_type(), + &outer_dt, + "ListView: build() must return the declared type, \ + not the RowConverter-flattened ListView", + ); + assert_eq!(built.len(), 2); + } + + /// Same regression through the `take_n` path (used by + /// `EmitTo::First(n)`), including the type of the *remaining* + /// values emitted by a subsequent `build`. + #[test] + fn take_n_preserves_list_view_of_dictionary_schema() { + let (outer_dt, input) = list_view_of_dict_input(); + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + col.vectorized_append(&input, &[0, 1]).unwrap(); + + let taken = col.take_n(1); + assert_eq!( + taken.data_type(), + &outer_dt, + "ListView: take_n() must return the declared type", + ); + assert_eq!(taken.len(), 1); + + let rest = col.build(); + assert_eq!( + rest.data_type(), + &outer_dt, + "ListView: build() after take_n must also preserve the type", + ); + assert_eq!(rest.len(), 1); + } + + /// `LargeListView` fails the same way as `ListView` + /// per the review; cover both `build` and `take_n`. + #[test] + fn build_and_take_n_preserve_large_list_view_of_dictionary_schema() { + use arrow::array::{DictionaryArray, LargeListViewArray, StringArray}; + use arrow::buffer::ScalarBuffer; + + let item_field = Arc::new(Field::new("item", dict_utf8(), true)); + let outer_dt = DataType::LargeListView(Arc::clone(&item_field)); + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let values = Arc::new(StringArray::from(vec!["a", "b", "c"])); + let keys = Int32Array::from(vec![0, 1, 2]); + let dict = DictionaryArray::::try_new(keys, values).unwrap(); + let offsets = ScalarBuffer::::from(vec![0, 2]); + let sizes = ScalarBuffer::::from(vec![2, 1]); + let list = LargeListViewArray::try_new( + Arc::clone(&item_field), + offsets, + sizes, + Arc::new(dict), + None, + ) + .unwrap(); + let input: ArrayRef = Arc::new(list); + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + col.vectorized_append(&input, &[0, 1]).unwrap(); + + let taken = col.take_n(1); + assert_eq!( + taken.data_type(), + &outer_dt, + "LargeListView: take_n() must return the declared type", + ); + + let rest = col.build(); + assert_eq!( + rest.data_type(), + &outer_dt, + "LargeListView: build() must return the declared type", + ); + assert_eq!(rest.len(), 1); + } + + /// Group-identity must survive the dictionary flatten + re-encode + /// round trip: appending the same logical list twice (with distinct + /// dictionary key mappings) must map to one group, a different list + /// to another. Mirrors the review reproducer's GROUP BY semantics + /// (2 distinct groups from 3 input rows). + #[test] + fn list_view_of_dict_groups_by_logical_value() { + use arrow::array::{DictionaryArray, ListViewArray, StringArray}; + use arrow::buffer::ScalarBuffer; + + let item_field = Arc::new(Field::new("item", dict_utf8(), true)); + let outer_dt = DataType::ListView(Arc::clone(&item_field)); + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + // Rows: ["a","b"], ["a","b"], ["c"] → 2 distinct groups. + let values = Arc::new(StringArray::from(vec!["a", "b", "a", "b", "c"])); + let keys = Int32Array::from(vec![0, 1, 2, 3, 4]); + let dict = DictionaryArray::::try_new(keys, values).unwrap(); + let offsets = ScalarBuffer::::from(vec![0, 2, 4]); + let sizes = ScalarBuffer::::from(vec![2, 2, 1]); + let list = ListViewArray::try_new( + Arc::clone(&item_field), + offsets, + sizes, + Arc::new(dict), + None, + ) + .unwrap(); + let input: ArrayRef = Arc::new(list); + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + // Append row 0 as group 0. + col.vectorized_append(&input, &[0]).unwrap(); + // Row 1 must compare equal to group 0 (same logical value). + assert!( + col.equal_to(0, &input, 1), + "identical logical lists must be equal regardless of dict keys", + ); + // Row 2 must not. + assert!( + !col.equal_to(0, &input, 2), + "different logical lists must not be equal", + ); + + col.vectorized_append(&input, &[2]).unwrap(); + assert_eq!(col.len(), 2, "3 input rows → 2 distinct groups"); + + let built = col.build(); + assert_eq!(built.data_type(), &outer_dt); + assert_eq!(built.len(), 2); + } + + /// End-to-end regression for `Map>` when + /// arrow-row supports it. Same intent as the LargeList test. + #[test] + fn build_preserves_map_of_dictionary_schema() { + use arrow::array::{ + DictionaryArray, Int32Array, MapArray, StringArray, StructArray, + }; + use arrow::buffer::OffsetBuffer; + + let key_field = Arc::new(Field::new("keys", DataType::Int32, false)); + let value_field = Arc::new(Field::new("values", dict_utf8(), true)); + let entries_field = Arc::new(Field::new( + "entries", + DataType::Struct(vec![(*key_field).clone(), (*value_field).clone()].into()), + false, + )); + let outer_dt = DataType::Map(Arc::clone(&entries_field), false); + + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + + // One map entry: {1 -> "a"}. + let keys = Arc::new(Int32Array::from(vec![1])) as ArrayRef; + let values_arr = Arc::new(StringArray::from(vec!["a"])); + let value_keys = Int32Array::from(vec![0]); + let value_dict = + DictionaryArray::::try_new(value_keys, values_arr).unwrap(); + let entries = StructArray::try_new( + vec![(*key_field).clone(), (*value_field).clone()].into(), + vec![keys, Arc::new(value_dict)], + None, + ) + .unwrap(); + let offsets = OffsetBuffer::::from_lengths([1]); + let map = + MapArray::try_new(Arc::clone(&entries_field), offsets, entries, None, false) + .unwrap(); + + col.vectorized_append(&(Arc::new(map) as ArrayRef), &[0]) + .unwrap(); + let built = col.build(); + assert_eq!( + built.data_type(), + &outer_dt, + "Map<..., Dict>: build() must preserve the declared type", + ); + } + + // ---- Union / RunEndEncoded defensive rejection ----------------- + // + // Before this PR both types were routed to `GroupValuesRows` + // (`group_column_supported_type` didn't have a nested branch). This + // PR added `is_nested`-based dispatch to `RowsGroupColumn`, which + // would opt them in — but the arrow-row round-trip for these two + // families hasn't been covered by our tests. Reject them here so + // the pre-PR routing is preserved; drop the blacklist when the + // round-trip matrix grows to include them. + + #[test] + fn supports_type_rejects_union() { + use arrow::datatypes::UnionFields; + + let fields = UnionFields::try_new( + vec![0_i8, 1_i8], + vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ], + ) + .unwrap(); + let dt = DataType::Union(fields, arrow::datatypes::UnionMode::Dense); + assert!( + !RowsGroupColumn::supports_type(&dt), + "Union must fall back to GroupValuesRows until arrow-row \ + round-trip is covered by our tests", + ); + } + + #[test] + fn supports_type_rejects_run_end_encoded_with_nested_values() { + // REE with `is_nested() = true` (nested values) is what this PR + // could otherwise opt into RowsGroupColumn; keep it on + // GroupValuesRows. + let list_of_i32 = + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))); + let dt = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new("values", list_of_i32, true)), + ); + assert!(!RowsGroupColumn::supports_type(&dt)); + } + + #[test] + fn supports_type_rejects_run_end_encoded_with_scalar_values() { + // REE with scalar values is `is_nested() == false`, so + // `group_column_supported_type` never routes it to us via the + // nested branch anyway — but pin the invariant explicitly so a + // future refactor doesn't accidentally opt it in. + let dt = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new("values", DataType::Utf8, true)), + ); + assert!(!RowsGroupColumn::supports_type(&dt)); + } + + #[test] + fn supports_type_rejects_ree_hidden_under_outer_wrapper() { + // REE buried under a struct or list: still rejected because + // the wrapper's decoder recurses through the REE branch we + // haven't validated. + let ree = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new("values", DataType::Utf8, true)), + ); + let outer = DataType::Struct(vec![Field::new("f", ree, true)].into()); + assert!(!RowsGroupColumn::supports_type(&outer)); + } + + #[test] + fn supports_type_accepts_plain_list_and_struct_still() { + // Sanity: the defensive Union/REE blacklist must not accidentally + // catch the well-tested list-likes / structs that this column + // exists to serve. + let list_of_int = + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))); + assert!(RowsGroupColumn::supports_type(&list_of_int)); + + let struct_of_prims = DataType::Struct( + vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ] + .into(), + ); + assert!(RowsGroupColumn::supports_type(&struct_of_prims)); + } +} diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index a3bd31f76c233..5fefe9df3e849 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -17,7 +17,8 @@ use crate::aggregates::group_values::GroupValues; use arrow::array::{ - Array, ArrayRef, ListArray, PrimitiveArray, RunArray, StructArray, + Array, ArrayRef, FixedSizeListArray, LargeListArray, LargeListViewArray, ListArray, + ListViewArray, MapArray, PrimitiveArray, RunArray, StructArray, downcast_run_end_index, }; use arrow::compute::cast; @@ -239,7 +240,7 @@ impl GroupValues for GroupValuesRows { // https://github.com/apache/datafusion/issues/7647 for (field, array) in self.schema.fields.iter().zip(&mut output) { let expected = field.data_type(); - *array = dictionary_encode_if_necessary(array, expected)?; + *array = encode_array_if_necessary(array, expected)?; } self.group_values = Some(group_values); @@ -259,7 +260,17 @@ impl GroupValues for GroupValuesRows { } } -fn dictionary_encode_if_necessary( +/// Re-apply dictionary / run-end encoding to `array` so it matches `expected`. +/// +/// Arrow's [`RowConverter`] flattens dictionary and run-end-encoded values to +/// their plain value type during row encoding (at [`RowConverter::append`]), +/// so any group-value array produced from the row format is in that plain +/// type and must be re-encoded to match the schema's expected type before +/// being returned. Shared with the generic row-backed `GroupColumn`. +/// +/// [`RowConverter`]: arrow::row::RowConverter +/// [`RowConverter::append`]: arrow::row::RowConverter::append +pub(crate) fn encode_array_if_necessary( array: &ArrayRef, expected: &DataType, ) -> Result { @@ -270,7 +281,7 @@ fn dictionary_encode_if_necessary( .iter() .zip(struct_array.columns()) .map(|(expected_field, column)| { - dictionary_encode_if_necessary(column, expected_field.data_type()) + encode_array_if_necessary(column, expected_field.data_type()) }) .collect::>>()?; @@ -286,13 +297,82 @@ fn dictionary_encode_if_necessary( Ok(Arc::new(ListArray::try_new( Arc::::clone(expected_field), list.offsets().clone(), - dictionary_encode_if_necessary( - list.values(), - expected_field.data_type(), - )?, + encode_array_if_necessary(list.values(), expected_field.data_type())?, + list.nulls().cloned(), + )?)) + } + (DataType::LargeList(expected_field), &DataType::LargeList(_)) => { + let list = array.as_any().downcast_ref::().unwrap(); + + Ok(Arc::new(LargeListArray::try_new( + Arc::::clone(expected_field), + list.offsets().clone(), + encode_array_if_necessary(list.values(), expected_field.data_type())?, + list.nulls().cloned(), + )?)) + } + (DataType::ListView(expected_field), &DataType::ListView(_)) => { + // arrow-row's `decode_list_view` applies the dictionary-flatten + // `corrected_type` to the child, so a `ListView>` + // decodes as `ListView` and the child must be + // re-encoded here (same as `List` above, plus the `sizes` + // buffer that view-lists carry). + let list = array.as_any().downcast_ref::().unwrap(); + + Ok(Arc::new(ListViewArray::try_new( + Arc::::clone(expected_field), + list.offsets().clone(), + list.sizes().clone(), + encode_array_if_necessary(list.values(), expected_field.data_type())?, + list.nulls().cloned(), + )?)) + } + (DataType::LargeListView(expected_field), &DataType::LargeListView(_)) => { + let list = array.as_any().downcast_ref::().unwrap(); + + Ok(Arc::new(LargeListViewArray::try_new( + Arc::::clone(expected_field), + list.offsets().clone(), + list.sizes().clone(), + encode_array_if_necessary(list.values(), expected_field.data_type())?, + list.nulls().cloned(), + )?)) + } + ( + DataType::FixedSizeList(expected_field, expected_size), + &DataType::FixedSizeList(_, _), + ) => { + let list = array.as_any().downcast_ref::().unwrap(); + + Ok(Arc::new(FixedSizeListArray::try_new( + Arc::::clone(expected_field), + *expected_size, + encode_array_if_necessary(list.values(), expected_field.data_type())?, list.nulls().cloned(), )?)) } + (DataType::Map(expected_entries_field, ordered), &DataType::Map(_, _)) => { + let map = array.as_any().downcast_ref::().unwrap(); + // Re-encode the entries `StructArray` (which holds key/value + // columns) against the expected entries field's struct type. + let entries_as_ref: ArrayRef = Arc::new(map.entries().clone()); + let entries = encode_array_if_necessary( + &entries_as_ref, + expected_entries_field.data_type(), + )?; + let entries = entries + .as_any() + .downcast_ref::() + .expect("Map entries recurse must yield a StructArray") + .clone(); + Ok(Arc::new(MapArray::try_new( + Arc::::clone(expected_entries_field), + map.offsets().clone(), + entries, + map.nulls().cloned(), + *ordered, + )?)) + } (DataType::Dictionary(_, _), _) => Ok(cast(array.as_ref(), expected)?), ( DataType::RunEndEncoded(run_ends_field, expected_values_field), @@ -304,7 +384,7 @@ fn dictionary_encode_if_necessary( .as_any() .downcast_ref::>() .unwrap(); - let values = dictionary_encode_if_necessary( + let values = encode_array_if_necessary( &(Arc::clone(run_array.values()) as ArrayRef), expected_values_field.data_type(), )?; From 5914809d2a54749c160089d83e5f3997c526c142 Mon Sep 17 00:00:00 2001 From: RIchard Baah <137434454+Rich-T-kid@users.noreply.github.com> Date: Sun, 16 Aug 2026 02:55:39 +0000 Subject: [PATCH 2/4] Feat: add dictionaries as a supported group column type (#23187) - works towards closing #22682. - replacement for https://github.com/apache/datafusion/pull/21765 - - https://github.com/apache/datafusion/pull/21765#pullrequestreview-4546759419 This PR introduces a specialized `GroupColumn` implementation for dictionary-typed columns inside `GroupValuesColumn`, allowing dictionary columns to participate in the columnar, vectorized aggregation path instead of the row-based fallback. **The Implementation is only about 175+ lines of code**. the remaining LOC is adding extensive test at the `GroupColumn` trait level as well as testing the `GroupValuesColumn` GroupValues trait and how it inter-opts with multi-dictionary group by's. - Adds a `DictionaryGroupValueBuilder` struct implementing the `GroupColumn` trait for `Dictionary`-typed group-by columns, supporting a configurable subset of value types - Extends the type-check gate in `GroupValuesColumn::try_new` (the `matches!` block) to accept `Dictionary(_, value_type)` where `value_type` is already supported. - Adds schema-level support so emitted dictionary group key columns round-trip through the output schema correctly - [removes casting ](https://github.com/apache/datafusion/blob/9e8dd76d6deb6736c51962d9c97e04be4e3f1fc9/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs#L1200)thats done for each dictionary array in `emit` yes. a majority of this PR is test no. this is a pure perf boost for users. (cherry picked from commit c1b39bd503cefcc170748216495191b3cdf567ae) --- .../group_values/multi_group_by/dictionary.rs | 841 ++++++++++++++++++ .../group_values/multi_group_by/mod.rs | 360 +++++++- 2 files changed, 1183 insertions(+), 18 deletions(-) create mode 100644 datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs new file mode 100644 index 0000000000000..501b13d0cd183 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs @@ -0,0 +1,841 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::aggregates::group_values::multi_group_by::GroupColumn; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanBufferBuilder, DictionaryArray, Int64Array, + PrimitiveArray, +}; +use arrow::compute::take; +use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType, Field}; +use arrow::error::ArrowError; +use datafusion_common::hash_utils::{RandomState, create_hashes}; +use datafusion_common::{DataFusionError, Result, exec_err}; +use datafusion_execution::memory_pool::proxy::HashTableAllocExt; +use hashbrown::hash_table::HashTable; +use std::marker::PhantomData; +use std::mem::size_of; +use std::sync::Arc; + +use crate::aggregates::AGGREGATION_HASH_SEED; + +/// [`GroupColumn`] for dictionary-encoded columns with key type `K`. +/// +/// `inner` holds one slot per distinct value seen across all batches. +/// `group_to_inner[group_idx]` maps each group to its slot in `inner`, +/// so groups with the same value share a slot rather than duplicating data. +pub struct DictionaryGroupValuesColumn { + /// Deduplicated store of distinct values. + inner: Box, + /// Unary null array (length 1) reused for every null appended to `inner`. + null_array: ArrayRef, + /// Maps each group index to its slot in `inner`. + group_to_inner: Vec, + /// Lookup table mapping `(value_hash, inner_slot)` for each non-null distinct value. + value_dedup: HashTable<(u64, usize)>, + /// Tracked allocation size of `value_dedup` for memory accounting via `size()`. + value_dedup_size: usize, + /// Slot in `inner` for the null group; `None` until the first null is seen. + null_inner_slot: Option, + /// Hash seed — must match `create_hashes` so hashes are consistent across calls. + random_state: RandomState, + /// Reusable scratch buffer mapping `val_idx → inner_slot` across batches. + val_to_inner: Vec, + /// Reusable hash buffer for the dictionary values array. + val_hashes: Vec, + _phantom: PhantomData, +} + +impl DictionaryGroupValuesColumn { + pub fn new(inner: Box, field: &Field) -> Self { + let null_array = arrow::array::new_null_array(field.data_type(), 1); + Self { + inner, + null_array, + group_to_inner: Vec::new(), + value_dedup: HashTable::new(), + value_dedup_size: 0, + null_inner_slot: None, + random_state: AGGREGATION_HASH_SEED, + val_to_inner: Vec::default(), + val_hashes: Vec::default(), + _phantom: PhantomData, + } + } + + /// Build a `DictionaryArray` from `values` (all inner slots) and the + /// per-group slot mapping. The null inner slot, if any, is excluded from + /// the values array and its groups emit a null key — so it never consumes + /// a key index regardless of where it sits in `inner`. + fn into_dict( + values: ArrayRef, + group_to_inner: &[usize], + null_inner_slot: Option, + ) -> ArrayRef { + let Some(null_slot) = null_inner_slot else { + // Fast path: no null group — raw slot indices are valid keys. + let keys: PrimitiveArray = group_to_inner + .iter() + .map(|&slot| Some(K::Native::usize_as(slot))) + .collect(); + return Arc::new(DictionaryArray::::new(keys, values)); + }; + + // Build a compact remap: each non-null slot gets a contiguous key + // starting from 0; the null slot is skipped entirely. + let n = values.len(); + let mut remap = vec![0usize; n]; + let mut next = 0usize; + for (i, mapped) in remap.iter_mut().enumerate() { + if i != null_slot { + *mapped = next; + next += 1; + } + } + + let keys: PrimitiveArray = group_to_inner + .iter() + .map(|&slot| { + if slot == null_slot { + None + } else { + Some(K::Native::usize_as(remap[slot])) + } + }) + .collect(); + + // Compact values array: drop the null slot so key indices stay tight. + let compact_indices: Int64Array = (0..n) + .filter(|&i| i != null_slot) + .map(|i| i as i64) + .collect(); + let compact = + take(&*values, &compact_indices, None).expect("compact values in into_dict"); + Arc::new(DictionaryArray::::new(keys, compact)) + } + + // https://github.com/apache/datafusion/issues/23127 + // Null groups emit a null key (None), not a slot index, so the null inner + // slot never consumes a key index regardless of its position in inner. + fn check_key_overflow(&self) -> Result<()> { + let non_null_count = self.inner.len() - self.null_inner_slot.is_some() as usize; + if !Self::key_type_fits(non_null_count) { + return exec_err!( + "Dictionary key type {:?} cannot represent {} distinct values", + K::DATA_TYPE, + non_null_count + ); + } + Ok(()) + } + + fn key_type_fits(num_values: usize) -> bool { + let max: usize = match K::DATA_TYPE { + DataType::Int8 => i8::MAX as usize, + DataType::Int16 => i16::MAX as usize, + DataType::Int32 => i32::MAX as usize, + DataType::Int64 => i64::MAX as usize, + DataType::UInt8 => u8::MAX as usize, + DataType::UInt16 => u16::MAX as usize, + DataType::UInt32 => u32::MAX as usize, + DataType::UInt64 => usize::MAX, + _ => return false, + }; + num_values == 0 || num_values - 1 <= max + } + + fn hash_values(&mut self, values: &ArrayRef) { + self.val_hashes.clear(); + self.val_hashes.resize(values.len(), 0); + create_hashes( + std::slice::from_ref(values), + &self.random_state, + &mut self.val_hashes, + ) + .unwrap(); + } + + fn find_or_insert_value( + &mut self, + dict_values: &ArrayRef, + val_idx: usize, + hash: u64, + ) -> Result { + let inner = &*self.inner; + let existing = self + .value_dedup + .find(hash, |&(entry_hash, slot)| { + entry_hash == hash && inner.equal_to(slot, dict_values, val_idx) + }) + .map(|&(_, slot)| slot); + + match existing { + Some(slot) => Ok(slot), + None => { + let slot = self.inner.len(); + self.inner.append_val(dict_values, val_idx)?; + self.value_dedup.insert_accounted( + (hash, slot), + |&(entry_hash, _)| entry_hash, + &mut self.value_dedup_size, + ); + Ok(slot) + } + } + } + + fn find_or_insert_null(&mut self) -> Result { + if let Some(slot) = self.null_inner_slot { + return Ok(slot); + } + let slot = self.inner.len(); + self.inner.append_val(&self.null_array, 0)?; + self.null_inner_slot = Some(slot); + Ok(slot) + } + + fn build_lookup_table( + &self, + dict_values: &ArrayRef, + val_hashes: &[u64], + ) -> Vec { + let num_distinct = dict_values.len(); + let mut table = vec![usize::MAX; num_distinct + 1]; + let inner = &*self.inner; + for val_idx in 0..num_distinct { + if dict_values.is_null(val_idx) { + table[val_idx] = self.null_inner_slot.unwrap_or(usize::MAX); + } else { + let hash = val_hashes[val_idx]; + if let Some(&(_, slot)) = + self.value_dedup.find(hash, |&(entry_hash, slot)| { + entry_hash == hash && inner.equal_to(slot, dict_values, val_idx) + }) + { + table[val_idx] = slot; + } + } + } + table[num_distinct] = self.null_inner_slot.unwrap_or(usize::MAX); + table + } + + /// Per-row fallback for `vectorized_equal_to` used when the number of rows + /// to check is smaller than the dictionary cardinality, making the O(D) + /// lookup-table build more expensive than direct value comparison. + /// + /// `#[cold]` + `#[inline(never)]` keeps this code out of the hot + /// lookup-table loops in `vectorized_equal_to` so LLVM can pipeline them. + #[cold] + #[inline(never)] + fn equal_to_per_row( + &self, + lhs_rows: &[usize], + dict_values: &ArrayRef, + dict: &DictionaryArray, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + let group_to_inner = self.group_to_inner.as_slice(); + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + if !equal_to_results.get_bit(idx) { + continue; + } + let lhs_slot = group_to_inner[lhs_row]; + let equal = match dict.key(rhs_row) { + None => self.inner.equal_to(lhs_slot, &self.null_array, 0), + Some(val_idx) if dict_values.is_null(val_idx) => { + self.inner.equal_to(lhs_slot, &self.null_array, 0) + } + Some(val_idx) => self.inner.equal_to(lhs_slot, dict_values, val_idx), + }; + if !equal { + equal_to_results.set_bit(idx, false); + } + } + } +} + +impl GroupColumn + for DictionaryGroupValuesColumn +{ + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + let lhs_slot = self.group_to_inner[lhs_row]; + let dict = array.as_dictionary::(); + match dict.key(rhs_row) { + None => self.inner.equal_to(lhs_slot, &self.null_array, 0), + Some(val_idx) if dict.values().is_null(val_idx) => { + self.inner.equal_to(lhs_slot, &self.null_array, 0) + } + Some(val_idx) => self.inner.equal_to(lhs_slot, dict.values(), val_idx), + } + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + let dict = array.as_dictionary::(); + let inner_slot = match dict.key(row) { + None => self.find_or_insert_null()?, + Some(val_idx) if dict.values().is_null(val_idx) => { + self.find_or_insert_null()? + } + Some(val_idx) => { + let dict_values = dict.values(); + let single = dict_values.slice(val_idx, 1); + self.hash_values(&single); + self.find_or_insert_value(dict_values, val_idx, self.val_hashes[0])? + } + }; + self.group_to_inner.push(inner_slot); + self.check_key_overflow() + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + let dict = array.as_dictionary::(); + let dict_keys = dict.keys(); + let dict_values = dict.values(); + let num_distinct = dict_values.len(); + + // The fallback is in a separate #[cold] function so its code does not + // appear inline here and cannot prevent LLVM from pipelining / unrolling + // the hot lookup-table loops below. + if rhs_rows.len() < num_distinct { + self.equal_to_per_row( + lhs_rows, + dict_values, + dict, + rhs_rows, + equal_to_results, + ); + return; + } + + let mut val_hashes = vec![0u64; dict_values.len()]; + create_hashes( + std::slice::from_ref(dict_values), + &self.random_state, + &mut val_hashes, + ) + .unwrap(); + let lookup = self.build_lookup_table(dict_values, &val_hashes); + + let group_to_inner = self.group_to_inner.as_slice(); + + if dict_keys.null_count() == 0 { + // No null keys : skip the get_bit guard: we only ever write false, + // so overwriting an already-false bit is a no-op. + let raw_keys = dict_keys.values(); + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + let rhs_slot = lookup[raw_keys[rhs_row].as_usize()]; + if rhs_slot == usize::MAX || group_to_inner[lhs_row] != rhs_slot { + equal_to_results.set_bit(idx, false); + } + } + } else { + let null_buf = dict_keys.nulls().unwrap(); + let raw_keys = dict_keys.values(); + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + if equal_to_results.get_bit(idx) { + let val_idx = if null_buf.is_null(rhs_row) { + num_distinct + } else { + raw_keys[rhs_row].as_usize() + }; + let rhs_slot = lookup[val_idx]; + if rhs_slot == usize::MAX || group_to_inner[lhs_row] != rhs_slot { + equal_to_results.set_bit(idx, false); + } + } + } + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + let dict = array.as_dictionary::(); + let dict_keys = dict.keys(); + let dict_values = dict.values(); + let num_distinct = dict_values.len(); + + self.hash_values(dict_values); + self.val_to_inner.clear(); + self.val_to_inner.resize(num_distinct, usize::MAX); + + self.group_to_inner.try_reserve(rows.len()).map_err(|e| { + DataFusionError::ArrowError( + Box::new(ArrowError::MemoryError(e.to_string())), + None, + ) + })?; + + if dict_keys.null_count() == 0 { + let raw_keys = dict_keys.values(); + for &row in rows { + let val_idx = raw_keys[row].as_usize(); + if self.val_to_inner[val_idx] == usize::MAX { + // A non-null key can still point to a null value in the values array. + self.val_to_inner[val_idx] = if dict_values.is_null(val_idx) { + self.find_or_insert_null()? + } else { + self.find_or_insert_value( + dict_values, + val_idx, + self.val_hashes[val_idx], + )? + }; + } + self.group_to_inner.push(self.val_to_inner[val_idx]); + } + } else { + let raw_keys = dict_keys.values(); + let null_buf = dict_keys.nulls().unwrap(); + for &row in rows { + let slot = if null_buf.is_null(row) { + self.find_or_insert_null()? + } else { + let val_idx = raw_keys[row].as_usize(); + if self.val_to_inner[val_idx] == usize::MAX { + self.val_to_inner[val_idx] = if dict_values.is_null(val_idx) { + self.find_or_insert_null()? + } else { + self.find_or_insert_value( + dict_values, + val_idx, + self.val_hashes[val_idx], + )? + }; + } + self.val_to_inner[val_idx] + }; + self.group_to_inner.push(slot); + } + } + + self.check_key_overflow() + } + + fn len(&self) -> usize { + self.group_to_inner.len() + } + + fn size(&self) -> usize { + self.inner.size() + + self.value_dedup_size + + self.group_to_inner.capacity() * size_of::() + + self.val_to_inner.capacity() * size_of::() + + self.val_hashes.capacity() * size_of::() + + self.null_array.get_array_memory_size() + + size_of::() + } + + fn build(self: Box) -> ArrayRef { + let null_inner_slot = self.null_inner_slot; + let values = self.inner.build(); + Self::into_dict(values, &self.group_to_inner, null_inner_slot) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + let old_inner_len = self.inner.len(); + let all_inner_values = self.inner.take_n(old_inner_len); + + let mut emit_old_to_new = vec![usize::MAX; old_inner_len]; + let mut emit_new_to_old: Vec = Vec::new(); + for &old in &self.group_to_inner[..n] { + // Null groups emit a null key (None) and need no slot in the + // values array, so excluding them keeps key indices tight and + // prevents overflow at key-type capacity. + if all_inner_values.is_null(old) { + continue; + } + if emit_old_to_new[old] == usize::MAX { + emit_old_to_new[old] = emit_new_to_old.len(); + emit_new_to_old.push(old); + } + } + let emit_indices = + Int64Array::from_iter(emit_new_to_old.iter().map(|&i| i as i64)); + let compact_emit_values = + take(&*all_inner_values, &emit_indices, None).expect("take emit values"); + let emitted_keys: PrimitiveArray = self.group_to_inner[..n] + .iter() + .map(|&old| { + if all_inner_values.is_null(old) { + None + } else { + Some(K::Native::usize_as(emit_old_to_new[old])) + } + }) + .collect(); + let emitted: ArrayRef = + Arc::new(DictionaryArray::::new(emitted_keys, compact_emit_values)); + + // Null deferred to last so null_inner_slot is always the highest index + // and check_key_overflow can subtract it without a false overflow. + let remaining = self.group_to_inner[n..].to_vec(); + let mut old_to_new = vec![usize::MAX; old_inner_len]; + let mut new_to_old: Vec = Vec::new(); + let mut null_old_slot: Option = None; + for &old in &remaining { + if all_inner_values.is_null(old) { + if null_old_slot.is_none() { + null_old_slot = Some(old); + } + continue; + } + if old_to_new[old] == usize::MAX { + old_to_new[old] = new_to_old.len(); + new_to_old.push(old); + } + } + if let Some(old) = null_old_slot { + old_to_new[old] = new_to_old.len(); + new_to_old.push(old); + } + + self.value_dedup = HashTable::new(); + self.value_dedup_size = 0; + self.null_inner_slot = None; + + self.hash_values(&all_inner_values); + + for (new_slot, &old_slot) in new_to_old.iter().enumerate() { + if all_inner_values.is_null(old_slot) { + self.inner + .append_val(&self.null_array, 0) + .expect("append null failed in take_n"); + self.null_inner_slot = Some(new_slot); + } else { + self.inner + .append_val(&all_inner_values, old_slot) + .expect("append value failed in take_n"); + self.value_dedup.insert_accounted( + (self.val_hashes[old_slot], new_slot), + |&(entry_hash, _)| entry_hash, + &mut self.value_dedup_size, + ); + } + } + + self.group_to_inner = remaining.iter().map(|&old| old_to_new[old]).collect(); + self.check_key_overflow().expect("key overflow in take_n"); + + emitted + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::aggregates::group_values::multi_group_by::bytes::ByteGroupValueBuilder; + use arrow::array::{ + Array, ArrayRef, BooleanBufferBuilder, DictionaryArray, Int32Array, StringArray, + UInt8Array, + }; + use arrow::compute::cast; + use arrow::datatypes::{DataType, Int8Type, Int32Type, UInt8Type}; + use datafusion_physical_expr::binary_map::OutputType; + use std::sync::Arc; + + fn utf8_col() -> DictionaryGroupValuesColumn { + let f = Field::new("", DataType::Utf8, true); + DictionaryGroupValuesColumn::new( + Box::new(ByteGroupValueBuilder::::new(OutputType::Utf8)), + &f, + ) + } + + fn int8_col() -> DictionaryGroupValuesColumn { + let f = Field::new("", DataType::Utf8, true); + DictionaryGroupValuesColumn::new( + Box::new(ByteGroupValueBuilder::::new(OutputType::Utf8)), + &f, + ) + } + + fn uint8_col() -> DictionaryGroupValuesColumn { + let f = Field::new("", DataType::Utf8, true); + DictionaryGroupValuesColumn::new( + Box::new(ByteGroupValueBuilder::::new(OutputType::Utf8)), + &f, + ) + } + + fn i32_dict(keys: &[Option], values: &[Option<&str>]) -> ArrayRef { + Arc::new(DictionaryArray::::new( + Int32Array::from(keys.to_vec()), + Arc::new(StringArray::from(values.to_vec())), + )) + } + + fn i8_dict(keys: &[Option], values: &[Option<&str>]) -> ArrayRef { + use arrow::array::Int8Array; + Arc::new(DictionaryArray::::new( + Int8Array::from(keys.to_vec()), + Arc::new(StringArray::from(values.to_vec())), + )) + } + + fn u8_dict(keys: &[Option], values: &[Option<&str>]) -> ArrayRef { + Arc::new(DictionaryArray::::new( + UInt8Array::from(keys.to_vec()), + Arc::new(StringArray::from(values.to_vec())), + )) + } + + fn str_values(arr: &ArrayRef) -> Vec> { + let plain = cast(arr.as_ref(), &DataType::Utf8).unwrap(); + plain + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .map(|v| v.map(|s| s.to_owned())) + .collect() + } + + fn bool_vec(buf: &BooleanBufferBuilder) -> Vec { + (0..buf.len()).map(|i| buf.get_bit(i)).collect() + } + + fn all_true(len: usize) -> BooleanBufferBuilder { + let mut buf = BooleanBufferBuilder::new(len); + buf.append_n(len, true); + buf + } + + // Builds an Int8-keyed dict of `end-start` distinct strings "v{start}".."v{end-1}". + fn distinct_i8_dict(start: usize, end: usize) -> ArrayRef { + let strs: Vec = (start..end).map(|i| format!("v{i}")).collect(); + let refs: Vec> = strs.iter().map(|s| Some(s.as_str())).collect(); + i8_dict( + &(0..strs.len()).map(|i| Some(i as i8)).collect::>(), + &refs, + ) + } + + // Builds a UInt8-keyed dict of `count` distinct strings "u0".."u{count-1}". + fn distinct_u8_dict(count: usize) -> ArrayRef { + let strs: Vec = (0..count).map(|i| format!("u{i}")).collect(); + let refs: Vec> = strs.iter().map(|s| Some(s.as_str())).collect(); + u8_dict( + &(0..count).map(|i| Some(i as u8)).collect::>(), + &refs, + ) + } + + #[test] + fn repeated_values_are_deduplicated_in_inner_store() { + let mut col = utf8_col(); + let arr = i32_dict( + &[Some(0), Some(1), Some(0), Some(1), Some(0)], + &[Some("a"), Some("b")], + ); + col.vectorized_append(&arr, &[0, 1, 2, 3, 4]).unwrap(); + let out = Box::new(col).build(); + assert_eq!(out.as_dictionary::().values().len(), 2); + assert_eq!( + str_values(&out), + vec![ + Some("a".into()), + Some("b".into()), + Some("a".into()), + Some("b".into()), + Some("a".into()), + ] + ); + } + + #[test] + fn null_key_and_null_valued_entry_both_map_to_null_group() { + let mut col = utf8_col(); + let input = i32_dict(&[None, Some(0), Some(1)], &[None, Some("b")]); + for row in 0..3 { + col.append_val(&input, row).unwrap(); + } + assert!(col.equal_to(0, &input, 1)); + assert!(!col.equal_to(0, &input, 2)); + let out = Box::new(col).build(); + assert_eq!(out.as_dictionary::().values().len(), 1); + assert_eq!(str_values(&out), vec![None, None, Some("b".into())]); + } + + #[test] + fn take_n_compacts_emitted_values_and_remaps_remaining_slots() { + let mut col = utf8_col(); + let b1 = i32_dict( + &[Some(0), Some(1), None, Some(2)], + &[Some("a"), Some("b"), Some("c")], + ); + col.vectorized_append(&b1, &[0, 1, 2, 3]).unwrap(); + + let emitted = col.take_n(2); + assert_eq!(emitted.as_dictionary::().values().len(), 2); + assert_eq!( + str_values(&emitted), + vec![Some("a".into()), Some("b".into())] + ); + + let b2 = i32_dict(&[None, Some(0)], &[Some("z")]); + col.vectorized_append(&b2, &[0, 1]).unwrap(); + + let mut buf = all_true(2); + col.vectorized_equal_to(&[0, 1], &b2, &[0, 1], &mut buf); + assert_eq!(bool_vec(&buf), vec![true, false]); + + let out = Box::new(col).build(); + assert_eq!( + str_values(&out), + vec![None, Some("c".into()), None, Some("z".into())] + ); + } + + #[test] + fn vectorized_equal_to_does_not_use_stale_hashes_from_prior_append() { + let mut col = utf8_col(); + col.vectorized_append(&i32_dict(&[Some(0)], &[Some("a"), Some("b")]), &[0]) + .unwrap(); + let batch2 = i32_dict(&[Some(1)], &[Some("z"), Some("a")]); + let mut buf = all_true(1); + col.vectorized_equal_to(&[0], &batch2, &[0], &mut buf); + assert_eq!(bool_vec(&buf), vec![true]); + } + + #[test] + fn null_does_not_consume_a_key_slot_int8_null_first_mid_and_last() { + let rows128 = (0..128).collect::>(); + + let mut col = int8_col(); // null-last: 128 non-null + null — ok + col.vectorized_append(&distinct_i8_dict(0, 128), &rows128) + .unwrap(); + col.append_val(&i8_dict(&[None], &[Some("x")]), 0).unwrap(); + + let mut col = int8_col(); // null-first: null + 128 non-null — ok; 129th — error + col.append_val(&i8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_i8_dict(0, 128), &rows128) + .unwrap(); + assert!( + col.append_val(&i8_dict(&[Some(0)], &[Some("overflow")]), 0) + .is_err() + ); + + let mut col = int8_col(); // null-mid: 100 + null + 28 = 128 total — ok; 129th — error + col.vectorized_append(&distinct_i8_dict(0, 100), &(0..100).collect::>()) + .unwrap(); + col.append_val(&i8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_i8_dict(100, 128), &(0..28).collect::>()) + .unwrap(); + assert!( + col.append_val(&i8_dict(&[Some(0)], &[Some("v128")]), 0) + .is_err() + ); + + let mut col = int8_col(); // build() null-first: 128 values (null excluded), null → None + col.append_val(&i8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_i8_dict(0, 128), &rows128) + .unwrap(); + let out = Box::new(col).build(); + assert_eq!(out.as_dictionary::().values().len(), 128); + assert_eq!(str_values(&out)[0], None); + assert_eq!(str_values(&out)[1], Some("v0".into())); + } + + #[test] + fn null_does_not_consume_a_key_slot_uint8_null_first_and_last() { + let rows256 = (0..256).collect::>(); + + let mut col = uint8_col(); // null-first: null + 256 non-null — ok; 257th — error + col.append_val(&u8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_u8_dict(256), &rows256) + .unwrap(); + assert!( + col.append_val(&u8_dict(&[Some(0)], &[Some("overflow")]), 0) + .is_err() + ); + + let mut col = uint8_col(); // build() null-first: 256 values (null excluded), last correct + col.append_val(&u8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_u8_dict(256), &rows256) + .unwrap(); + let out = Box::new(col).build(); + assert_eq!(out.as_dictionary::().values().len(), 256); + assert_eq!(str_values(&out)[0], None); + assert_eq!(str_values(&out)[256], Some("u255".into())); + } + + #[test] + fn take_n_null_does_not_steal_key_slot_at_capacity() { + let rows128 = (0..128).collect::>(); + + // Int8 null-first + 128 non-null; emit all 129 — no panic, null → None + let mut col = int8_col(); + col.append_val(&i8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_i8_dict(0, 128), &rows128) + .unwrap(); + let emitted = col.take_n(129); + assert!(emitted.as_dictionary::().key(0).is_none()); + assert_eq!(str_values(&emitted)[1], Some("v0".into())); + assert_eq!(str_values(&emitted)[128], Some("v127".into())); + + // UInt8 null-first + 256 non-null; emit all 257 — last must be "u255" not "u0" (wrap guard) + let mut col = uint8_col(); + col.append_val(&u8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_u8_dict(256), &(0..256).collect::>()) + .unwrap(); + let emitted = col.take_n(257); + assert!(emitted.as_dictionary::().key(0).is_none()); + assert_eq!(str_values(&emitted)[1], Some("u0".into())); + assert_eq!(str_values(&emitted)[256], Some("u255".into())); + } + + #[test] + fn take_n_repeated_emissions_null_at_int8_capacity() { + let rows128 = (0..128).collect::>(); + let mut col = int8_col(); + col.vectorized_append(&distinct_i8_dict(0, 128), &rows128) + .unwrap(); + col.append_val(&i8_dict(&[None], &[Some("x")]), 0).unwrap(); + col.vectorized_append(&distinct_i8_dict(0, 128), &rows128) + .unwrap(); + + let first_half = col.take_n(64); + assert_eq!(str_values(&first_half)[0], Some("v0".into())); + assert_eq!(str_values(&first_half)[63], Some("v63".into())); + assert_eq!(first_half.as_dictionary::().values().len(), 64); + + let second_half = col.take_n(64); + assert_eq!(str_values(&second_half)[0], Some("v64".into())); + assert_eq!(str_values(&second_half)[63], Some("v127".into())); + + let null_group = col.take_n(1); + assert!(null_group.as_dictionary::().key(0).is_none()); + + let out = Box::new(col).build(); + assert_eq!(str_values(&out)[0], Some("v0".into())); + assert_eq!(str_values(&out)[127], Some("v127".into())); + assert_eq!(out.as_dictionary::().values().len(), 128); + } +} 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 2923726ad1b66..692af387e6034 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 @@ -20,6 +20,7 @@ mod boolean; mod bytes; pub mod bytes_view; +mod dictionary; pub mod primitive; pub mod row_backed; @@ -32,7 +33,6 @@ use crate::aggregates::group_values::multi_group_by::{ row_backed::RowsGroupColumn, }; 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, @@ -43,7 +43,7 @@ use arrow::datatypes::{ }; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; -use datafusion_common::{Result, internal_datafusion_err, not_impl_err}; +use datafusion_common::{Result, not_impl_err}; use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; use datafusion_expr::EmitTo; use datafusion_physical_expr::binary_map::OutputType; @@ -979,7 +979,7 @@ fn group_column_supported_type(data_type: &DataType) -> bool { | DataType::Utf8View | DataType::BinaryView | DataType::Boolean - ) + ) || matches!(data_type, DataType::Dictionary(_,v ) if group_column_supported_type(v)) } /// Build a [`GroupColumn`] for a single schema field. @@ -1091,6 +1091,33 @@ fn make_group_column(field: &Field) -> Result> { v.push(Box::new(BooleanGroupValueBuilder::::new())); } } + DataType::Dictionary(ref key_dt, ref value_dt) => { + let new_field = Field::new("", *value_dt.clone(), true); + let inner = make_group_column(&new_field)?; + macro_rules! dict_col { + ($T:ty) => { + Box::new(dictionary::DictionaryGroupValuesColumn::<$T>::new( + inner, &new_field, + )) + }; + } + let col: Box = match key_dt.as_ref() { + DataType::Int8 => dict_col!(Int8Type), + DataType::Int16 => dict_col!(Int16Type), + DataType::Int32 => dict_col!(Int32Type), + DataType::Int64 => dict_col!(Int64Type), + DataType::UInt8 => dict_col!(UInt8Type), + DataType::UInt16 => dict_col!(UInt16Type), + DataType::UInt32 => dict_col!(UInt32Type), + DataType::UInt64 => dict_col!(UInt64Type), + _ => { + return not_impl_err!( + "Dictionary key type {key_dt} not supported in GroupValuesColumn" + ); + } + }; + v.push(col) + } // Generic fallback for nested types (Struct / List / LargeList / // FixedSizeList, recursively) that lack a type-specialized builder but // can be encoded by arrow's row format. This is what lets a mixed @@ -1140,7 +1167,7 @@ impl GroupValues for GroupValuesColumn { } fn emit(&mut self, emit_to: EmitTo) -> Result> { - let mut output = match emit_to { + let output = match emit_to { EmitTo::All => { // Replace the column builders with a fresh set so the // aggregator is immediately reusable after the drain. @@ -1230,20 +1257,6 @@ impl GroupValues for GroupValuesColumn { } }; - // TODO: Materialize dictionaries in group keys (#7647) - for (field, array) in self.schema.fields.iter().zip(&mut output) { - let expected = field.data_type(); - if let DataType::Dictionary(_, v) = expected { - let actual = array.data_type(); - if v.as_ref() != actual { - return Err(internal_datafusion_err!( - "Converted group rows expected dictionary of {v} got {actual}" - )); - } - *array = cast(array.as_ref(), expected)?; - } - } - Ok(output) } @@ -1578,6 +1591,267 @@ mod tests { assert_eq!(v, vec![3, 4, 5, 6]); } + #[test] + fn group_column_supported_type_matches_make_group_column() { + let supported_cases: Vec = vec![ + DataType::Int8, + DataType::Int64, + DataType::UInt64, + DataType::Float32, + DataType::Float64, + DataType::Float16, + DataType::Decimal128(38, 10), + DataType::Decimal256(76, 10), + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Binary, + DataType::LargeBinary, + DataType::BinaryView, + DataType::FixedSizeBinary(16), + // Zero-width FixedSizeBinary is valid per the Arrow spec + DataType::FixedSizeBinary(0), + DataType::Boolean, + DataType::Date32, + DataType::Date64, + DataType::Time32(arrow::datatypes::TimeUnit::Second), + DataType::Time32(arrow::datatypes::TimeUnit::Millisecond), + DataType::Time64(arrow::datatypes::TimeUnit::Microsecond), + DataType::Time64(arrow::datatypes::TimeUnit::Nanosecond), + DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None), + DataType::Duration(arrow::datatypes::TimeUnit::Second), + DataType::Duration(arrow::datatypes::TimeUnit::Millisecond), + DataType::Duration(arrow::datatypes::TimeUnit::Microsecond), + DataType::Duration(arrow::datatypes::TimeUnit::Nanosecond), + DataType::Interval(arrow::datatypes::IntervalUnit::YearMonth), + DataType::Interval(arrow::datatypes::IntervalUnit::DayTime), + DataType::Interval(arrow::datatypes::IntervalUnit::MonthDayNano), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int64)), + DataType::Dictionary( + Box::new(DataType::UInt16), + Box::new(DataType::LargeUtf8), + ), + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Timestamp( + arrow::datatypes::TimeUnit::Nanosecond, + None, + )), + ), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Float16)), + ]; + + for dt in &supported_cases { + assert!( + group_column_supported_type(dt), + "expected group_column_supported_type=true for {dt:?}" + ); + let field = Field::new("col", dt.clone(), true); + make_group_column(&field).unwrap_or_else(|e| { + panic!( + "group_column_supported_type accepted {dt:?} but make_group_column rejected: {e}" + ) + }); + } + + let unsupported_cases: Vec = vec![ + // Invalid Time-unit combinations: Time32 is defined only for + // Second / Millisecond and Time64 only for Microsecond / + // Nanosecond. The TimeUnit enum allows constructing the other + // combinations programmatically, but they are not valid Arrow + // types and must be rejected by both group_column_supported_type + // and the dispatcher. + DataType::Time64(arrow::datatypes::TimeUnit::Second), + DataType::Time64(arrow::datatypes::TimeUnit::Millisecond), + DataType::Time32(arrow::datatypes::TimeUnit::Microsecond), + DataType::Time32(arrow::datatypes::TimeUnit::Nanosecond), + // A negative width is representable in the DataType but is not + // a valid Arrow type; no array can be constructed for it. + DataType::FixedSizeBinary(-5), + ]; + + for dt in &unsupported_cases { + assert!( + !group_column_supported_type(dt), + "expected group_column_supported_type=false for {dt:?}" + ); + let field = Field::new("col", dt.clone(), true); + assert!( + make_group_column(&field).is_err(), + "group_column_supported_type rejected {dt:?} but make_group_column accepted it" + ); + } + } + + // `Duration` group keys stay on the `GroupValuesColumn` fast path, dedup + // (including nulls), and round-trip with the `Duration` type preserved. + #[test] + fn test_group_values_column_duration() { + use arrow::datatypes::TimeUnit; + + let schema = Arc::new(Schema::new(vec![ + Field::new("d", DataType::Duration(TimeUnit::Microsecond), true), + Field::new("i", DataType::Int64, true), + ])); + assert!(supported_schema(&schema)); + let mut group_values = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + + // (d, i) rows, where row 3 repeats row 0 and row 4 repeats the null pair. + let d: ArrayRef = Arc::new(DurationMicrosecondArray::from(vec![ + Some(10), + None, + Some(20), + Some(10), + None, + ])); + let i: ArrayRef = Arc::new(Int64Array::from(vec![ + Some(1), + None, + Some(2), + Some(1), + None, + ])); + let mut groups = Vec::new(); + group_values.intern(&[d, i], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 2, 0, 1]); + + let emitted = group_values.emit(EmitTo::All).unwrap(); + assert_eq!(emitted.len(), 2); + // The Duration column round-trips as Duration on emit, not bare i64. + assert_eq!( + emitted[0].data_type(), + &DataType::Duration(TimeUnit::Microsecond) + ); + let actual = emitted[0] + .as_any() + .downcast_ref::() + .expect("emitted column should be a DurationMicrosecondArray"); + // Three groups in first-seen order: 10, null, 20. + assert_eq!(actual.len(), 3); + assert_eq!(actual.value(0), 10); + assert!(actual.is_null(1)); + 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]); + } + + // `(Interval, Int32)` keys for each of the three interval units: null keys + // dedup, the Int32 key splits equal intervals, and emit gives back Interval. + #[test] + fn test_group_values_column_interval() { + use arrow::datatypes::{ + ArrowPrimitiveType, IntervalDayTime, IntervalDayTimeType, + IntervalMonthDayNano, IntervalMonthDayNanoType, IntervalUnit, + IntervalYearMonthType, + }; + + fn check(unit: IntervalUnit, value: T::Native) { + let schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Interval(unit), true), + Field::new("n", DataType::Int32, true), + ])); + assert!(supported_schema(&schema), "{unit:?} schema not supported"); + let mut group_values = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + + let i: ArrayRef = Arc::new(PrimitiveArray::::from_iter([ + Some(value), + None, + Some(value), + None, + Some(value), + ])); + let n: ArrayRef = Arc::new(Int32Array::from(vec![3, 3, 3, 3, 4])); + let mut groups = Vec::new(); + group_values.intern(&[i, n], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 0, 1, 2], "{unit:?}"); + + let emitted = group_values.emit(EmitTo::All).unwrap(); + assert_eq!(emitted.len(), 2); + // The emitted key keeps its Interval type, not the bare native. + assert_eq!(emitted[0].data_type(), &DataType::Interval(unit)); + let actual = emitted[0] + .as_any() + .downcast_ref::>() + .unwrap_or_else(|| panic!("emitted column should be a {unit:?} array")); + // Three groups in first-seen order: value, null, value (n=4). + assert_eq!(actual.len(), 3, "{unit:?}"); + assert_eq!(actual.value(0), value, "{unit:?}"); + assert!(actual.is_null(1), "{unit:?}"); + assert_eq!(actual.value(2), value, "{unit:?}"); + 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], "{unit:?}"); + } + + check::(IntervalUnit::YearMonth, 13); + check::(IntervalUnit::DayTime, IntervalDayTime::new(1, 500)); + check::( + IntervalUnit::MonthDayNano, + IntervalMonthDayNano::new(1, 0, 0), + ); } + #[test] fn test_split_vec_min_alloc_split_off_branch() { // remaining < n → split_off+replace branch (allocates remaining elements) @@ -1612,6 +1886,56 @@ mod tests { assert_eq!(v, vec![1, 2, 3]); } + // https://github.com/apache/datafusion/issues/23127 + // validate DictionaryGroupColumn deduplicates values — only k distinct keys appear + // in the values array even when there are more than 128 groups total. + #[test] + fn multi_col_groupby_dict_many_groups_two_values() { + use arrow::array::{AsArray, DictionaryArray, Int8Array}; + use arrow::datatypes::Int8Type; + + let n_groups = 129_usize; + let dict_vocab: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog"])); + + // Each row has a unique label (forcing a new group) and alternates + // between the two dictionary values. Int8 keys are used; only 2 + // distinct values exist so the key type never overflows. + let labels: ArrayRef = Arc::new(StringArray::from( + (0..n_groups).map(|i| format!("g{i}")).collect::>(), + )); + let dict_keys = Int8Array::from( + (0..n_groups) + .map(|i| Some((i % 2) as i8)) + .collect::>(), + ); + let categories: ArrayRef = Arc::new(DictionaryArray::::new( + dict_keys, + Arc::clone(&dict_vocab), + )); + + let schema = Arc::new(Schema::new(vec![ + Field::new("label", DataType::Utf8, false), + Field::new( + "category", + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), + false, + ), + ])); + + let mut gv = GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + gv.intern(&[labels, categories], &mut vec![]).unwrap(); + let out = gv.emit(EmitTo::All).unwrap(); + + assert_eq!(out[0].len(), n_groups); + assert!(matches!( + out[1].data_type(), + DataType::Dictionary(k, v) + if k.as_ref() == &DataType::Int8 && v.as_ref() == &DataType::Utf8 + )); + // Both vectorized and streaming paths now deduplicate dict values. + assert_eq!(out[1].as_dictionary::().values().len(), 2); + } + #[test] fn test_intern_for_vectorized_group_values() { let data_set = VectorizedTestDataSet::new(); From 8d1a677f98beb624b103e7b586f27a1d286e9bf7 Mon Sep 17 00:00:00 2001 From: Max Burke Date: Tue, 11 Aug 2026 06:07:56 -0700 Subject: [PATCH 3/4] Add FixedSizeBinary support for MultiGroupBy (#23646) - Closes #23645 - part of https://github.com/apache/datafusion/issues/22715 Multi-Group-By has cases for regular Binary/LargeBinary types, but not FixedSizeBinary Yes. No Co-authored-by: Claude Fable 5 (cherry picked from commit 3f0a95336a3f005a182971aff887ce8a7a6661fa) --- .../multi_group_by/fixed_size_binary.rs | 514 ++++++++++++++++++ .../group_values/multi_group_by/mod.rs | 102 +++- .../sqllogictest/test_files/aggregate.slt | 25 + 3 files changed, 634 insertions(+), 7 deletions(-) create mode 100644 datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs new file mode 100644 index 0000000000000..3ad19588d25d2 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs @@ -0,0 +1,514 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::aggregates::group_values::multi_group_by::{ + GroupColumn, Nulls, nulls_equal_to, split_vec_min_alloc, +}; +use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeBinaryArray, +}; +use arrow::buffer::{Buffer, NullBuffer}; +use datafusion_common::utils::proxy::VecAllocExt; +use datafusion_common::{Result, exec_datafusion_err}; +use std::sync::Arc; + +/// An implementation of [`GroupColumn`] for `FixedSizeBinary` values +/// +/// Stores the group values in a single flat buffer, `byte_width` bytes per +/// value, in a way that allows: +/// +/// 1. Efficient comparison of incoming rows to existing rows +/// 2. Efficient construction of the final output array (the buffer is handed +/// to [`FixedSizeBinaryArray`] as-is, no offsets needed) +/// +/// Null values occupy `byte_width` zeroed bytes in the buffer so that the +/// value of row `i` is always stored at `i * byte_width..(i + 1) * byte_width`. +pub struct FixedSizeBinaryGroupValueBuilder { + /// The width in bytes of each value, from `DataType::FixedSizeBinary` + byte_width: usize, + /// The flattened group values, `byte_width` bytes per value + buffer: Vec, + /// The number of group values stored + /// + /// Tracked explicitly rather than derived from `buffer.len()` because + /// `byte_width` may be `0` + len: usize, + /// Null state (null rows still occupy `byte_width` bytes in `buffer`) + nulls: MaybeNullBufferBuilder, +} + +impl FixedSizeBinaryGroupValueBuilder { + /// Create a new builder for values of `byte_width` bytes each + /// + /// `byte_width` is the width carried by `DataType::FixedSizeBinary` and + /// must be non-negative (negative widths are rejected by the dispatch in + /// `make_group_column`) + pub fn new(byte_width: i32) -> Self { + debug_assert!(byte_width >= 0); + Self { + byte_width: byte_width as usize, + buffer: Vec::new(), + len: 0, + nulls: MaybeNullBufferBuilder::new(), + } + } + + fn do_append_val_inner(&mut self, array: &FixedSizeBinaryArray, row: usize) { + if array.is_null(row) { + self.nulls.append(true); + // Null rows still occupy `byte_width` (zeroed) bytes in the + // buffer so the value offset stays a function of the row index + self.buffer.resize(self.buffer.len() + self.byte_width, 0); + } else { + self.nulls.append(false); + self.buffer.extend_from_slice(array.value(row)); + } + self.len += 1; + } + + fn do_equal_to_inner( + &self, + lhs_row: usize, + array: &FixedSizeBinaryArray, + rhs_row: usize, + ) -> bool { + let exist_null = self.nulls.is_null(lhs_row); + let input_null = array.is_null(rhs_row); + if let Some(result) = nulls_equal_to(exist_null, input_null) { + return result; + } + // Otherwise, we need to check their values + self.value(lhs_row) == array.value(rhs_row) + } + + /// return the current value of the specified row irrespective of null + /// (null rows store `byte_width` zeroed bytes) + pub fn value(&self, row: usize) -> &[u8] { + let start = row * self.byte_width; + &self.buffer[start..start + self.byte_width] + } + + /// Assemble an output array from `values` + `nulls` parts + /// + /// Uses `try_new_with_len` rather than `try_new` because the length + /// cannot be derived from the values buffer when `byte_width == 0` + fn build_array( + byte_width: usize, + values: Vec, + nulls: Option, + len: usize, + ) -> ArrayRef { + let array = FixedSizeBinaryArray::try_new_with_len( + byte_width as i32, + Buffer::from(values), + nulls, + len, + ) + .expect("buffer, nulls and len kept consistent on append"); + Arc::new(array) + } +} + +impl GroupColumn for FixedSizeBinaryGroupValueBuilder { + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + self.do_equal_to_inner(lhs_row, array.as_fixed_size_binary(), rhs_row) + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + let arr = array.as_fixed_size_binary(); + debug_assert_eq!(arr.value_size(), self.byte_width); + self.do_append_val_inner(arr, row); + Ok(()) + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + let array = array.as_fixed_size_binary(); + + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + // Has found not equal to in previous column, don't need to check + if !equal_to_results.get_bit(idx) { + continue; + } + + if !self.do_equal_to_inner(lhs_row, array, rhs_row) { + equal_to_results.set_bit(idx, false); + } + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + let arr = array.as_fixed_size_binary(); + debug_assert_eq!(arr.value_size(), self.byte_width); + + let reserve_bytes = rows.len() * self.byte_width; + self.buffer.try_reserve(reserve_bytes).map_err(|e| { + exec_datafusion_err!("failed to reserve {reserve_bytes} bytes: {e}") + })?; + + let null_count = array.null_count(); + let num_rows = array.len(); + let all_null_or_non_null = if null_count == 0 { + Nulls::None + } else if null_count == num_rows { + Nulls::All + } else { + Nulls::Some + }; + + match all_null_or_non_null { + Nulls::Some => { + for &row in rows { + self.do_append_val_inner(arr, row); + } + } + + Nulls::None => { + self.nulls.append_n(rows.len(), false); + for &row in rows { + self.buffer.extend_from_slice(arr.value(row)); + } + self.len += rows.len(); + } + + Nulls::All => { + self.nulls.append_n(rows.len(), true); + self.buffer + .resize(self.buffer.len() + rows.len() * self.byte_width, 0); + self.len += rows.len(); + } + } + + Ok(()) + } + + fn len(&self) -> usize { + self.len + } + + fn size(&self) -> usize { + self.buffer.allocated_size() + self.nulls.allocated_size() + } + + fn build(self: Box) -> ArrayRef { + let Self { + byte_width, + buffer, + len, + nulls, + } = *self; + + Self::build_array(byte_width, buffer, nulls.build(), len) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + debug_assert!(self.len >= n); + + let null_buffer = self.nulls.take_n(n); + let first_n = split_vec_min_alloc(&mut self.buffer, n * self.byte_width); + self.len -= n; + + Self::build_array(self.byte_width, first_n, null_buffer, n) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use crate::aggregates::group_values::multi_group_by::fixed_size_binary::FixedSizeBinaryGroupValueBuilder; + use arrow::array::{ArrayRef, BooleanBufferBuilder, FixedSizeBinaryArray}; + + use super::GroupColumn; + + fn make_true_buffer(n: usize) -> BooleanBufferBuilder { + let mut buf = BooleanBufferBuilder::new(n); + buf.append_n(n, true); + buf + } + + fn to_vec(buf: &BooleanBufferBuilder) -> Vec { + (0..buf.len()).map(|i| buf.get_bit(i)).collect() + } + + fn make_array(values: Vec>, byte_width: i32) -> ArrayRef { + Arc::new( + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + values.into_iter(), + byte_width, + ) + .unwrap(), + ) + } + + #[test] + fn test_fixed_size_binary_equal_to() { + let append = |builder: &mut FixedSizeBinaryGroupValueBuilder, + builder_array: &ArrayRef, + append_rows: &[usize]| { + for &index in append_rows { + builder.append_val(builder_array, index).unwrap(); + } + }; + + let equal_to = + |builder: &FixedSizeBinaryGroupValueBuilder, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + let iter = lhs_rows.iter().zip(rhs_rows.iter()); + for (idx, (&lhs_row, &rhs_row)) in iter.enumerate() { + equal_to_results + .set_bit(idx, builder.equal_to(lhs_row, input_array, rhs_row)); + } + }; + + test_fixed_size_binary_equal_to_internal(append, equal_to); + } + + #[test] + fn test_fixed_size_binary_vectorized_equal_to() { + let append = |builder: &mut FixedSizeBinaryGroupValueBuilder, + builder_array: &ArrayRef, + append_rows: &[usize]| { + builder + .vectorized_append(builder_array, append_rows) + .unwrap(); + }; + + let equal_to = + |builder: &FixedSizeBinaryGroupValueBuilder, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + builder.vectorized_equal_to( + lhs_rows, + input_array, + rhs_rows, + equal_to_results, + ); + }; + + test_fixed_size_binary_equal_to_internal(append, equal_to); + } + + fn test_fixed_size_binary_equal_to_internal(mut append: A, mut equal_to: E) + where + A: FnMut(&mut FixedSizeBinaryGroupValueBuilder, &ArrayRef, &[usize]), + E: FnMut( + &FixedSizeBinaryGroupValueBuilder, + &[usize], + &ArrayRef, + &[usize], + &mut BooleanBufferBuilder, + ), + { + // Will cover such cases: + // - exist null, input not null + // - exist null, input null; values not equal + // - exist null, input null; values equal + // - exist not null, input null + // - exist not null, input not null; values not equal + // - exist not null, input not null; values equal + + // Define FixedSizeBinaryGroupValueBuilder + let mut builder = FixedSizeBinaryGroupValueBuilder::new(3); + let builder_array = make_array( + vec![ + None, + None, + None, + Some(b"foo".as_slice()), + Some(b"bar".as_slice()), + Some(b"baz".as_slice()), + ], + 3, + ); + append(&mut builder, &builder_array, &[0, 1, 2, 3, 4, 5]); + + // Define input array; the value behind the null at row 3 happens to + // match the existing group value to make sure nulls win over values + let input_array = make_array( + vec![ + Some(b"foo".as_slice()), + None, + None, + None, + Some(b"foo".as_slice()), + Some(b"baz".as_slice()), + ], + 3, + ); + + // Check + let mut equal_to_results = make_true_buffer(builder.len()); + equal_to( + &builder, + &[0, 1, 2, 3, 4, 5], + &input_array, + &[0, 1, 2, 3, 4, 5], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(!results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(!results[3]); + assert!(!results[4]); + assert!(results[5]); + } + + #[test] + fn test_fixed_size_binary_vectorized_operation_special_case() { + // Test the special `all nulls` or `not nulls` input array case + // for vectorized append and equal to + + let mut builder = FixedSizeBinaryGroupValueBuilder::new(2); + + // All nulls input array + let all_nulls_input_array = make_array(vec![None, None, None, None, None], 2); + builder + .vectorized_append(&all_nulls_input_array, &[0, 1, 2, 3, 4]) + .unwrap(); + + let mut equal_to_results = make_true_buffer(all_nulls_input_array.len()); + builder.vectorized_equal_to( + &[0, 1, 2, 3, 4], + &all_nulls_input_array, + &[0, 1, 2, 3, 4], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(results[3]); + assert!(results[4]); + + // All not nulls input array + let all_not_nulls_input_array = make_array( + vec![ + Some(b"v1".as_slice()), + Some(b"v2".as_slice()), + Some(b"v3".as_slice()), + Some(b"v4".as_slice()), + Some(b"v5".as_slice()), + ], + 2, + ); + builder + .vectorized_append(&all_not_nulls_input_array, &[0, 1, 2, 3, 4]) + .unwrap(); + + let mut equal_to_results = make_true_buffer(all_not_nulls_input_array.len()); + builder.vectorized_equal_to( + &[5, 6, 7, 8, 9], + &all_not_nulls_input_array, + &[0, 1, 2, 3, 4], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(results[3]); + assert!(results[4]); + } + + #[test] + fn test_fixed_size_binary_take_n() { + let mut builder = FixedSizeBinaryGroupValueBuilder::new(2); + let array = make_array(vec![Some(b"aa".as_slice()), None], 2); + // aa, null, null + builder.append_val(&array, 0).unwrap(); + builder.append_val(&array, 1).unwrap(); + builder.append_val(&array, 1).unwrap(); + + // (aa, null) remaining: null + let output = builder.take_n(2); + assert_eq!(&output, &array); + assert_eq!(builder.len(), 1); + + // null, aa, null, aa + builder.append_val(&array, 0).unwrap(); + builder.append_val(&array, 1).unwrap(); + builder.append_val(&array, 0).unwrap(); + + // (null, aa) remaining: (null, aa) + let output = builder.take_n(2); + let expected = make_array(vec![None, Some(b"aa".as_slice())], 2); + assert_eq!(&output, &expected); + assert_eq!(builder.len(), 2); + + // take the remaining (null, aa) + let output = builder.take_n(2); + assert_eq!(&output, &expected); + assert_eq!(builder.len(), 0); + } + + #[test] + fn test_fixed_size_binary_build() { + let mut builder = FixedSizeBinaryGroupValueBuilder::new(2); + let array = make_array( + vec![Some(b"aa".as_slice()), None, Some(b"bb".as_slice())], + 2, + ); + builder.vectorized_append(&array, &[0, 1, 2]).unwrap(); + assert_eq!(builder.len(), 3); + + let output = Box::new(builder).build(); + assert_eq!(&output, &array); + } + + #[test] + fn test_zero_width_fixed_size_binary() { + // A zero byte width is valid per the Arrow spec; the builder must + // track its length without relying on the (empty) values buffer + let mut builder = FixedSizeBinaryGroupValueBuilder::new(0); + let array = make_array(vec![Some(b"".as_slice()), None, Some(b"".as_slice())], 0); + + builder.vectorized_append(&array, &[0, 1, 2]).unwrap(); + assert_eq!(builder.len(), 3); + + // Empty values compare equal, null only equals null + assert!(builder.equal_to(0, &array, 2)); + assert!(builder.equal_to(1, &array, 1)); + assert!(!builder.equal_to(1, &array, 0)); + + let output = builder.take_n(2); + let expected = make_array(vec![Some(b"".as_slice()), None], 0); + assert_eq!(&output, &expected); + assert_eq!(builder.len(), 1); + + let output = Box::new(builder).build(); + let expected = make_array(vec![Some(b"".as_slice())], 0); + assert_eq!(&output, &expected); + } +} 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 692af387e6034..d7e374f0bccb9 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 @@ -21,6 +21,7 @@ mod boolean; mod bytes; pub mod bytes_view; mod dictionary; +mod fixed_size_binary; pub mod primitive; pub mod row_backed; @@ -29,8 +30,9 @@ use std::mem::{self, size_of}; use crate::aggregates::group_values::GroupValues; use crate::aggregates::group_values::multi_group_by::{ boolean::BooleanGroupValueBuilder, bytes::ByteGroupValueBuilder, - bytes_view::ByteViewGroupValueBuilder, primitive::PrimitiveGroupValueBuilder, - row_backed::RowsGroupColumn, + bytes_view::ByteViewGroupValueBuilder, + fixed_size_binary::FixedSizeBinaryGroupValueBuilder, + primitive::PrimitiveGroupValueBuilder, row_backed::RowsGroupColumn, }; use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; use arrow::datatypes::{ @@ -964,6 +966,11 @@ fn group_column_supported_type(data_type: &DataType) -> bool { | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary + // Only non-negative widths: a negative width is not a valid + // Arrow type (no array can be constructed for it), and the + // dispatcher in `make_group_column` rejects it. Keep the two + // in lockstep. + | DataType::FixedSizeBinary(0..) | DataType::Date32 | DataType::Date64 // Only the semantically valid Time variants per the Arrow spec. @@ -1078,6 +1085,11 @@ fn make_group_column(field: &Field) -> Result> { OutputType::Binary, ))); } + // A negative width is not a valid Arrow type; it falls to the `_` + // arm below, matching `group_column_supported_type`. + DataType::FixedSizeBinary(byte_width @ 0..) => { + v.push(Box::new(FixedSizeBinaryGroupValueBuilder::new(byte_width))); + } DataType::Utf8View => { v.push(Box::new(ByteViewGroupValueBuilder::::new())); } @@ -1136,7 +1148,6 @@ fn make_group_column(field: &Field) -> Result> { Ok(v.into_iter().next().unwrap()) } - impl GroupValues for GroupValuesColumn { fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { // `try_new` and the reset points in `emit` / `clear_shrink` keep @@ -1305,7 +1316,11 @@ enum Nulls { mod tests { use std::{collections::HashMap, sync::Arc}; - use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray, StringViewArray}; + use arrow::array::{ + Array, ArrayRef, DurationMicrosecondArray, FixedSizeBinaryArray, Float16Array, + Int32Array, Int64Array, PrimitiveArray, 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; @@ -1316,8 +1331,8 @@ mod tests { }; use super::{ - GroupIndexView, group_column_supported_type, make_group_column, split_vec_min_alloc, - supported_schema, + GroupIndexView, group_column_supported_type, make_group_column, + split_vec_min_alloc, supported_schema, }; /// A mixed group-by key of several native columns plus one nested column @@ -1850,7 +1865,8 @@ mod tests { check::( IntervalUnit::MonthDayNano, IntervalMonthDayNano::new(1, 0, 0), - ); } + ); + } #[test] fn test_split_vec_min_alloc_split_off_branch() { @@ -1949,6 +1965,78 @@ mod tests { check_result(&actual_batch, &data_set.expected_batch); } + #[test] + fn test_intern_for_fixed_size_binary_group_values() { + // Two-column group by `(FixedSizeBinary(2), Int64)` exercising the + // vectorized intern path end-to-end (hashing included), with nulls, + // within-batch repeats and across-batch repeats. + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::FixedSizeBinary(2), true), + Field::new("b", DataType::Int64, true), + ])); + let mut group_values = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + + fn fsb(values: Vec>) -> ArrayRef { + Arc::new( + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + values.into_iter(), + 2, + ) + .unwrap(), + ) + } + + let batch1: Vec = vec![ + fsb(vec![Some(b"aa"), Some(b"aa"), None, None, Some(b"bb")]), + Arc::new(Int64Array::from(vec![ + Some(1), + Some(1), + None, + Some(2), + None, + ])), + ]; + // Mix of groups repeated from batch1 and new groups + let batch2: Vec = vec![ + fsb(vec![Some(b"aa"), Some(b"cc"), None, Some(b"bb")]), + Arc::new(Int64Array::from(vec![Some(1), Some(1), None, Some(3)])), + ]; + + group_values.intern(&batch1, &mut vec![]).unwrap(); + group_values.intern(&batch2, &mut vec![]).unwrap(); + + let actual_batch = group_values.emit(EmitTo::All).unwrap(); + let actual_batch = + RecordBatch::try_new(Arc::clone(&schema), actual_batch).unwrap(); + + let expected_batch = RecordBatch::try_new( + schema, + vec![ + fsb(vec![ + Some(b"aa"), + None, + None, + Some(b"bb"), + Some(b"cc"), + Some(b"bb"), + ]), + Arc::new(Int64Array::from(vec![ + Some(1), + None, + Some(2), + None, + Some(1), + Some(3), + ])), + ], + ) + .unwrap(); + + assert_eq!(actual_batch.num_rows(), expected_batch.num_rows()); + check_result(&actual_batch, &expected_batch); + } + #[test] fn test_emit_first_n_for_vectorized_group_values() { let data_set = VectorizedTestDataSet::new(); diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 25b69d16dd035..3e6c16e12595f 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -5478,6 +5478,31 @@ SELECT id, MAX(value) FROM fixed_size_binary_views GROUP BY id ORDER BY id; 3 000101 4 NULL +# Group by a FixedSizeBinary column +# (exercises the FixedSizeBinary `GroupColumn` in `GroupValuesColumn`) +query ?I +SELECT value, COUNT(*) FROM fixed_size_binary_views GROUP BY value ORDER BY value; +---- +000101 2 +000102 1 +000103 3 +000104 2 +000109 1 +NULL 2 + +# Multi-column group by including a FixedSizeBinary column +query ?II +SELECT value, id, COUNT(*) FROM fixed_size_binary_views GROUP BY value, id ORDER BY value, id; +---- +000101 2 1 +000101 3 1 +000102 1 1 +000103 1 3 +000104 1 2 +000109 2 1 +NULL 1 1 +NULL 4 1 + statement ok DROP VIEW fixed_size_binary_views; From 2c00fafdcb0afb1900bb25299b1e279004068229 Mon Sep 17 00:00:00 2001 From: RIchard Baah Date: Mon, 17 Aug 2026 14:47:49 -0400 Subject: [PATCH 4/4] [cherry-pick] feat: Support IEEE 754 negative zero semantics (#22835) Cherry-pick of upstream commit 7dd1c6a2c68072eb7cacd3c56adf53a21a448863. --- datafusion/common/src/hash_utils.rs | 10 +- datafusion/common/src/utils/mod.rs | 89 +++++++ datafusion/functions-nested/src/except.rs | 15 +- datafusion/functions-nested/src/set_ops.rs | 36 ++- datafusion/physical-expr-common/src/datum.rs | 18 +- .../group_values/multi_group_by/mod.rs | 52 +++- .../group_values/multi_group_by/primitive.rs | 25 +- .../group_values/multi_group_by/row_backed.rs | 9 +- .../src/aggregates/group_values/row.rs | 8 + .../group_values/single_group_by/primitive.rs | 28 ++- datafusion/physical-plan/src/joins/utils.rs | 33 ++- .../test_files/array/array_distinct.slt | 42 ++++ .../test_files/array/array_except.slt | 48 ++++ .../test_files/array/array_union.slt | 60 +++++ .../sqllogictest/test_files/negative_zero.slt | 231 ++++++++++++++++++ 15 files changed, 660 insertions(+), 44 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/negative_zero.slt diff --git a/datafusion/common/src/hash_utils.rs b/datafusion/common/src/hash_utils.rs index fcc2e919b6cc2..02db75498af49 100644 --- a/datafusion/common/src/hash_utils.rs +++ b/datafusion/common/src/hash_utils.rs @@ -188,10 +188,16 @@ macro_rules! hash_float_value { ($(($t:ty, $i:ty)),+) => { $(impl HashValue for $t { fn hash_one(&self, state: &RandomState) -> u64 { - state.hash_one(<$i>::from_ne_bytes(self.to_ne_bytes())) + // +0.0 and -0.0 differ only in the sign bit but compare equal + // under IEEE 754; normalize -0.0 → +0.0 so Hash agrees with Eq. + let bits = <$i>::from_ne_bytes(self.to_ne_bytes()); + let bits = if bits << 1 == 0 { 0 } else { bits }; + state.hash_one(bits) } fn hash_write(&self, hasher: &mut impl Hasher) { - hasher.write(&self.to_ne_bytes()) + let bits = <$i>::from_ne_bytes(self.to_ne_bytes()); + let bits: $i = if bits << 1 == 0 { 0 } else { bits }; + hasher.write(&bits.to_ne_bytes()) } })+ }; diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index 9e7f6cdf4ebe1..99205e411d67c 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -1255,6 +1255,95 @@ fn fsl_values_row_number(list_size: i32, array_len: usize) -> Result Ok(PrimitiveArray::new(rows_number.into(), None)) } +/// Replace `-0.0` with `+0.0` in any `Float16`, `Float32`, or `Float64` array. +/// For non-float arrays returns the input unchanged. NaN payloads are +/// preserved. +/// +/// Arrow's comparison kernels (`arrow::compute::kernels::cmp::eq` etc.) and +/// row-encoding (`arrow::row::RowConverter`) use IEEE 754 totalOrder +/// semantics, which treats `-0.0` and `+0.0` as distinct. SQL semantics +/// (PostgreSQL / IEEE 754 equality) require them to compare equal, so +/// callers normalize before invoking those kernels. +/// +/// The common case - no `-0.0` present - is allocation-free: a single +/// read-only scan of the underlying buffer (auto-vectorizable to an +/// OR-reduction) decides whether to fall through to the rewriting path. +/// Only arrays that actually contain `-0.0` pay for a new buffer. +pub fn normalize_float_zero(array: &ArrayRef) -> ArrayRef { + use arrow::array::{Float16Array, Float32Array, Float64Array}; + use arrow::datatypes::{Float16Type, Float32Type, Float64Type}; + // -0.0 has only the sign bit set; no other finite or NaN value shares + // this bit pattern, so a strict-equality scan reliably gates the rewrite. + const NEG_ZERO_F16_BITS: u16 = half::f16::NEG_ZERO.to_bits(); + const NEG_ZERO_F32_BITS: u32 = (-0.0_f32).to_bits(); + const NEG_ZERO_F64_BITS: u64 = (-0.0_f64).to_bits(); + match array.data_type() { + DataType::Float32 => { + let arr: &Float32Array = array.as_primitive::(); + if !arr + .values() + .iter() + .any(|v| v.to_bits() == NEG_ZERO_F32_BITS) + { + return Arc::clone(array); + } + let normalized: Float32Array = + arr.unary(|v| if v.to_bits() << 1 == 0 { 0.0_f32 } else { v }); + Arc::new(normalized) + } + DataType::Float64 => { + let arr: &Float64Array = array.as_primitive::(); + if !arr + .values() + .iter() + .any(|v| v.to_bits() == NEG_ZERO_F64_BITS) + { + return Arc::clone(array); + } + let normalized: Float64Array = + arr.unary(|v| if v.to_bits() << 1 == 0 { 0.0_f64 } else { v }); + Arc::new(normalized) + } + DataType::Float16 => { + let arr: &Float16Array = array.as_primitive::(); + if !arr + .values() + .iter() + .any(|v| v.to_bits() == NEG_ZERO_F16_BITS) + { + return Arc::clone(array); + } + let normalized: Float16Array = arr.unary(|v| { + if v.to_bits() << 1 == 0 { + half::f16::from_bits(0) + } else { + v + } + }); + Arc::new(normalized) + } + _ => Arc::clone(array), + } +} + +/// Replace `-0.0` with `+0.0` in `Float16`, `Float32`, or `Float64` scalar +/// values. Other variants are returned unchanged. See [`normalize_float_zero`] +/// for context. +pub fn normalize_float_zero_scalar(scalar: ScalarValue) -> ScalarValue { + match scalar { + ScalarValue::Float32(Some(v)) if v.to_bits() << 1 == 0 => { + ScalarValue::Float32(Some(0.0)) + } + ScalarValue::Float64(Some(v)) if v.to_bits() << 1 == 0 => { + ScalarValue::Float64(Some(0.0)) + } + ScalarValue::Float16(Some(v)) if v.to_bits() << 1 == 0 => { + ScalarValue::Float16(Some(half::f16::from_bits(0))) + } + other => other, + } +} + #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/datafusion/functions-nested/src/except.rs b/datafusion/functions-nested/src/except.rs index 12ed6c2e186f4..dbf815c0ec539 100644 --- a/datafusion/functions-nested/src/except.rs +++ b/datafusion/functions-nested/src/except.rs @@ -27,7 +27,7 @@ use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::compute::take; use arrow::datatypes::{DataType, FieldRef}; use arrow::row::{RowConverter, SortField}; -use datafusion_common::utils::{ListCoercion, take_function_args}; +use datafusion_common::utils::{ListCoercion, normalize_float_zero, take_function_args}; use datafusion_common::{HashSet, Result, internal_err}; use datafusion_expr::{ ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, @@ -169,16 +169,21 @@ fn general_except( ) -> Result> { let converter = RowConverter::new(vec![SortField::new(l.value_type())])?; + // Normalize -0.0 → +0.0 so RowConverter (IEEE 754 totalOrder) groups + // ±0 together for both the rhs lookup set and the lhs probe. + let l_values_norm = normalize_float_zero(l.values()); + let r_values_norm = normalize_float_zero(r.values()); + // Only convert the visible portion of the values array. For sliced // ListArrays, values() returns the full underlying array but only // elements between the first and last offset are referenced. let l_first = l.offsets()[0].as_usize(); let l_len = l.offsets()[l.len()].as_usize() - l_first; - let l_values = converter.convert_columns(&[l.values().slice(l_first, l_len)])?; + let l_values = converter.convert_columns(&[l_values_norm.slice(l_first, l_len)])?; let r_first = r.offsets()[0].as_usize(); let r_len = r.offsets()[r.len()].as_usize() - r_first; - let r_values = converter.convert_columns(&[r.values().slice(r_first, r_len)])?; + let r_values = converter.convert_columns(&[r_values_norm.slice(r_first, r_len)])?; let mut offsets = Vec::::with_capacity(l.len() + 1); offsets.push(OffsetSize::usize_as(0)); @@ -223,11 +228,11 @@ fn general_except( } else if OffsetSize::IS_LARGE { let indices = UInt64Array::from(indices.into_iter().map(|i| i as u64).collect::>()); - take(l.values().as_ref(), &indices, None)? + take(l_values_norm.as_ref(), &indices, None)? } else { let indices = UInt32Array::from(indices.into_iter().map(|i| i as u32).collect::>()); - take(l.values().as_ref(), &indices, None)? + take(l_values_norm.as_ref(), &indices, None)? }; Ok(GenericListArray::::new( diff --git a/datafusion/functions-nested/src/set_ops.rs b/datafusion/functions-nested/src/set_ops.rs index 2ad08e2d43c02..2214d3d35bb7b 100644 --- a/datafusion/functions-nested/src/set_ops.rs +++ b/datafusion/functions-nested/src/set_ops.rs @@ -28,7 +28,7 @@ use arrow::datatypes::DataType::{LargeList, List, Null}; use arrow::datatypes::{DataType, Field, FieldRef}; use arrow::row::{RowConverter, SortField}; use datafusion_common::cast::{as_large_list_array, as_list_array}; -use datafusion_common::utils::ListCoercion; +use datafusion_common::utils::{ListCoercion, normalize_float_zero}; use datafusion_common::{ Result, assert_eq_or_internal_err, exec_err, internal_err, utils::take_function_args, }; @@ -351,21 +351,28 @@ fn generic_set_lists( let converter = RowConverter::new(vec![SortField::new(l.value_type())])?; + // Normalize -0.0 → +0.0 so RowConverter (which uses IEEE 754 totalOrder + // and treats ±0 as distinct) groups them together. Use the normalized + // arrays for both row conversion and the final output values. + let l_values_norm = normalize_float_zero(l.values()); + let r_values_norm = normalize_float_zero(r.values()); + // Only convert the visible portion of the values array. For sliced // ListArrays, values() returns the full underlying array but only // elements between the first and last offset are referenced. let l_first = l.offsets()[0].as_usize(); let l_len = l.offsets()[l.len()].as_usize() - l_first; - let rows_l = converter.convert_columns(&[l.values().slice(l_first, l_len)])?; + let l_values = l_values_norm.slice(l_first, l_len); + let rows_l = converter.convert_columns(&[Arc::clone(&l_values)])?; let r_first = r.offsets()[0].as_usize(); let r_len = r.offsets()[r.len()].as_usize() - r_first; - let rows_r = converter.convert_columns(&[r.values().slice(r_first, r_len)])?; + let r_values = r_values_norm.slice(r_first, r_len); + let rows_r = converter.convert_columns(&[Arc::clone(&r_values)])?; - // Combine the *sliced* value arrays so 0-based indices from the row - // converter map directly into the concatenated array. - let l_values = l.values().slice(l_first, l_len); - let r_values = r.values().slice(r_first, r_len); + // Indices from the row converter are 0-based in the per-side slice; + // concatenating those same slices lets indices map directly into the + // combined values array. let combined_values = concat(&[l_values.as_ref(), r_values.as_ref()])?; let r_offset = l_len; @@ -558,13 +565,18 @@ fn general_array_distinct( let converter = RowConverter::new(vec![SortField::new(dt.clone())])?; + // Normalize -0.0 → +0.0 so RowConverter (which uses IEEE 754 totalOrder + // and treats ±0 as distinct) groups them together, and so the output + // carries the canonical sign. + let values_norm = normalize_float_zero(array.values()); + // Only convert the visible portion of the values array. For sliced // ListArrays, values() returns the full underlying array but only // elements between the first and last offset are referenced. let first_offset = value_offsets[0].as_usize(); let visible_len = value_offsets[array.len()].as_usize() - first_offset; let rows = - converter.convert_columns(&[array.values().slice(first_offset, visible_len)])?; + converter.convert_columns(&[values_norm.slice(first_offset, visible_len)])?; let mut indices: Vec = Vec::with_capacity(rows.num_rows()); let mut seen = HashSet::new(); @@ -593,19 +605,19 @@ fn general_array_distinct( } // Gather distinct values in a single pass, using the computed `indices`. - // Indices are absolute positions in array.values() (first_offset was added - // back when collecting them), so we can take directly from the full values. + // Indices are absolute positions in the (normalized) values array, so we + // can take directly from the full values. // Use UInt64Array for LargeList to support values arrays exceeding u32::MAX. let final_values = if indices.is_empty() { new_empty_array(&dt) } else if OffsetSize::IS_LARGE { let indices = UInt64Array::from(indices.into_iter().map(|i| i as u64).collect::>()); - take(array.values().as_ref(), &indices, None)? + take(values_norm.as_ref(), &indices, None)? } else { let indices = UInt32Array::from(indices.into_iter().map(|i| i as u32).collect::>()); - take(array.values().as_ref(), &indices, None)? + take(values_norm.as_ref(), &indices, None)? }; Ok(Arc::new(GenericListArray::::try_new( diff --git a/datafusion/physical-expr-common/src/datum.rs b/datafusion/physical-expr-common/src/datum.rs index bd5790507f662..d23fb30db6c4a 100644 --- a/datafusion/physical-expr-common/src/datum.rs +++ b/datafusion/physical-expr-common/src/datum.rs @@ -23,6 +23,7 @@ use arrow::compute::kernels::cmp::{ }; use arrow::compute::{SortOptions, ilike, like, nilike, nlike}; use arrow::error::ArrowError; +use datafusion_common::utils::{normalize_float_zero, normalize_float_zero_scalar}; use datafusion_common::{Result, ScalarValue}; use datafusion_common::{arrow_datafusion_err, assert_or_internal_err, internal_err}; use datafusion_expr_common::columnar_value::ColumnarValue; @@ -84,7 +85,22 @@ pub fn apply_cmp( } }; - apply(lhs, rhs, |l, r| Ok(Arc::new(f(l, r)?))) + // Arrow's comparison kernels use IEEE 754 totalOrder semantics for + // floats, which treats `-0.0` and `+0.0` as distinct. Normalize float + // operands so SQL semantics (`+0.0 == -0.0`) hold. No-op for + // non-float types. + let lhs = normalize_cmp_input(lhs); + let rhs = normalize_cmp_input(rhs); + apply(&lhs, &rhs, |l, r| Ok(Arc::new(f(l, r)?))) + } +} + +fn normalize_cmp_input(cv: &ColumnarValue) -> ColumnarValue { + match cv { + ColumnarValue::Array(a) => ColumnarValue::Array(normalize_float_zero(a)), + ColumnarValue::Scalar(s) => { + ColumnarValue::Scalar(normalize_float_zero_scalar(s.clone())) + } } } 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 d7e374f0bccb9..e16f461e38f2a 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,10 +36,13 @@ use crate::aggregates::group_values::multi_group_by::{ }; use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; 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, + 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, TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; @@ -959,9 +962,11 @@ fn group_column_supported_type(data_type: &DataType) -> bool { | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 + | DataType::Float16 | DataType::Float32 | DataType::Float64 | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) | DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary @@ -983,6 +988,8 @@ fn group_column_supported_type(data_type: &DataType) -> bool { | DataType::Time64(TimeUnit::Microsecond) | DataType::Time64(TimeUnit::Nanosecond) | DataType::Timestamp(_, _) + | DataType::Duration(_) + | DataType::Interval(_) | DataType::Utf8View | DataType::BinaryView | DataType::Boolean @@ -1017,6 +1024,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) } @@ -1065,6 +1075,34 @@ 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::Duration(t) => match t { + TimeUnit::Second => { + instantiate_primitive!(v, nullable, DurationSecondType, data_type) + } + TimeUnit::Millisecond => { + instantiate_primitive!(v, nullable, DurationMillisecondType, data_type) + } + TimeUnit::Microsecond => { + instantiate_primitive!(v, nullable, DurationMicrosecondType, data_type) + } + TimeUnit::Nanosecond => { + instantiate_primitive!(v, nullable, DurationNanosecondType, data_type) + } + }, + DataType::Interval(u) => match u { + IntervalUnit::YearMonth => { + instantiate_primitive!(v, nullable, IntervalYearMonthType, data_type) + } + IntervalUnit::DayTime => { + instantiate_primitive!(v, nullable, IntervalDayTimeType, data_type) + } + IntervalUnit::MonthDayNano => { + instantiate_primitive!(v, nullable, IntervalMonthDayNanoType, data_type) + } + }, DataType::Utf8 => { v.push(Box::new(ByteGroupValueBuilder::::new( OutputType::Utf8, @@ -2090,8 +2128,10 @@ mod tests { // 6. Only decrease group indices in non-inlined group index view // 7. Erase all things - let field = Field::new_list_field(DataType::Int32, true); - let schema = Arc::new(Schema::new_with_metadata(vec![field], HashMap::new())); + let schema = Arc::new(Schema::new_with_metadata( + vec![] as Vec, + HashMap::new(), + )); let mut group_values = GroupValuesColumn::::try_new(schema).unwrap(); // Insert group index views and check if success to insert diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs index 4aae996f6811d..5a54b87c0015f 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::aggregates::group_values::HashValue; use crate::aggregates::group_values::multi_group_by::{ GroupColumn, Nulls, nulls_equal_to, split_vec_min_alloc, }; @@ -50,6 +51,7 @@ pub struct PrimitiveGroupValueBuilder PrimitiveGroupValueBuilder where T: ArrowPrimitiveType, + T::Native: HashValue, { /// Create a new `PrimitiveGroupValueBuilder` pub fn new(data_type: DataType) -> Self { @@ -90,7 +92,9 @@ where } else { unsafe { *array_values.get_unchecked(rhs_row) } }; - if left.is_eq(right) { + // `left` was already canonicalized on append; canonicalize the + // input so ±0 (and any future equivalence class) compares equal. + if left.is_eq(right.canonicalize()) { cmp_buf[i / 8] |= 1 << (i % 8); } } @@ -132,7 +136,7 @@ where continue; } - if !self.group_values[lhs_row].is_eq(array.value(rhs_row)) { + if !self.group_values[lhs_row].is_eq(array.value(rhs_row).canonicalize()) { equal_to_results.set_bit(idx, false); } } @@ -141,6 +145,8 @@ where impl GroupColumn for PrimitiveGroupValueBuilder +where + T::Native: HashValue, { fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { // Perf: skip null check (by short circuit) if input is not nullable @@ -153,7 +159,8 @@ impl GroupColumn // Otherwise, we need to check their values } - self.group_values[lhs_row].is_eq(array.as_primitive::().value(rhs_row)) + self.group_values[lhs_row] + .is_eq(array.as_primitive::().value(rhs_row).canonicalize()) } fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { @@ -164,10 +171,12 @@ impl GroupColumn self.group_values.push(T::default_value()); } else { self.nulls.append(false); - self.group_values.push(array.as_primitive::().value(row)); + self.group_values + .push(array.as_primitive::().value(row).canonicalize()); } } else { - self.group_values.push(array.as_primitive::().value(row)); + self.group_values + .push(array.as_primitive::().value(row).canonicalize()); } Ok(()) @@ -213,7 +222,7 @@ impl GroupColumn self.group_values.push(T::default_value()); } else { self.nulls.append(false); - self.group_values.push(arr.value(row)); + self.group_values.push(arr.value(row).canonicalize()); } } } @@ -221,7 +230,7 @@ impl GroupColumn (true, Nulls::None) => { self.nulls.append_n(rows.len(), false); for &row in rows { - self.group_values.push(arr.value(row)); + self.group_values.push(arr.value(row).canonicalize()); } } @@ -233,7 +242,7 @@ impl GroupColumn (false, _) => { for &row in rows { - self.group_values.push(arr.value(row)); + self.group_values.push(arr.value(row).canonicalize()); } } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs index 29beb3bd66229..085510d998976 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs @@ -28,7 +28,7 @@ //! //! [`GroupValuesColumn`] can only be used when *every* column of the group-by //! key has a [`GroupColumn`] implementation; otherwise the whole aggregation -//! falls back to the row-wise [`GroupValuesRows`], which is materially slower +//! falls back to the row-wise `GroupValuesRows`, which is materially slower //! and heavier for the columns that *would* have qualified for the column-wise //! fast path. By providing a generic fallback `GroupColumn`, a schema like //! `GROUP BY int_col, struct_col` keeps `int_col` on its fast native builder @@ -41,10 +41,9 @@ //! raw input columns via `create_hashes`, which already supports nested types. //! Equality is decided here by comparing arrow-row bytes. For the two to agree //! on group identity, values that this column considers equal must hash equal — -//! see the float `-0.0` / `NaN` note on [`RowsGroupColumn`]. +//! see the float `-0.0` / `NaN` note on [`RowsGroupColumn`] below. //! //! [`GroupValuesColumn`]: crate::aggregates::group_values::multi_group_by::GroupValuesColumn -//! [`GroupValuesRows`]: crate::aggregates::group_values::GroupValuesRows use crate::aggregates::group_values::multi_group_by::GroupColumn; use crate::aggregates::group_values::row::encode_array_if_necessary; @@ -70,11 +69,9 @@ use datafusion_common::{DataFusionError, Result}; /// encoding, which treats `-0.0` and `+0.0` as distinct and canonicalizes /// `NaN`. Because hashing is performed separately (on the raw input array), a /// caller must ensure the two agree — e.g. by normalizing `-0.0 → +0.0` on the -/// input columns before hashing when a float leaf is present (as -/// [`GroupValuesRows`] does). See the module docs. +/// input columns before hashing when a float leaf is present. See the module docs. /// /// [row format]: arrow::row -/// [`GroupValuesRows`]: crate::aggregates::group_values::GroupValuesRows pub struct RowsGroupColumn { /// Single-field row converter for this column's data type. row_converter: RowConverter, diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 5fefe9df3e849..cbd7a609c5caa 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -27,6 +27,7 @@ use arrow::row::{RowConverter, Rows, SortField}; use datafusion_common::Result; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; +use datafusion_common::utils::normalize_float_zero; use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; use datafusion_expr::EmitTo; use hashbrown::hash_table::HashTable; @@ -117,6 +118,13 @@ impl GroupValuesRows { impl GroupValues for GroupValuesRows { fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { + // Normalize -0.0 → +0.0 so RowConverter (IEEE 754 totalOrder) and + // primitive hashing both group ±0 together. No-op for non-float + // columns. + let normalized_cols: Vec = + cols.iter().map(normalize_float_zero).collect(); + let cols = normalized_cols.as_slice(); + // Convert the group keys into the row format let group_rows = &mut self.rows_buffer; group_rows.clear(); diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index efaf7eba0f1b5..73b7d918e753a 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs @@ -34,8 +34,21 @@ use std::mem::size_of; use std::sync::Arc; /// A trait to allow hashing of floating point numbers -pub(crate) trait HashValue { +pub trait HashValue { fn hash(&self, state: &RandomState) -> u64; + + /// Return a canonical representative whose bit pattern is identical for + /// all values that should be grouped together. Default is the identity; + /// floats override this to fold `-0.0` into `+0.0` so the bit-equal + /// `is_eq` check used during insertion treats them as the same group. + /// NaN payload bits are preserved. + #[inline] + fn canonicalize(self) -> Self + where + Self: Sized, + { + self + } } macro_rules! hash_integer { @@ -62,13 +75,20 @@ macro_rules! hash_float { $(impl HashValue for $t { #[cfg(not(feature = "force_hash_collisions"))] fn hash(&self, state: &RandomState) -> u64 { - state.hash_one(self.to_bits()) + state.hash_one(self.canonicalize().to_bits()) } #[cfg(feature = "force_hash_collisions")] fn hash(&self, _state: &RandomState) -> u64 { 0 } + + #[inline] + fn canonicalize(self) -> Self { + let bits = self.to_bits(); + let bits = if bits << 1 == 0 { 0 } else { bits }; + Self::from_bits(bits) + } })+ }; } @@ -126,6 +146,10 @@ where group_id }), Some(key) => { + // Fold equivalence-class duplicates (e.g. `-0.0` → `+0.0`) + // so the bit-equal `is_eq` matches and the stored value is + // the canonical representative. + let key = key.canonicalize(); let state = &self.random_state; let hash = key.hash(state); let insert = self.map.entry( diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index b4aa295562b67..9e0d7ae5580af 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -43,7 +43,7 @@ pub use crate::joins::{JoinOn, JoinOnRef}; use arrow::array::{ Array, ArrowPrimitiveType, BooleanBufferBuilder, NativeAdapter, PrimitiveArray, RecordBatch, RecordBatchOptions, UInt32Array, UInt32Builder, UInt64Array, - builder::UInt64Builder, downcast_array, new_null_array, + builder::UInt64Builder, downcast_array, make_array, new_null_array, }; use arrow::array::{ ArrayRef, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, Date64Array, @@ -65,6 +65,7 @@ use datafusion_common::cast::as_boolean_array; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; use datafusion_common::stats::Precision; +use datafusion_common::utils::normalize_float_zero; use datafusion_common::{ DataFusionError, JoinSide, JoinType, NullEquality, Result, SharedResult, not_impl_err, plan_err, @@ -1931,6 +1932,25 @@ fn eq_dyn_null( }; return Ok(compare_op_for_nested(op, &left, &right)?); } + // Arrow's `eq` / `not_distinct` use IEEE 754 totalOrder semantics for + // floats, so `-0.0` and `+0.0` would compare unequal. Normalize float + // operands first; non-float types dispatch directly to avoid the + // `make_array(to_data())` round-trip. + if !matches!( + left.data_type(), + DataType::Float16 | DataType::Float32 | DataType::Float64 + ) { + return match null_equality { + NullEquality::NullEqualsNothing => eq(&left, &right), + NullEquality::NullEqualsNull => not_distinct(&left, &right), + }; + } + let left_arr: ArrayRef = make_array(left.to_data()); + let right_arr: ArrayRef = make_array(right.to_data()); + let left_norm = normalize_float_zero(&left_arr); + let right_norm = normalize_float_zero(&right_arr); + let left = left_norm.as_ref(); + let right = right_norm.as_ref(); match null_equality { NullEquality::NullEqualsNothing => eq(&left, &right), NullEquality::NullEqualsNull => not_distinct(&left, &right), @@ -1979,7 +1999,16 @@ impl JoinKeyComparator { .zip(right_arrays.iter()) .zip(sort_options.iter()) .map(|((l, r), opts)| { - let inner = make_comparator(l.as_ref(), r.as_ref(), *opts)?; + // `make_comparator` uses IEEE 754 totalOrder for floats and + // treats `-0.0` / `+0.0` as distinct. Normalize float arrays + // so SMJ / piecewise-merge equi-keys honor SQL equality; + // no-op (Arc::clone) for non-floats and for float arrays + // that contain no `-0.0`. `normalize_float_zero` preserves + // null positions, so the original null masks below remain + // valid. + let l_norm = normalize_float_zero(l); + let r_norm = normalize_float_zero(r); + let inner = make_comparator(l_norm.as_ref(), r_norm.as_ref(), *opts)?; if null_equality == NullEquality::NullEqualsNothing { let ln = l.logical_nulls().filter(|n| n.null_count() > 0); let rn = r.logical_nulls().filter(|n| n.null_count() > 0); diff --git a/datafusion/sqllogictest/test_files/array/array_distinct.slt b/datafusion/sqllogictest/test_files/array/array_distinct.slt index 88ffdf7f2ff78..7b7033139d767 100644 --- a/datafusion/sqllogictest/test_files/array/array_distinct.slt +++ b/datafusion/sqllogictest/test_files/array/array_distinct.slt @@ -210,5 +210,47 @@ select array_compact(arrow_cast(make_array(NULL, NULL, NULL), 'FixedSizeList(3, ---- [] +# Negative zero (-0.0) and positive zero (+0.0) compare equal under +# IEEE 754, but their bit patterns differ. array_distinct must normalize +# the sign so the canonical representative (+0.0) is used; otherwise +# group-by / dedup hashing on the raw bits keeps both as distinct +# elements. PostgreSQL / IEEE 754 expected output below. + +# array_distinct collapses +0.0 and -0.0 into a single element. +query ? +select array_distinct([0.0, -0.0]); +---- +[0.0] + +# General case with extra elements. +query ? +select array_distinct([0.0, -0.0, 0.0, 1.0, -0.0]); +---- +[0.0, 1.0] + +# array_length(array_distinct(...)) for {+0.0, -0.0, +0.0} must be 1. +query I +select array_length(array_distinct([0.0, -0.0, 0.0])); +---- +1 + +# Float32 list. +query ? +select array_distinct(arrow_cast([0.0, -0.0], 'List(Float32)')); +---- +[0.0] + +# LargeList(Float64). +query ? +select array_distinct(arrow_cast([0.0, -0.0], 'LargeList(Float64)')); +---- +[0.0] + +# FixedSizeList(Float64). +query ? +select array_distinct(arrow_cast([0.0, -0.0], 'FixedSizeList(2, Float64)')); +---- +[0.0] + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/array/array_except.slt b/datafusion/sqllogictest/test_files/array/array_except.slt index a718723e58c38..1d41a5a79d15f 100644 --- a/datafusion/sqllogictest/test_files/array/array_except.slt +++ b/datafusion/sqllogictest/test_files/array/array_except.slt @@ -156,4 +156,52 @@ select array_except(arrow_cast([1, 2, 3, 4], 'FixedSizeList(4, Int64)'), arrow_c [1, 2] +# Negative zero (-0.0) and positive zero (+0.0) compare equal under +# IEEE 754, but their bit patterns differ. array_except must treat +# +0.0 and -0.0 as the same element when subtracting. PostgreSQL / +# IEEE 754 expected output below. + +# -0.0 in rhs removes +0.0 in lhs. +query ? +select array_except([0.0], [-0.0]); +---- +[] + +# Reverse direction. +query ? +select array_except([-0.0], [0.0]); +---- +[] + +# -0.0 in rhs also removes the +0.0 element from the lhs. +query ? +select array_except([0.0, -0.0], [-0.0]); +---- +[] + +# +0.0 in rhs also removes the -0.0 element from the lhs. +query ? +select array_except([0.0, -0.0], [0.0]); +---- +[] + +# More general case with extra unmatched element. +query ? +select array_except([0.0, -0.0, 1.0], [-0.0]); +---- +[1.0] + +# Float32 list. +query ? +select array_except(arrow_cast([0.0, -0.0], 'List(Float32)'), arrow_cast([0.0], 'List(Float32)')); +---- +[] + +# LargeList(Float64). +query ? +select array_except(arrow_cast([0.0, -0.0], 'LargeList(Float64)'), arrow_cast([-0.0], 'LargeList(Float64)')); +---- +[] + + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/array/array_union.slt b/datafusion/sqllogictest/test_files/array/array_union.slt index 6a0fdc546e7d7..edb90705af940 100644 --- a/datafusion/sqllogictest/test_files/array/array_union.slt +++ b/datafusion/sqllogictest/test_files/array/array_union.slt @@ -236,4 +236,64 @@ select array_except([1, 2], arrow_cast(null, 'List(Int64)')); NULL +# Negative zero (-0.0) and positive zero (+0.0) compare equal under +# IEEE 754, but their bit patterns differ. array_union and array_intersect +# must normalize the sign for dedup / matching; the canonical +# representative is +0.0. PostgreSQL / IEEE 754 expected output below. + +# array_union with +0.0 / -0.0 +query ? +select array_union([0.0], [-0.0]); +---- +[0.0] + +query ? +select array_union([0.0, 1.0], [-0.0]); +---- +[0.0, 1.0] + +query ? +select array_union([0.0, -0.0, 1.0], [-0.0, 1.0]); +---- +[0.0, 1.0] + +# Float32 list. +query ? +select array_union(arrow_cast([0.0], 'List(Float32)'), arrow_cast([-0.0], 'List(Float32)')); +---- +[0.0] + +# LargeList(Float64). +query ? +select array_union(arrow_cast([0.0], 'LargeList(Float64)'), arrow_cast([-0.0], 'LargeList(Float64)')); +---- +[0.0] + + +# array_intersect with +0.0 / -0.0 +# +0.0 in lhs matches -0.0 in rhs. +query ? +select array_intersect([0.0, 1.0], [-0.0]); +---- +[0.0] + +# Either +0.0 or -0.0 in lhs matches +0.0 in rhs (canonicalized to +0.0). +query ? +select array_intersect([0.0, -0.0], [0.0]); +---- +[0.0] + +# Same with -0.0 in rhs. +query ? +select array_intersect([0.0, -0.0], [-0.0]); +---- +[0.0] + +# Float32 list. +query ? +select array_intersect(arrow_cast([0.0], 'List(Float32)'), arrow_cast([-0.0], 'List(Float32)')); +---- +[0.0] + + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/negative_zero.slt b/datafusion/sqllogictest/test_files/negative_zero.slt new file mode 100644 index 0000000000000..8ea1122880e14 --- /dev/null +++ b/datafusion/sqllogictest/test_files/negative_zero.slt @@ -0,0 +1,231 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +########## +## Negative Zero (-0.0) vs. Positive Zero (+0.0) Behavior +########## +# +# IEEE 754 specifies +0.0 == -0.0 (they compare equal). PostgreSQL follows +# this and treats them as the same value for DISTINCT, GROUP BY, UNION, +# INTERSECT, EXCEPT, and equality predicates. The bit patterns differ in +# the sign bit, so any code path that hashes / compares the raw bits (e.g. +# `f64::to_bits` or `f64::to_ne_bytes`) will treat them as distinct values +# and must be normalized before grouping / dedup. +# +# Note: the sqllogictest formatter renders both `-0.0` and `+0.0` as `0`, +# so the visible scalar values look identical in the expected output. The +# behavior is asserted via row counts and via auxiliary `1.0 / a` +# (`Infinity` vs `-Infinity`) columns that expose the sign. + +##### +## Equality and ordering predicates +##### + +# +0.0 == -0.0 is TRUE; +0.0 < -0.0 and +0.0 > -0.0 are both FALSE. +query BBB +SELECT 0.0 = -0.0 AS eq, 0.0 < -0.0 AS lt, 0.0 > -0.0 AS gt; +---- +true false false + +# 0.0 IS DISTINCT FROM -0.0 must be FALSE because the values are equal. +query B +SELECT 0.0 IS DISTINCT FROM -0.0 AS is_distinct; +---- +false + +##### +## SELECT DISTINCT with +0.0 / -0.0 (Float64) +##### + +# DISTINCT must collapse +0.0 and -0.0 into a single row. +query R rowsort +SELECT DISTINCT a +FROM (SELECT 0.0 AS a UNION ALL SELECT -0.0 UNION ALL SELECT 0.0); +---- +0 + +# Same query, with `1.0 / a` to expose the sign in the projection. The +# tuples `(+0.0, +Infinity)` and `(-0.0, -Infinity)` are not equal — the +# zero columns compare equal but `+Infinity != -Infinity` — so DISTINCT +# keeps both rows. PG returns the same two rows. +query RR rowsort +SELECT DISTINCT a, 1.0 / a AS inv +FROM (SELECT 0.0 AS a UNION ALL SELECT -0.0 UNION ALL SELECT 0.0); +---- +0 -Infinity +0 Infinity + +# COUNT(DISTINCT) over {+0.0, -0.0} must return 1. +query I +SELECT COUNT(DISTINCT a) +FROM (SELECT 0.0 AS a UNION ALL SELECT -0.0 UNION ALL SELECT 0.0); +---- +1 + +# GROUP BY must put +0.0 and -0.0 in the same group. +query RRI rowsort +SELECT a, 1.0 / a AS inv, COUNT(*) +FROM (SELECT 0.0 AS a UNION ALL SELECT -0.0 UNION ALL SELECT 0.0) +GROUP BY a; +---- +0 Infinity 3 + +# Multi-column DISTINCT; (+0.0, 1) and (-0.0, 1) must collapse. +query RI rowsort +SELECT DISTINCT a, b +FROM (SELECT 0.0 AS a, 1 AS b UNION ALL SELECT -0.0, 1 UNION ALL SELECT 0.0, 2); +---- +0 1 +0 2 + +##### +## SELECT DISTINCT with +0.0 / -0.0 (Float32 / REAL) +##### + +# DISTINCT for Float32: same collapse to a single row. +query R rowsort +SELECT DISTINCT a +FROM ( + SELECT arrow_cast(0.0, 'Float32') AS a + UNION ALL SELECT arrow_cast(-0.0, 'Float32') + UNION ALL SELECT arrow_cast(0.0, 'Float32') +); +---- +0 + +# COUNT(DISTINCT) for Float32: must be 1. +query I +SELECT COUNT(DISTINCT a) +FROM ( + SELECT arrow_cast(0.0, 'Float32') AS a + UNION ALL SELECT arrow_cast(-0.0, 'Float32') +); +---- +1 + +##### +## UNION (set semantics) with +0.0 / -0.0 +##### + +# UNION (DISTINCT) must collapse +0.0 / -0.0 into a single row. +query R rowsort +SELECT 0.0 AS a UNION SELECT -0.0 UNION SELECT 0.0; +---- +0 + +# UNION ALL preserves every input row regardless of sign — baseline. +query R rowsort +SELECT 0.0 AS a UNION ALL SELECT -0.0 UNION ALL SELECT 0.0; +---- +0 +0 +0 + +# UNION on Float32 must also collapse to a single row. +query R rowsort +SELECT arrow_cast(0.0, 'Float32') AS a +UNION +SELECT arrow_cast(-0.0, 'Float32'); +---- +0 + +##### +## INTERSECT with +0.0 / -0.0 +##### + +# INTERSECT treats +0.0 and -0.0 as equal — one matching row. +query R rowsort +SELECT 0.0 AS a INTERSECT SELECT -0.0; +---- +0 + +# INTERSECT ALL with multiplicities min(1,1) = 1. +query R rowsort +SELECT 0.0 AS a INTERSECT ALL SELECT -0.0; +---- +0 + +# INTERSECT for Float32: same matching behavior. +query R rowsort +SELECT arrow_cast(0.0, 'Float32') AS a +INTERSECT +SELECT arrow_cast(-0.0, 'Float32'); +---- +0 + +##### +## EXCEPT with +0.0 / -0.0 +##### + +# EXCEPT treats +0.0 and -0.0 as equal — zero rows after subtraction. +query R rowsort +SELECT 0.0 AS a EXCEPT SELECT -0.0; +---- + +# Reverse direction: also zero rows. +query R rowsort +SELECT -0.0 AS a EXCEPT SELECT 0.0; +---- + +# EXCEPT for Float32: zero rows. +query R rowsort +SELECT arrow_cast(0.0, 'Float32') AS a +EXCEPT +SELECT arrow_cast(-0.0, 'Float32'); +---- + +# EXCEPT ALL with matching multiplicities: zero rows. +query R rowsort +SELECT 0.0 AS a EXCEPT ALL SELECT -0.0; +---- + +##### +## INNER JOIN ON equality with +0.0 / -0.0 +##### + +# Equi-join on a = b matches +0.0 against -0.0. +query RR +SELECT t1.a, t2.b +FROM (SELECT 0.0 AS a) t1 +JOIN (SELECT -0.0 AS b) t2 ON t1.a = t2.b; +---- +0 0 + +# Sort-merge join must also match +0.0 against -0.0. SMJ builds equi-key +# matchers via `JoinKeyComparator`, which calls Arrow's `make_comparator` +# (IEEE 754 totalOrder); without normalization, +0.0 and -0.0 produce +# different orderings and miss the match. +statement ok +set datafusion.optimizer.prefer_hash_join = false; + +query RR +SELECT t1.a, t2.b +FROM (SELECT 0.0 AS a) t1 +JOIN (SELECT -0.0 AS b) t2 ON t1.a = t2.b; +---- +0 0 + +# Float32 SMJ equi-join. +query RR +SELECT t1.a, t2.b +FROM (SELECT arrow_cast(0.0, 'Float32') AS a) t1 +JOIN (SELECT arrow_cast(-0.0, 'Float32') AS b) t2 ON t1.a = t2.b; +---- +0 0 + +statement ok +reset datafusion.optimizer.prefer_hash_join;