diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index f275d777c3279..edba17b05283d 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,6 +29,7 @@ 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; @@ -923,6 +925,15 @@ macro_rules! instantiate_primitive { /// 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 @@ -1067,6 +1078,14 @@ fn make_group_column(field: &Field) -> Result> { v.push(Box::new(BooleanGroupValueBuilder::::new())); } } + // 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!( @@ -1273,6 +1292,248 @@ mod tests { GroupIndexView, group_column_supported_type, make_group_column, 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 must be smaller than + // encoding every column into rows. + 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 + ); + assert!( + column_size < rows_size, + "expected column-wise path ({column_size}) to use less memory than \ + the all-rows fallback ({rows_size})" + ); + + // 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. + 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 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..6743c6bd3acd7 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs @@ -0,0 +1,373 @@ +// 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 the +//! entire key 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::dictionary_encode_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, +} + +impl RowsGroupColumn { + /// Returns whether `data_type` can be handled by this generic column, i.e. + /// whether arrow's [`RowConverter`] can encode it. + pub fn supports_type(data_type: &DataType) -> bool { + 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); + dictionary_encode_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); + } +} diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 4976a098ecee5..0d5e6aed1d6d4 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -267,7 +267,15 @@ impl GroupValues for GroupValuesRows { } } -fn dictionary_encode_if_necessary( +/// Re-apply dictionary / run-end encoding to `array` so it matches `expected`. +/// +/// Arrow's [`RowConverter`] decodes dictionary and run-end-encoded values to +/// their plain value type on the way out, so any group-value array produced +/// from the row format must be re-encoded to the schema's expected type before +/// being returned. Shared with the generic row-backed `GroupColumn`. +/// +/// [`RowConverter`]: arrow::row::RowConverter +pub(crate) fn dictionary_encode_if_necessary( array: &ArrayRef, expected: &DataType, ) -> Result {