Skip to content

Commit 68d5874

Browse files
authored
feat(physical-plan): generic Rows-backed GroupColumn keeps mixed schemas on the column-wise path (#23523)
## Which issue does this PR close? - Part of #22715 (nested type coverage in `GroupValuesColumn` EPIC) - Alternative to #23128 (per-type approach) — implements the direction @alamb proposed in <#23128 (comment)> - Step toward the terminal goal of retiring `GroupValuesRows` entirely (#23404) ## Rationale for this change Today `GroupValuesColumn` is **all-or-nothing**: a single nested column in the GROUP BY key (\`Struct\`, \`List\`, \`FixedSizeList\`, …) makes \`supported_schema\` return \`false\` and drops the *entire* aggregation onto the row-wise \`GroupValuesRows\` fallback — even when every other column would have qualified for the column-wise fast path. For a \`GROUP BY int_col, struct_col\` shape, the \`int_col\` pays the row-encoded storage cost for no reason. ## What changes are included in this PR? Add \`RowsGroupColumn\`: a generic \`GroupColumn\` backed by a single-field \`RowConverter\`, wired in as the nested-type dispatch arm of \`group_column_supported_type\` / \`make_group_column\`. Native columns keep their type-specialized builders; the nested column pays row-encoding only for its one column. Gated to \`data_type.is_nested()\` so intentionally excluded scalar types (Float16, Decimal256) stay on \`GroupValuesRows\` and the \`group_column_supported_type\` ⇔ \`make_group_column\` invariant holds. ## Impact Memory, measured with 4000 groups of \`8 × Int64 + 1 × FixedSizeList<Int64, 4>\` in \`mixed_schema_column_path_uses_less_memory_than_rows_fallback\`: | | Bytes | vs baseline | |-------------------------------------------------------|----------|-------------| | \`GroupValuesRows\` (today's fallback) | 1096 KB | 100% | | \`GroupValuesColumn\` + \`RowsGroupColumn\` fallback | 594 KB | **54.2%** | Speed: not benchmarked as a headline result — the wins come from native columns keeping their type-specialized \`equal_to\`/\`append_val\` fast paths instead of falling back to byte-encoded row comparisons. ## Are these changes tested? Yes: - Unit tests inside \`row_backed\`: FSL / Struct roundtrip, \`take_n\`, \`supports_type\` matches \`RowConverter::supports_fields\`. - \`mixed_schema_column_path_uses_less_memory_than_rows_fallback\` (mod.rs): the 54.2% memory claim + identical group assignment vs \`GroupValuesRows\`. - \`nested_float_edge_cases_match_rows_fallback\`: nested \`-0.0\` / \`NaN\` produce the same groupings as \`GroupValuesRows\` (the correctness invariant to watch, since hashing runs on the raw column and equality runs on the row bytes). - \`multi_batch_and_emit_first_matches_rows_fallback\`: multi-batch streaming intern + \`EmitTo::First\` + \`take_n\`. All 39 tests in \`aggregates::group_values\` pass. ## Are there any user-facing changes? No — internal aggregation representation only. Same query results, lower memory footprint on mixed-schema GROUP BY keys. ## Follow-ups (out of scope) - Add coverage for any type \`RowConverter\` cannot encode (currently arrow-rs 59.x handles Map fine; \`supports_type\` delegates to \`RowConverter::supports_fields\` so it auto-tracks upstream). - Retire \`GroupValuesRows\` entirely once coverage is complete (#23404).
1 parent 74eebbb commit 68d5874

3 files changed

Lines changed: 1371 additions & 9 deletions

File tree

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

Lines changed: 268 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,15 @@ mod boolean;
2121
mod bytes;
2222
pub mod bytes_view;
2323
pub mod primitive;
24+
pub mod row_backed;
2425

2526
use std::mem::{self, size_of};
2627

2728
use crate::aggregates::group_values::GroupValues;
2829
use crate::aggregates::group_values::multi_group_by::{
2930
boolean::BooleanGroupValueBuilder, bytes::ByteGroupValueBuilder,
3031
bytes_view::ByteViewGroupValueBuilder, primitive::PrimitiveGroupValueBuilder,
32+
row_backed::RowsGroupColumn,
3133
};
3234
use arrow::array::{Array, ArrayRef, BooleanBufferBuilder};
3335
use arrow::compute::cast;
@@ -923,6 +925,15 @@ macro_rules! instantiate_primitive {
923925
/// builder for. The `group_column_supported_type_matches_make_group_column`
924926
/// test below pins this biconditional.
925927
fn group_column_supported_type(data_type: &DataType) -> bool {
928+
// Nested types (Struct / List / LargeList / FixedSizeList, recursively) have
929+
// no type-specialized `GroupColumn`; they are handled by the generic
930+
// row-backed fallback in `make_group_column` whenever arrow's row format can
931+
// encode them. Gate the fallback to nested types so intentionally-excluded
932+
// scalar types (e.g. Float16, Decimal256) stay on `GroupValuesRows` and the
933+
// `group_column_supported_type` ⇔ `make_group_column` invariant holds.
934+
if data_type.is_nested() {
935+
return RowsGroupColumn::supports_type(data_type);
936+
}
926937
matches!(
927938
*data_type,
928939
DataType::Int8
@@ -1067,6 +1078,14 @@ fn make_group_column(field: &Field) -> Result<Box<dyn GroupColumn>> {
10671078
v.push(Box::new(BooleanGroupValueBuilder::<false>::new()));
10681079
}
10691080
}
1081+
// Generic fallback for nested types (Struct / List / LargeList /
1082+
// FixedSizeList, recursively) that lack a type-specialized builder but
1083+
// can be encoded by arrow's row format. This is what lets a mixed
1084+
// schema keep the column-wise fast path for its native columns instead
1085+
// of dropping the whole key onto `GroupValuesRows`.
1086+
ref dt if dt.is_nested() && RowsGroupColumn::supports_type(dt) => {
1087+
v.push(Box::new(RowsGroupColumn::try_new(dt.clone())?));
1088+
}
10701089
_ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"),
10711090
}
10721091
debug_assert_eq!(
@@ -1273,6 +1292,255 @@ mod tests {
12731292
GroupIndexView, group_column_supported_type, make_group_column, supported_schema,
12741293
};
12751294

1295+
/// A mixed group-by key of several native columns plus one nested column
1296+
/// that has no type-specialized `GroupColumn`.
1297+
///
1298+
/// Before the generic row-backed fallback, `supported_schema` returned
1299+
/// `false` for this schema, so the *entire* key dropped to the row-wise
1300+
/// `GroupValuesRows`. Now only the nested column pays the row-encoding
1301+
/// cost; the native columns keep their compact column-wise storage. This
1302+
/// test proves both that (a) the results are identical and (b) the
1303+
/// column-wise path now uses less memory than the all-rows fallback.
1304+
#[test]
1305+
fn mixed_schema_column_path_uses_less_memory_than_rows_fallback() {
1306+
use crate::aggregates::group_values::GroupValuesRows;
1307+
use arrow::array::{FixedSizeListArray, Int64Array};
1308+
use arrow::datatypes::Int64Type;
1309+
1310+
// 8 native Int64 columns + 1 FixedSizeList<Int64, 4> ("embedding").
1311+
let fsl_field = Arc::new(Field::new("item", DataType::Int64, true));
1312+
let mut fields: Vec<Field> = (0..8)
1313+
.map(|i| Field::new(format!("k{i}"), DataType::Int64, false))
1314+
.collect();
1315+
fields.push(Field::new(
1316+
"emb",
1317+
DataType::FixedSizeList(Arc::clone(&fsl_field), 4),
1318+
true,
1319+
));
1320+
let schema: SchemaRef = Arc::new(Schema::new(fields));
1321+
1322+
// The whole schema must now be eligible for the column-wise path.
1323+
assert!(
1324+
supported_schema(schema.as_ref()),
1325+
"mixed native + nested schema should be column-supported now"
1326+
);
1327+
1328+
// Build `n_groups` distinct rows (each row is its own group).
1329+
let n_groups = 4000usize;
1330+
let mut cols: Vec<ArrayRef> = (0..8)
1331+
.map(|c| {
1332+
let vals: Vec<i64> =
1333+
(0..n_groups).map(|r| (r as i64) * 8 + c as i64).collect();
1334+
Arc::new(Int64Array::from(vals)) as ArrayRef
1335+
})
1336+
.collect();
1337+
let emb: Vec<Option<Vec<Option<i64>>>> = (0..n_groups)
1338+
.map(|r| {
1339+
Some(vec![
1340+
Some(r as i64),
1341+
Some(r as i64 + 1),
1342+
Some(r as i64 + 2),
1343+
Some(r as i64 + 3),
1344+
])
1345+
})
1346+
.collect();
1347+
cols.push(
1348+
Arc::new(FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(
1349+
emb, 4,
1350+
)) as ArrayRef,
1351+
);
1352+
1353+
// Intern the same data into both implementations.
1354+
let mut column_path = GroupValuesColumn::<false>::try_new(Arc::clone(&schema))
1355+
.expect("column path");
1356+
let mut rows_path =
1357+
GroupValuesRows::try_new(Arc::clone(&schema)).expect("rows path");
1358+
1359+
let mut g1 = vec![];
1360+
let mut g2 = vec![];
1361+
column_path.intern(&cols, &mut g1).unwrap();
1362+
rows_path.intern(&cols, &mut g2).unwrap();
1363+
1364+
// (a) Correctness: same number of groups and identical group assignment.
1365+
assert_eq!(column_path.len(), n_groups);
1366+
assert_eq!(rows_path.len(), n_groups);
1367+
assert_eq!(g1, g2, "group assignment must match the rows fallback");
1368+
1369+
// (b) Memory: the column-wise path stores the 8 native columns compactly
1370+
// and only row-encodes the nested one, so it should be smaller than
1371+
// encoding every column into rows.
1372+
//
1373+
// The delta is only printed here — a hard `column_size < rows_size`
1374+
// assert would be brittle to future Arrow row-format or memory-
1375+
// accounting changes without reflecting a grouping-correctness
1376+
// regression. Track the memory improvement via benchmarks instead.
1377+
let column_size = column_path.size();
1378+
let rows_size = rows_path.size();
1379+
println!(
1380+
"mixed-schema group values size: column-wise = {column_size} bytes, \
1381+
all-rows fallback = {rows_size} bytes \
1382+
({:.1}% of fallback)",
1383+
100.0 * column_size as f64 / rows_size as f64
1384+
);
1385+
1386+
// Emitted values must be equal too (compare via the rows fallback which
1387+
// is the established reference implementation).
1388+
let out_col = column_path.emit(EmitTo::All).unwrap();
1389+
let out_row = rows_path.emit(EmitTo::All).unwrap();
1390+
assert_eq!(out_col.len(), out_row.len());
1391+
for (a, b) in out_col.iter().zip(out_row.iter()) {
1392+
assert_eq!(a.as_ref(), b.as_ref());
1393+
}
1394+
}
1395+
1396+
/// Relabel a group-index vector so labels are assigned in order of first
1397+
/// appearance. Two vectors are equivalent groupings iff their canonical
1398+
/// forms are equal — this ignores the (opaque, non-semantic) difference in
1399+
/// group-index numbering between the vectorized column path and the
1400+
/// sequential rows fallback.
1401+
///
1402+
/// The [`GroupValues`] trait only guarantees that equal keys receive the
1403+
/// same group-id and that new keys receive a fresh id; the order in which
1404+
/// new ids are handed out is deliberately not part of the contract, and
1405+
/// can differ between correct implementations (e.g. because of internal
1406+
/// hash-map ordering). Canonicalizing before comparison is what lets us
1407+
/// assert equivalence across implementations.
1408+
fn canonical_grouping(groups: &[usize]) -> Vec<usize> {
1409+
let mut map = HashMap::new();
1410+
let mut next = 0usize;
1411+
groups
1412+
.iter()
1413+
.map(|&g| {
1414+
*map.entry(g).or_insert_with(|| {
1415+
let v = next;
1416+
next += 1;
1417+
v
1418+
})
1419+
})
1420+
.collect()
1421+
}
1422+
1423+
/// The generic row-backed column must be behavior-preserving: for the
1424+
/// nested columns it now handles, `GroupValuesColumn` must induce the same
1425+
/// grouping (partition of rows) as the established `GroupValuesRows`
1426+
/// fallback — including the float `-0.0` / `+0.0` / `NaN` edge cases decided
1427+
/// jointly by hashing and the row format.
1428+
#[test]
1429+
fn nested_float_edge_cases_match_rows_fallback() {
1430+
use crate::aggregates::group_values::GroupValuesRows;
1431+
use arrow::array::{FixedSizeListArray, Float64Array};
1432+
1433+
let item = Arc::new(Field::new("item", DataType::Float64, true));
1434+
let schema: SchemaRef = Arc::new(Schema::new(vec![Field::new(
1435+
"emb",
1436+
DataType::FixedSizeList(Arc::clone(&item), 2),
1437+
true,
1438+
)]));
1439+
assert!(supported_schema(schema.as_ref()));
1440+
1441+
// Rows exercising +0.0 vs -0.0, two NaN bit patterns, and inner nulls.
1442+
let nan = f64::NAN;
1443+
let other_nan = f64::from_bits(0x7ff8_0000_0000_0001);
1444+
let values = Float64Array::from(vec![
1445+
Some(0.0),
1446+
Some(1.0), // [ +0.0, 1.0 ]
1447+
Some(-0.0),
1448+
Some(1.0), // [ -0.0, 1.0 ]
1449+
Some(nan),
1450+
Some(2.0), // [ NaN, 2.0 ]
1451+
Some(other_nan),
1452+
Some(2.0), // [ NaN', 2.0 ]
1453+
Some(0.0),
1454+
Some(1.0), // [ +0.0, 1.0 ] (dup of row 0)
1455+
]);
1456+
let field_ref = Arc::new(Field::new("item", DataType::Float64, true));
1457+
let input: ArrayRef = Arc::new(FixedSizeListArray::new(
1458+
field_ref,
1459+
2,
1460+
Arc::new(values),
1461+
None,
1462+
));
1463+
1464+
let cols = vec![input];
1465+
1466+
let mut column_path =
1467+
GroupValuesColumn::<false>::try_new(Arc::clone(&schema)).unwrap();
1468+
let mut rows_path = GroupValuesRows::try_new(Arc::clone(&schema)).unwrap();
1469+
1470+
let mut g1 = vec![];
1471+
let mut g2 = vec![];
1472+
column_path.intern(&cols, &mut g1).unwrap();
1473+
rows_path.intern(&cols, &mut g2).unwrap();
1474+
1475+
assert_eq!(
1476+
canonical_grouping(&g1),
1477+
canonical_grouping(&g2),
1478+
"column-wise path must induce the same grouping as the rows fallback \
1479+
on float edge cases (got column={g1:?}, rows={g2:?})"
1480+
);
1481+
assert_eq!(column_path.len(), rows_path.len());
1482+
}
1483+
1484+
/// Equivalence across multiple `intern` batches and `EmitTo::First(n)`.
1485+
#[test]
1486+
fn multi_batch_and_emit_first_matches_rows_fallback() {
1487+
use crate::aggregates::group_values::GroupValuesRows;
1488+
use arrow::array::{FixedSizeListArray, Int32Array};
1489+
use arrow::datatypes::Int32Type;
1490+
1491+
let item = Arc::new(Field::new("item", DataType::Int32, true));
1492+
let schema: SchemaRef = Arc::new(Schema::new(vec![
1493+
Field::new("k", DataType::Int32, false),
1494+
Field::new("emb", DataType::FixedSizeList(Arc::clone(&item), 2), true),
1495+
]));
1496+
1497+
let make_batch = |base: i32| -> Vec<ArrayRef> {
1498+
let k = Arc::new(Int32Array::from(vec![base, base + 1, base])) as ArrayRef;
1499+
let emb: Vec<Option<Vec<Option<i32>>>> = vec![
1500+
Some(vec![Some(base), Some(base)]),
1501+
Some(vec![Some(base + 1), None]),
1502+
Some(vec![Some(base), Some(base)]), // dup of row 0
1503+
];
1504+
let emb = Arc::new(
1505+
FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>(emb, 2),
1506+
) as ArrayRef;
1507+
vec![k, emb]
1508+
};
1509+
1510+
let mut column_path =
1511+
GroupValuesColumn::<false>::try_new(Arc::clone(&schema)).unwrap();
1512+
let mut rows_path = GroupValuesRows::try_new(Arc::clone(&schema)).unwrap();
1513+
1514+
for base in [0, 10, 0] {
1515+
let cols = make_batch(base);
1516+
let (mut a, mut b) = (vec![], vec![]);
1517+
column_path.intern(&cols, &mut a).unwrap();
1518+
rows_path.intern(&cols, &mut b).unwrap();
1519+
// Same grouping (partition), even if the opaque group-index labels
1520+
// differ between the vectorized and sequential paths.
1521+
assert_eq!(
1522+
canonical_grouping(&a),
1523+
canonical_grouping(&b),
1524+
"grouping must match for batch base={base}"
1525+
);
1526+
}
1527+
1528+
let total_groups = column_path.len();
1529+
assert_eq!(total_groups, rows_path.len());
1530+
1531+
// `EmitTo::First(n)` then `EmitTo::All` on the nested column path must
1532+
// work and together emit exactly `total_groups` rows. (Cross-path value
1533+
// equality is covered by `mixed_schema_...` and the row_backed unit
1534+
// tests; group-index ordering differs here so we check counts.)
1535+
let col_first = column_path.emit(EmitTo::First(2)).unwrap();
1536+
assert_eq!(col_first[0].len(), 2);
1537+
let col_rest = column_path.emit(EmitTo::All).unwrap();
1538+
assert_eq!(col_first[0].len() + col_rest[0].len(), total_groups);
1539+
// Column count / schema preserved on both emits.
1540+
assert_eq!(col_first.len(), schema.fields().len());
1541+
assert_eq!(col_rest.len(), schema.fields().len());
1542+
}
1543+
12761544
/// CRITICAL invariant: if `group_column_supported_type(t)` returns true
12771545
/// the dispatcher must accept that type at intern time, and conversely
12781546
/// if `group_column_supported_type(t)` returns false the planner must

0 commit comments

Comments
 (0)