This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch branch-54
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/branch-54 by this push:
new 7ceb0ff84 fix(scheduler): keep GROUP BY-less aggregates when
propagating empty stages (#2194) (#2235)
7ceb0ff84 is described below
commit 7ceb0ff84a5164b59d6dd6ab7c9857a4a202d80e
Author: Andy Grove <[email protected]>
AuthorDate: Thu Aug 6 07:54:55 2026 -0600
fix(scheduler): keep GROUP BY-less aggregates when propagating empty stages
(#2194) (#2235)
(cherry picked from commit 7eeff3bebf4c0bb8debe2cbfb4e65b345e5f50ca)
---
.../state/aqe/optimizer_rule/propagate_empty.rs | 78 +++++++++++++++++++++-
1 file changed, 77 insertions(+), 1 deletion(-)
diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/propagate_empty.rs
b/ballista/scheduler/src/state/aqe/optimizer_rule/propagate_empty.rs
index 87d52a5bd..621375886 100644
--- a/ballista/scheduler/src/state/aqe/optimizer_rule/propagate_empty.rs
+++ b/ballista/scheduler/src/state/aqe/optimizer_rule/propagate_empty.rs
@@ -101,6 +101,9 @@ impl PropagateEmptyExecRule {
Ok(Transformed::yes(limit.input().clone()))
} else if let Some(aggregation) = plan.downcast_ref::<AggregateExec>()
&& is_empty_exec!(aggregation.input())
+ // An aggregate with no GROUP BY emits exactly one row even over
zero
+ // input rows (`sum` -> NULL, `count` -> 0), so it must be
preserved.
+ && !aggregation.group_expr().is_empty()
{
empty_exec!(aggregation)
} else if let Some(repartition) =
plan.downcast_ref::<RepartitionExec>()
@@ -328,8 +331,12 @@ mod tests {
use datafusion::{
arrow::datatypes::{DataType, Field, Schema},
common::{ColumnStatistics, JoinType, NullEquality, Statistics},
+ physical_expr::aggregate::AggregateExprBuilder,
physical_plan::{
- expressions::Column, joins::PartitionMode,
test::exec::StatisticsExec,
+ aggregates::{AggregateMode, PhysicalGroupBy},
+ expressions::Column,
+ joins::PartitionMode,
+ test::exec::StatisticsExec,
},
};
@@ -871,6 +878,75 @@ mod tests {
assert!(result.downcast_ref::<StatisticsExec>().is_some());
}
+ // ── Aggregate
────────────────────────────────────────────────────────────
+
+ /// Build an `AggregateExec` over `input` with the given grouping.
+ fn aggregate(
+ input: Arc<dyn ExecutionPlan>,
+ group_by: PhysicalGroupBy,
+ mode: AggregateMode,
+ ) -> Arc<dyn ExecutionPlan> {
+ let aggr = AggregateExprBuilder::new(
+ datafusion::functions_aggregate::count::count_udaf(),
+ vec![Arc::new(Column::new("a", 0))],
+ )
+ .schema(input.schema())
+ .alias("count(a)")
+ .build()
+ .unwrap();
+
+ Arc::new(
+ AggregateExec::try_new(
+ mode,
+ group_by,
+ vec![Arc::new(aggr)],
+ vec![None],
+ input,
+ schema(),
+ )
+ .unwrap(),
+ )
+ }
+
+ #[test]
+ fn aggregate_with_grouping_over_empty_eliminated() {
+ // GROUP BY over zero rows produces zero groups — collapsing is
correct.
+ let group_by = PhysicalGroupBy::new_single(vec![(
+ Arc::new(Column::new("a", 0)),
+ "a".into(),
+ )]);
+ let plan = aggregate(
+ Arc::new(EmptyExec::new(schema())),
+ group_by,
+ AggregateMode::Single,
+ );
+ assert_empty(&transform(plan));
+ }
+
+ #[test]
+ fn aggregate_without_grouping_over_empty_is_preserved() {
+ // No GROUP BY means exactly one output row even over zero input rows
+ // (`sum` -> NULL, `count` -> 0). Collapsing to EmptyExec loses that
row.
+ let plan = aggregate(
+ Arc::new(EmptyExec::new(schema())),
+ PhysicalGroupBy::default(),
+ AggregateMode::Single,
+ );
+ assert_untouched(&transform(plan));
+ }
+
+ #[test]
+ fn partial_aggregate_without_grouping_over_empty_is_preserved() {
+ // Same holds for a partial aggregate: it emits one row of accumulator
+ // state per partition, which the final aggregate needs.
+ let plan = aggregate(
+ Arc::new(EmptyExec::new(schema())),
+ PhysicalGroupBy::default(),
+ AggregateMode::Partial,
+ );
+ assert_untouched(&transform(plan));
+ }
+
// ── unknown stats — never optimised ─────────────────────────────────────
#[test]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
