zhuqi-lucas commented on code in PR #23523:
URL: https://github.com/apache/datafusion/pull/23523#discussion_r3579536147
##########
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})"
+ );
+
+ // 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<usize> {
+ 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()
+ }
Review Comment:
Good point — expanded the doc in 132905ebf to spell out that
`GroupValues::intern` deliberately does not fix the order of fresh group-ids,
and canonicalizing before comparison is what lets us assert equivalence between
the vectorized column path and the sequential rows fallback.
##########
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`.
Review Comment:
Applied in 132905ebf.
##########
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");
Review Comment:
Thanks for the pointer — will track those. They don't block this PR since
arrow's current API is enough for our use, but future improvements upstream
would likely simplify `encode_array_if_necessary`.
##########
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")
+ }
+
+ /// Encode a whole incoming column into the row format.
+ fn convert(&self, array: &ArrayRef) -> Result<Rows> {
+ 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<Self>) -> ArrayRef {
+ self.rows_to_array(&self.group_values)
+ }
+
+ fn take_n(&mut self, n: usize) -> ArrayRef {
+ debug_assert!(n <= self.group_values.num_rows());
Review Comment:
Yes. `aggregates/order/mod.rs:80` always emits `EmitTo::First(n.min(max))`,
and every other `GroupColumn::take_n` impl (`bytes.rs`, `bytes_view.rs`,
`boolean.rs`, `primitive.rs`) has the same `debug_assert!(self.len() >= n)`. So
the assert is defensive documentation of the caller-side invariant, matching
the existing convention.
--
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]
