Rich-T-kid commented on code in PR #23523:
URL: https://github.com/apache/datafusion/pull/23523#discussion_r3573068220
##########
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:
> Arrow's [`RowConverter`] decodes dictionary and run-end-encoded values to
their plain value type on the way out
nit: Id say RowConverter decodes dictionarys/REE on their way in [since it
occurs
on](https://arrow.apache.org/rust/arrow_row/struct.RowConverter.html#method.append)`RowConverter::append()`.
##########
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:
the comment above mentions casting arrays to REE but the functions is called
`Dictionary_encode_if_necessary`.
```suggestion
pub(crate) fn encode_array_if_necessary(
```
or something of the sort
##########
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:
I think the cast here shows that Dictionaries and REE arrays should have
their own specialized implementation of `GroupColumn`. For now I think its fine
but unlike other implementations like fixedSizeList that can go from
`arrow-array -> RowFormat -> arrow-array`
dictionary/REE need an extra cast step
`arrow-array -> rowFormat -> arrow-array -> (cast) -> arrow-array`
##########
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:
why are `Float16` & `Decimal256` arrays not supported for this optimization?
##########
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:
this is a nice way to guarantee `groupColumns` will always consume less
memory than the `GroupValueRows` implementation!
##########
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:
nit:
```suggestion
//! and only pays the row-encoding cost on `struct_col`, instead of dragging
both
//! columns onto `GroupValuesRows`.
```
##########
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:
You may be interested in these issues/ PR in arrow-rs
https://github.com/apache/arrow-rs/issues/10275
https://github.com/apache/arrow-rs/pull/10319
##########
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:
does DataFusion guarantee that it wont call `GroupValues::emit(emit::to(n))`
where n is possibly greater than the number of groups?
##########
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:
This took me a bit of time to understand why its needed. from my
understanding its because `GroupValues::intern()` doesn't specify the order in
which new IDs are given.
[docs](https://docs.rs/datafusion/latest/datafusion/physical_plan/aggregates/group_values/trait.GroupValues.html#tymethod.intern)
I think the comment is fine as it is but maybe adding something like
> `GroupValues` implementations only guarantee that equal rows receive
/// equal group ids and new rows receive a fresh id — the order in which
/// new ids are handed out is not part of the contract, and can differ
/// between correct implementations
can help readers know why this function has to exist. This is mostly a nit
--
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]
