zhuqi-lucas commented on code in PR #23523: URL: https://github.com/apache/datafusion/pull/23523#discussion_r3579534143
########## 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<Self> { + 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<Item = arrow::row::Row<'a>>, + ) -> 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") + } Review Comment: Fixed in 132905ebf — extended the recursion in `encode_array_if_necessary` to also cover `FixedSizeList`, `LargeList`, and `Map` (parallel to the existing `List` branch). Added `build_preserves_list_of_dictionary_schema` as a regression test. Side note: I discovered while writing tests that `RowConverter::convert_rows` currently rejects `FixedSizeList<Dict>` in arrow-rs itself (before our helper runs), so that specific nesting is unreachable through `supports_type` today. The new FSL branch is defensive coverage for parity with `List` in case that limitation lifts. ########## 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 Review Comment: You're right, thanks. Fixed the doc in 132905ebf to say 'flattens dict/REE values during row encoding (at `append`)' instead of 'on the way out'. ########## 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( Review Comment: Renamed to `encode_array_if_necessary` in 132905ebf. ########## 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`. Review Comment: Agreed — filing as a follow-up so this PR stays scoped to 'keep mixed nested schemas on the column-wise path'. A dedicated `DictGroupColumn` / `ReeGroupColumn` would skip the row-encode → cast round-trip entirely, which is a nice next step. ########## datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs: ########## @@ -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); + } Review Comment: `make_group_column` doesn't have a type-specialized `PrimitiveGroupValueBuilder` branch for them — they'd need `Float16Type` / `Decimal256Type` wired through the same pattern as `Float32Type` / `Int32Type` / etc. We could alternatively route them through `RowsGroupColumn` by relaxing the `is_nested()` gate, but I kept this PR scoped to nested types so the `group_column_supported_type ⇔ make_group_column` biconditional stays clean. Happy to file a follow-up either way if that's useful. ########## datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs: ########## @@ -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<Int64, 4> ("embedding"). + let fsl_field = Arc::new(Field::new("item", DataType::Int64, true)); + let mut fields: Vec<Field> = (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<ArrayRef> = (0..8) + .map(|c| { + let vals: Vec<i64> = + (0..n_groups).map(|r| (r as i64) * 8 + c as i64).collect(); + Arc::new(Int64Array::from(vals)) as ArrayRef + }) + .collect(); + let emb: Vec<Option<Vec<Option<i64>>>> = (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::<Int64Type, _, _>( + emb, 4, + )) as ArrayRef, + ); + + // Intern the same data into both implementations. + let mut column_path = GroupValuesColumn::<false>::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})" + ); Review Comment: Thanks! 🙂 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
