Skip to content

Commit efd3ae8

Browse files
committed
fix(physical-optimizer): decline eager aggregation for grouping-set aggregates
`EagerAggregation` rebuilds the top grouping with `PhysicalGroupBy::new_single`, which cannot express `null_expr`, `groups`, or `has_grouping_set`. Applied to a `GROUPING SETS`/`ROLLUP`/`CUBE` aggregate it therefore returns a plain grouping, dropping both the grouping-set rows and the `__grouping_id` column that `group_fields` appends for them. The rule opts out of `schema_check` (COUNT->SUM widens non-null to nullable), so nothing catches the narrowed output schema. A parent projection or window keeps binding aggregate outputs at their original indices and trips the `col.name() == matching_name` assertion in `ProjectionMapping`: Internal error: Assertion failed: col.name() == matching_name: Input field name sum(store_sales.ss_ext_sales_price) does not match with the projection expression sum(store_sales.ss_net_profit). Decline the rewrite when the partial aggregate's grouping reports `has_grouping_set()`. The guard is on the *partial* grouping because `as_final` always clears the flag, folding `__grouping_id` into its expression list, so the final aggregate's grouping cannot report it. Adds two regression tests over the existing beneficial-join fixture, grouped by `ROLLUP(d_name)` so the grouping set is the only difference: one asserts no pre-aggregation is pushed, one asserts the aggregate output schema (including `__grouping_id`) is unchanged. Both fail without the guard. The plain-grouping fixture is reused via a new `agg_over_join_grouped` helper, so the accept path stays covered. Refs spiceai/spiceai#11827
1 parent b8b9592 commit efd3ae8

1 file changed

Lines changed: 110 additions & 5 deletions

File tree

datafusion/physical-optimizer/src/eager_aggregation.rs

Lines changed: 110 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,14 @@
112112
//! could ride a heuristic side choice but is declined for now.
113113
//! * **Decimal `AVG`.** Only Float64 `AVG` is recombined; a decimal result would
114114
//! need its exact output scale reproduced in the division projection.
115+
//! * **`GROUPING SETS`/`ROLLUP`/`CUBE`** (TPC-DS q36). The rebuilt top grouping
116+
//! goes through `PhysicalGroupBy::new_single`, which cannot express
117+
//! `null_expr`/`groups`/`has_grouping_set`, so the grouping-set rows and the
118+
//! `__grouping_id` column would be dropped. Supporting it means preserving the
119+
//! partial grouping's full shape — remapping `null_expr` alongside `expr` into
120+
//! the rebuilt join's schema and keeping `__grouping_id` at its original output
121+
//! index, since parent projections and window `PARTITION BY`/`ORDER BY` bind
122+
//! `GROUPING(...)` results positionally.
115123
//! * **`schema_check` is disabled** (COUNT→SUM widens non-null→nullable). A
116124
//! `coalesce(_, 0)` in a top projection would let it stay enabled.
117125
@@ -239,6 +247,23 @@ fn try_push_aggregate(
239247
}};
240248
}
241249

250+
// `GROUPING SETS`/`ROLLUP`/`CUBE` are not supported. The rewrite rebuilds the
251+
// top grouping with `PhysicalGroupBy::new_single`, which carries neither
252+
// `null_expr`/`groups` nor `has_grouping_set` — so a grouping-set aggregate
253+
// would come back as a plain grouping, dropping both the extra grouping-set
254+
// rows and the `__grouping_id` column that `group_fields` appends for them.
255+
// Because `schema_check` is opted out of (see above), the framework does not
256+
// catch the narrowed schema: the parent projection/window keeps binding
257+
// aggregate outputs at their original indices and hits
258+
// `col.name() == matching_name` in `ProjectionMapping` (TPC-DS q36).
259+
//
260+
// Guard on the *partial* grouping: `as_final` always clears
261+
// `has_grouping_set` (folding `__grouping_id` into its expression list), so
262+
// `top_final`'s grouping cannot report it.
263+
if top_partial.group_expr().has_grouping_set() {
264+
decline!("grouping sets (ROLLUP/CUBE/GROUPING SETS) are not supported");
265+
}
266+
242267
// Between the partial aggregate and the join the planner emits a *chain* of
243268
// column-only `ProjectionExec`s (this rule runs before ProjectionPushdown
244269
// folds them into the join). Peel the whole chain down to the HashJoinExec,
@@ -1159,6 +1184,26 @@ mod tests {
11591184
fact_key_ndv: usize,
11601185
dim_rows: usize,
11611186
dim_ndv: usize,
1187+
) -> Arc<dyn ExecutionPlan> {
1188+
agg_over_join_grouped(
1189+
fact_key_ndv,
1190+
dim_rows,
1191+
dim_ndv,
1192+
PhysicalGroupBy::new_single(vec![(
1193+
Arc::new(Column::new("d_name", 3)) as Arc<dyn PhysicalExpr>,
1194+
"d_name".to_string(),
1195+
)]),
1196+
)
1197+
}
1198+
1199+
/// Like [`agg_over_join_dim`], but with the partial aggregate's grouping
1200+
/// supplied by the caller, so a test can vary only the `PhysicalGroupBy`
1201+
/// (e.g. a grouping set) while holding the plan shape and statistics fixed.
1202+
fn agg_over_join_grouped(
1203+
fact_key_ndv: usize,
1204+
dim_rows: usize,
1205+
dim_ndv: usize,
1206+
group: PhysicalGroupBy,
11621207
) -> Arc<dyn ExecutionPlan> {
11631208
let fact = stats_leaf(
11641209
vec![
@@ -1205,14 +1250,10 @@ mod tests {
12051250
.build()
12061251
.unwrap(),
12071252
);
1208-
let group = PhysicalGroupBy::new_single(vec![(
1209-
Arc::new(Column::new("d_name", 3)) as Arc<dyn PhysicalExpr>,
1210-
"d_name".to_string(),
1211-
)]);
12121253
let partial = Arc::new(
12131254
AggregateExec::try_new(
12141255
AggregateMode::Partial,
1215-
group.clone(),
1256+
group,
12161257
vec![Arc::clone(&sum_expr)],
12171258
vec![None],
12181259
join,
@@ -1234,6 +1275,31 @@ mod tests {
12341275
) as Arc<dyn ExecutionPlan>
12351276
}
12361277

1278+
/// [`agg_over_join`]'s plan, but grouped by `ROLLUP(d_name)` instead of a plain
1279+
/// `GROUP BY d_name`. Statistics, join shape and aggregate are identical, so the
1280+
/// cost gate reaches the same verdict as in [`fires_on_beneficial_join`] and the
1281+
/// grouping set is the only difference.
1282+
///
1283+
/// The partial grouping carries `has_grouping_set`, so both aggregates output
1284+
/// `[d_name, __grouping_id, sum(f_amount)]`.
1285+
fn rollup_agg_over_join() -> Arc<dyn ExecutionPlan> {
1286+
// `ROLLUP(d_name)`: the plain group, plus the all-NULL super-aggregate row.
1287+
let group = PhysicalGroupBy::new(
1288+
vec![(
1289+
Arc::new(Column::new("d_name", 3)) as Arc<dyn PhysicalExpr>,
1290+
"d_name".to_string(),
1291+
)],
1292+
vec![(
1293+
Arc::new(Literal::new(ScalarValue::Utf8(None))) as Arc<dyn PhysicalExpr>,
1294+
"d_name".to_string(),
1295+
)],
1296+
vec![vec![false], vec![true]],
1297+
true,
1298+
);
1299+
assert!(group.has_grouping_set());
1300+
agg_over_join_grouped(100, 100, 100, group)
1301+
}
1302+
12371303
fn run_rule(plan: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
12381304
let mut opts = ConfigOptions::default();
12391305
opts.optimizer.enable_eager_aggregation = true;
@@ -1266,6 +1332,45 @@ mod tests {
12661332
);
12671333
}
12681334

1335+
// Decline path: the same beneficial join shape, but grouped by `ROLLUP(d_name)`.
1336+
// The rewrite cannot express a grouping set, so the rule must leave the plan
1337+
// alone rather than silently degrade it to a plain `GROUP BY`.
1338+
#[test]
1339+
fn declines_grouping_set_aggregate() {
1340+
let optimized = run_rule(rollup_agg_over_join());
1341+
assert!(
1342+
!join_child_is_aggregate(&optimized),
1343+
"expected no pre-aggregation below the join for a grouping-set aggregate, got:\n{}",
1344+
plan_str(&optimized)
1345+
);
1346+
}
1347+
1348+
// The defect this guards against: rebuilding the top grouping with
1349+
// `PhysicalGroupBy::new_single` drops `__grouping_id`, narrowing the aggregate
1350+
// output schema. `schema_check` is disabled for this rule, so nothing else
1351+
// catches it — a parent projection keeps binding by the original index and
1352+
// fails with `Input field name ... does not match with the projection
1353+
// expression ...` (TPC-DS q36).
1354+
#[test]
1355+
fn grouping_set_aggregate_preserves_output_schema() {
1356+
let plan = rollup_agg_over_join();
1357+
let before = plan.schema();
1358+
assert_eq!(
1359+
before
1360+
.fields()
1361+
.iter()
1362+
.map(|f| f.name().as_str())
1363+
.collect::<Vec<_>>(),
1364+
vec!["d_name", "__grouping_id", "sum(f_amount)"],
1365+
"fixture should carry the grouping-set `__grouping_id` column"
1366+
);
1367+
assert_eq!(
1368+
run_rule(plan).schema(),
1369+
before,
1370+
"eager aggregation must not change the aggregate output schema"
1371+
);
1372+
}
1373+
12691374
// Executable example (canonical SUM push). Input is the plan the planner emits
12701375
// for `SELECT d_name, SUM(f_amount) FROM fact JOIN dim ON f_dim = d_id GROUP BY
12711376
// d_name`: a two-phase aggregate directly over the join. The rule pushes a

0 commit comments

Comments
 (0)