kosiew commented on code in PR #23523:
URL: https://github.com/apache/datafusion/pull/23523#discussion_r3612057072
##########
datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs:
##########
@@ -1273,6 +1292,255 @@ 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!(
Review Comment:
Nice test! One thought: this correctness test also hard asserts the current
relative `size()` estimates of `GroupValuesColumn` versus `GroupValuesRows`.
That provides useful evidence for the memory improvement, but it could become
brittle if Arrow row format or memory accounting changes without affecting
grouping correctness.
Would it make sense to keep the printed comparison together with the
`supported_schema` and output equivalence assertions here, and move the memory
delta check to a benchmark or relax it a bit? That would help keep the
correctness test focused while still demonstrating the memory benefit.
--
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]
