feat(physical-plan): generic Rows-backed GroupColumn keeps mixed schemas on the column-wise path

datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs

@@ -21,13 +21,15 @@ mod boolean;
 mod bytes;
 pub mod bytes_view;
 pub mod primitive;
+pub mod row_backed;
 
 use std::mem::{self, size_of};
 
 use crate::aggregates::group_values::GroupValues;
 use crate::aggregates::group_values::multi_group_by::{
     boolean::BooleanGroupValueBuilder, bytes::ByteGroupValueBuilder,
     bytes_view::ByteViewGroupValueBuilder, primitive::PrimitiveGroupValueBuilder,
+    row_backed::RowsGroupColumn,
 };
 use arrow::array::{Array, ArrayRef, BooleanBufferBuilder};
 use arrow::compute::cast;
@@ -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);
+    }
     matches!(
         *data_type,
         DataType::Int8
@@ -1067,6 +1078,14 @@ fn make_group_column(field: &Field) -> Result<Box<dyn GroupColumn>> {
                 v.push(Box::new(BooleanGroupValueBuilder::<false>::new()));
             }
         }
+        // Generic fallback for nested types (Struct / List / LargeList /
+        // FixedSizeList, recursively) that lack a type-specialized builder but
+        // can be encoded by arrow's row format. This is what lets a mixed
+        // schema keep the column-wise fast path for its native columns instead
+        // of dropping the whole key onto `GroupValuesRows`.
+        ref dt if dt.is_nested() && RowsGroupColumn::supports_type(dt) => {
+            v.push(Box::new(RowsGroupColumn::try_new(dt.clone())?));
+        }
         _ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"),
     }
     debug_assert_eq!(
@@ -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()
+    }
+
+    /// The generic row-backed column must be behavior-preserving: for the
+    /// nested columns it now handles, `GroupValuesColumn` must induce the same
+    /// grouping (partition of rows) as the established `GroupValuesRows`
+    /// fallback — including the float `-0.0` / `+0.0` / `NaN` edge cases decided
+    /// jointly by hashing and the row format.
+    #[test]
+    fn nested_float_edge_cases_match_rows_fallback() {
+        use crate::aggregates::group_values::GroupValuesRows;
+        use arrow::array::{FixedSizeListArray, Float64Array};
+
+        let item = Arc::new(Field::new("item", DataType::Float64, true));
+        let schema: SchemaRef = Arc::new(Schema::new(vec![Field::new(
+            "emb",
+            DataType::FixedSizeList(Arc::clone(&item), 2),
+            true,
+        )]));
+        assert!(supported_schema(schema.as_ref()));
+
+        // Rows exercising +0.0 vs -0.0, two NaN bit patterns, and inner nulls.
+        let nan = f64::NAN;
+        let other_nan = f64::from_bits(0x7ff8_0000_0000_0001);
+        let values = Float64Array::from(vec![
+            Some(0.0),
+            Some(1.0), // [ +0.0, 1.0 ]
+            Some(-0.0),
+            Some(1.0), // [ -0.0, 1.0 ]
+            Some(nan),
+            Some(2.0), // [ NaN,  2.0 ]
+            Some(other_nan),
+            Some(2.0), // [ NaN', 2.0 ]
+            Some(0.0),
+            Some(1.0), // [ +0.0, 1.0 ]  (dup of row 0)
+        ]);
+        let field_ref = Arc::new(Field::new("item", DataType::Float64, true));
+        let input: ArrayRef = Arc::new(FixedSizeListArray::new(
+            field_ref,
+            2,
+            Arc::new(values),
+            None,
+        ));
+
+        let cols = vec![input];
+
+        let mut column_path =
+            GroupValuesColumn::<false>::try_new(Arc::clone(&schema)).unwrap();
+        let mut rows_path = GroupValuesRows::try_new(Arc::clone(&schema)).unwrap();
+
+        let mut g1 = vec![];
+        let mut g2 = vec![];
+        column_path.intern(&cols, &mut g1).unwrap();
+        rows_path.intern(&cols, &mut g2).unwrap();
+
+        assert_eq!(
+            canonical_grouping(&g1),
+            canonical_grouping(&g2),
+            "column-wise path must induce the same grouping as the rows fallback \
+             on float edge cases (got column={g1:?}, rows={g2:?})"
+        );
+        assert_eq!(column_path.len(), rows_path.len());
+    }
+
+    /// Equivalence across multiple `intern` batches and `EmitTo::First(n)`.
+    #[test]
+    fn multi_batch_and_emit_first_matches_rows_fallback() {
+        use crate::aggregates::group_values::GroupValuesRows;
+        use arrow::array::{FixedSizeListArray, Int32Array};
+        use arrow::datatypes::Int32Type;
+
+        let item = Arc::new(Field::new("item", DataType::Int32, true));
+        let schema: SchemaRef = Arc::new(Schema::new(vec![
+            Field::new("k", DataType::Int32, false),
+            Field::new("emb", DataType::FixedSizeList(Arc::clone(&item), 2), true),
+        ]));
+
+        let make_batch = |base: i32| -> Vec<ArrayRef> {
+            let k = Arc::new(Int32Array::from(vec![base, base + 1, base])) as ArrayRef;
+            let emb: Vec<Option<Vec<Option<i32>>>> = vec![
+                Some(vec![Some(base), Some(base)]),
+                Some(vec![Some(base + 1), None]),
+                Some(vec![Some(base), Some(base)]), // dup of row 0
+            ];
+            let emb = Arc::new(
+                FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>(emb, 2),
+            ) as ArrayRef;
+            vec![k, emb]
+        };
+
+        let mut column_path =
+            GroupValuesColumn::<false>::try_new(Arc::clone(&schema)).unwrap();
+        let mut rows_path = GroupValuesRows::try_new(Arc::clone(&schema)).unwrap();
+
+        for base in [0, 10, 0] {
+            let cols = make_batch(base);
+            let (mut a, mut b) = (vec![], vec![]);
+            column_path.intern(&cols, &mut a).unwrap();
+            rows_path.intern(&cols, &mut b).unwrap();
+            // Same grouping (partition), even if the opaque group-index labels
+            // differ between the vectorized and sequential paths.
+            assert_eq!(
+                canonical_grouping(&a),
+                canonical_grouping(&b),
+                "grouping must match for batch base={base}"
+            );
+        }
+
+        let total_groups = column_path.len();
+        assert_eq!(total_groups, rows_path.len());
+
+        // `EmitTo::First(n)` then `EmitTo::All` on the nested column path must
+        // work and together emit exactly `total_groups` rows. (Cross-path value
+        // equality is covered by `mixed_schema_...` and the row_backed unit
+        // tests; group-index ordering differs here so we check counts.)
+        let col_first = column_path.emit(EmitTo::First(2)).unwrap();
+        assert_eq!(col_first[0].len(), 2);
+        let col_rest = column_path.emit(EmitTo::All).unwrap();
+        assert_eq!(col_first[0].len() + col_rest[0].len(), total_groups);
+        // Column count / schema preserved on both emits.
+        assert_eq!(col_first.len(), schema.fields().len());
+        assert_eq!(col_rest.len(), schema.fields().len());
+    }
+
     /// CRITICAL invariant: if `group_column_supported_type(t)` returns true
     /// the dispatcher must accept that type at intern time, and conversely
     /// if `group_column_supported_type(t)` returns false the planner must