kosiew commented on code in PR #23523: URL: https://github.com/apache/datafusion/pull/23523#discussion_r3622161281
########## datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs: ########## @@ -0,0 +1,420 @@ +// 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, +} + +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())]) Review Comment: `supports_type` currently delegates to `RowConverter::supports_fields`, but that is not quite strong enough for the `GroupColumn` contract. The type also needs to successfully round trip through `convert_rows` during `build()` / `take_n()`. With current arrow-rs, `FixedSizeList<Dictionary<Int32, Utf8>>` returns `true` from `RowsGroupColumn::supports_type`, so `group_column_supported_type` and `make_group_column` select `RowsGroupColumn`. Appending succeeds, but `build()` later panics in `rows_to_array` before `encode_array_if_necessary` runs: ```text row conversion during emit: InvalidArgumentError("FixedSizeListArray expected data type Dictionary(Int32, Utf8) got Utf8 for \"item\"") ``` I reproduced this with a one-row `FixedSizeList<Dict<Int32, Utf8>, 2>`: `supports_type` returned `true`, `vectorized_append` succeeded, and `build()` panicked at `row_backed.rs:118`. So I don't think the Copilot thread is fully addressed for the nested container cases discussed there. The new recursion fixes `List<Dict>`, where Arrow currently round trips correctly, but `FixedSizeList`, `LargeList`, and `Map` with dictionary leaves can still be routed into `RowsGroupColumn` and panic instead of falling back to `GroupValuesRows`. Could we make the support predicate match the actual round trip capability of this column rather than only `RowConverter::supports_fields`? Until arrow-rs supports these cases, I think we should reject the known unsupported nested shapes so `supported_schema` returns `false` and the existing `GroupValuesRows` fallback is used instead. It would also be great to add a regression test showing either that `RowsGroupColumn::supports_type` rejects the currently unsupported shape, or that `build()` succeeds if you decide to handle it directly. -- 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]
