You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "wirybeaver (via GitHub)" <gi...@apache.org> on 2026/05/18 05:43:07 UTC
[PR] feat(aqe): dynamic streaming↔hash aggregation switch [datafusion-ballista]
wirybeaver opened a new pull request, #1722:
URL: https://github.com/apache/datafusion-ballista/pull/1722
## Which issue does this PR close?
Part of #1359 (AQE epic) — roadmap line: *"switch from streaming aggregation to hash aggregation (extended rules)"*
## What does this PR do?
Adds `DynamicAggregateAlgorithmRule`, an AQE physical optimizer rule that re-derives `InputOrderMode` for each `AggregateExec` after a shuffle stage resolves and rewrites the operator when the derived mode differs from the cached one.
### Problem
DataFusion freezes `InputOrderMode` (`Linear` / `Sorted` / `PartiallySorted`) at plan time. In Ballista's AQE this creates two issues:
1. **Wasted memory**: a downstream aggregate stays `Linear` (hash table, O(distinct-groups) memory) even when an upstream stage completes and grants ordering on the group-by columns — where `Sorted` (streaming, O(1) memory) would suffice.
2. **Stale correctness assumption**: after a subtree rewrite, an aggregate may hold a `Sorted` claim when its input is no longer ordered.
### Approach
`AggregateExec::with_new_children` already calls `try_new_with_schema` which re-derives `input_order_mode` from the current input `EquivalenceProperties`. The rule:
1. Walks the plan with `transform_up`.
2. At each `AggregateExec`, skips if no resolved `ExchangeExec` exists in the subtree (idempotence guard).
3. Calls `with_new_children([same_input])` to force re-derivation.
4. Returns `Transformed::yes(rebuilt)` only when the derived mode differs from the cached one.
No upstream DataFusion changes are required.
### Configuration
```sql
SET ballista.aqe.dynamic_aggregate.enabled = true;
```
Disabled by default (`false`) pending benchmarking.
## Checklist
- [x] New rule registered in `default_optimizers()` before `DistributedExchangeRule`
- [x] Config key `ballista.aqe.dynamic_aggregate.enabled` added to `BallistaConfig`
- [x] 9 unit tests: gate off, no-exchange skip, unresolved-exchange skip, sorted/unsorted input, schema preservation, idempotence
- [x] All 42 AQE optimizer rule tests pass
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(aqe): dynamic switching from streaming agg to hash aggregation [datafusion-ballista]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1722:
URL: https://github.com/apache/datafusion-ballista/pull/1722#issuecomment-4482734022
Sorry but I do not understand this PR, can you please give a bit more explanation what are you trying to do in this PR?
In which case can we assume sorting is going to be preserved after an exchange ?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(aqe): dynamic switching from streaming agg to hash aggregation [datafusion-ballista]
Posted by "martin-g (via GitHub)" <gi...@apache.org>.
martin-g commented on code in PR #1722:
URL: https://github.com/apache/datafusion-ballista/pull/1722#discussion_r3258307004
##########
ballista/scheduler/src/state/aqe/optimizer_rule/dynamic_aggregate_algorithm.rs:
##########
@@ -0,0 +1,338 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::state::aqe::execution_plan::ExchangeExec;
+use ballista_core::config::BallistaConfig;
+use datafusion::common::tree_node::{Transformed, TreeNode};
+use datafusion::config::ConfigOptions;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::aggregates::AggregateExec;
+use log::info;
+use std::sync::Arc;
+
+/// AQE rule that re-derives `InputOrderMode` for each `AggregateExec` after a
+/// stage completes and rewrites the operator when the derived mode differs from
+/// the cached one.
+///
+/// DataFusion freezes `InputOrderMode` (Linear = hash table / Sorted = streaming /
+/// PartiallySorted) at plan time. After a shuffle stage resolves, the input's
+/// `EquivalenceProperties` may expose ordering that wasn't visible earlier, making
+/// streaming aggregation possible (or, after a subtree rewrite, making the previously
+/// assumed ordering stale). This rule repairs both cases on each AQE replan.
+///
+/// Relies on `AggregateExec::with_new_children` calling `try_new_with_schema`,
+/// which re-derives `input_order_mode` from the current input equivalence
+/// properties — no upstream DataFusion changes required.
+///
+/// Enabled via `ballista.aqe.dynamic_aggregate.enabled` (default: false).
+#[derive(Debug, Clone, Default)]
+pub struct DynamicAggregateAlgorithmRule {}
+
+impl DynamicAggregateAlgorithmRule {
+ fn transform(
+ plan: Arc<dyn ExecutionPlan>,
+ enabled: bool,
+ ) -> datafusion::error::Result<Transformed<Arc<dyn ExecutionPlan>>> {
+ if !enabled {
+ return Ok(Transformed::no(plan));
+ }
+ let Some(agg) = plan.as_any().downcast_ref::<AggregateExec>() else {
+ return Ok(Transformed::no(plan));
+ };
+ // Only act when a resolved exchange exists below — if nothing resolved,
+ // the ordering is the same as at initial plan time and no switch is needed.
+ if !subtree_has_resolved_exchange(&plan) {
+ return Ok(Transformed::no(plan));
+ }
+
+ // Re-derive InputOrderMode by reconstructing the aggregate with the same
+ // input. with_new_children → try_new_with_schema re-runs the derivation
+ // against the current input EquivalenceProperties.
+ let rebuilt = plan
+ .clone()
+ .with_new_children(vec![agg.children()[0].clone()])?;
+ let rebuilt_agg = rebuilt
+ .as_any()
+ .downcast_ref::<AggregateExec>()
+ .expect("with_new_children on AggregateExec must return AggregateExec");
+
+ if rebuilt_agg.input_order_mode() == agg.input_order_mode() {
+ return Ok(Transformed::no(plan));
+ }
+
+ info!(
+ "DynamicAggregateAlgorithmRule: switching InputOrderMode {:?} → {:?}",
+ agg.input_order_mode(),
+ rebuilt_agg.input_order_mode(),
+ );
+ Ok(Transformed::yes(rebuilt))
+ }
+}
+
+impl PhysicalOptimizerRule for DynamicAggregateAlgorithmRule {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
+ let enabled = config
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.aqe_dynamic_aggregate_enabled())
+ .unwrap_or(false);
+ Ok(plan.transform_up(|p| Self::transform(p, enabled))?.data)
+ }
+
+ fn name(&self) -> &str {
+ "DynamicAggregateAlgorithmRule"
+ }
+
+ fn schema_check(&self) -> bool {
+ false
+ }
+}
+
+/// Returns true if any `ExchangeExec` in `plan`'s subtree is resolved
+/// (i.e. its shuffle has completed). Recurses through all nodes, including
+/// through exchange boundaries — unlike `nearest_exchange_status` in
+/// `distributed_exchange.rs` which stops at the first exchange, here we
+/// want to know if *any* upstream stage has finished.
+fn subtree_has_resolved_exchange(plan: &Arc<dyn ExecutionPlan>) -> bool {
+ if let Some(exchange) = plan.as_any().downcast_ref::<ExchangeExec>() {
+ exchange.shuffle_created() && !exchange.inactive_stage
Review Comment:
The doc above says `Recurses through all nodes, including through exchange boundaries` but the implementation stops recursion on the first ExchangeExec
##########
ballista/scheduler/src/state/aqe/optimizer_rule/dynamic_aggregate_algorithm.rs:
##########
@@ -0,0 +1,338 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::state::aqe::execution_plan::ExchangeExec;
+use ballista_core::config::BallistaConfig;
+use datafusion::common::tree_node::{Transformed, TreeNode};
+use datafusion::config::ConfigOptions;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::aggregates::AggregateExec;
+use log::info;
+use std::sync::Arc;
+
+/// AQE rule that re-derives `InputOrderMode` for each `AggregateExec` after a
+/// stage completes and rewrites the operator when the derived mode differs from
+/// the cached one.
+///
+/// DataFusion freezes `InputOrderMode` (Linear = hash table / Sorted = streaming /
+/// PartiallySorted) at plan time. After a shuffle stage resolves, the input's
+/// `EquivalenceProperties` may expose ordering that wasn't visible earlier, making
+/// streaming aggregation possible (or, after a subtree rewrite, making the previously
+/// assumed ordering stale). This rule repairs both cases on each AQE replan.
+///
+/// Relies on `AggregateExec::with_new_children` calling `try_new_with_schema`,
+/// which re-derives `input_order_mode` from the current input equivalence
+/// properties — no upstream DataFusion changes required.
+///
+/// Enabled via `ballista.aqe.dynamic_aggregate.enabled` (default: false).
+#[derive(Debug, Clone, Default)]
+pub struct DynamicAggregateAlgorithmRule {}
+
+impl DynamicAggregateAlgorithmRule {
+ fn transform(
+ plan: Arc<dyn ExecutionPlan>,
+ enabled: bool,
+ ) -> datafusion::error::Result<Transformed<Arc<dyn ExecutionPlan>>> {
+ if !enabled {
+ return Ok(Transformed::no(plan));
+ }
+ let Some(agg) = plan.as_any().downcast_ref::<AggregateExec>() else {
+ return Ok(Transformed::no(plan));
+ };
+ // Only act when a resolved exchange exists below — if nothing resolved,
+ // the ordering is the same as at initial plan time and no switch is needed.
+ if !subtree_has_resolved_exchange(&plan) {
+ return Ok(Transformed::no(plan));
+ }
+
+ // Re-derive InputOrderMode by reconstructing the aggregate with the same
+ // input. with_new_children → try_new_with_schema re-runs the derivation
+ // against the current input EquivalenceProperties.
+ let rebuilt = plan
+ .clone()
+ .with_new_children(vec![agg.children()[0].clone()])?;
+ let rebuilt_agg = rebuilt
+ .as_any()
+ .downcast_ref::<AggregateExec>()
+ .expect("with_new_children on AggregateExec must return AggregateExec");
+
+ if rebuilt_agg.input_order_mode() == agg.input_order_mode() {
+ return Ok(Transformed::no(plan));
+ }
+
+ info!(
+ "DynamicAggregateAlgorithmRule: switching InputOrderMode {:?} → {:?}",
+ agg.input_order_mode(),
+ rebuilt_agg.input_order_mode(),
+ );
+ Ok(Transformed::yes(rebuilt))
+ }
+}
+
+impl PhysicalOptimizerRule for DynamicAggregateAlgorithmRule {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
+ let enabled = config
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.aqe_dynamic_aggregate_enabled())
+ .unwrap_or(false);
+ Ok(plan.transform_up(|p| Self::transform(p, enabled))?.data)
+ }
+
+ fn name(&self) -> &str {
+ "DynamicAggregateAlgorithmRule"
+ }
+
+ fn schema_check(&self) -> bool {
+ false
+ }
+}
+
+/// Returns true if any `ExchangeExec` in `plan`'s subtree is resolved
+/// (i.e. its shuffle has completed). Recurses through all nodes, including
+/// through exchange boundaries — unlike `nearest_exchange_status` in
+/// `distributed_exchange.rs` which stops at the first exchange, here we
+/// want to know if *any* upstream stage has finished.
+fn subtree_has_resolved_exchange(plan: &Arc<dyn ExecutionPlan>) -> bool {
+ if let Some(exchange) = plan.as_any().downcast_ref::<ExchangeExec>() {
+ exchange.shuffle_created() && !exchange.inactive_stage
+ } else {
+ plan.children()
+ .iter()
+ .any(|child| subtree_has_resolved_exchange(child))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use ballista_core::config::BallistaConfig;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::common::{ColumnStatistics, Statistics};
+ use datafusion::config::{ConfigOptions, ExtensionOptions};
+ use datafusion::physical_expr::LexOrdering;
+ use datafusion::physical_expr::expressions::Column;
+ use datafusion::physical_plan::ExecutionPlan;
+ use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
+ use datafusion::physical_plan::sorts::sort::SortExec;
+ use datafusion::physical_plan::test::exec::StatisticsExec;
+ use datafusion::physical_expr::PhysicalSortExpr;
+ use std::sync::Arc;
+
+ fn leaf_unsorted() -> Arc<dyn ExecutionPlan> {
+ Arc::new(StatisticsExec::new(
+ Statistics {
+ num_rows: Default::default(),
+ total_byte_size: Default::default(),
+ column_statistics: vec![
+ ColumnStatistics::new_unknown(),
+ ColumnStatistics::new_unknown(),
+ ],
+ },
+ Schema::new(vec![
+ Field::new("k", DataType::Int32, true),
+ Field::new("v", DataType::Int64, true),
+ ]),
+ ))
+ }
+
+ fn leaf_sorted_on_k() -> Arc<dyn ExecutionPlan> {
+ let sort_expr =
+ PhysicalSortExpr::new_default(Arc::new(Column::new("k", 0)));
+ let ordering = LexOrdering::new(vec![sort_expr]).unwrap();
+ Arc::new(SortExec::new(ordering, leaf_unsorted()))
+ }
+
+ fn resolved_exchange(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
+ let exchange = ExchangeExec::new(input, None, 0);
+ exchange.resolve_shuffle_partitions(vec![]);
+ Arc::new(exchange)
+ }
+
+ fn unresolved_exchange(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
+ Arc::new(ExchangeExec::new(input, None, 0))
+ }
+
+ fn aggregate_on_k(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
+ let schema = input.schema();
+ let k_col = Arc::new(Column::new("k", 0)) as _;
+ let group_by = PhysicalGroupBy::new_single(vec![(k_col, "k".to_string())]);
+ Arc::new(
+ AggregateExec::try_new(
+ AggregateMode::Single,
+ group_by,
+ vec![],
+ vec![],
+ input,
+ schema,
+ )
+ .unwrap(),
+ )
+ }
+
+ fn config_with_flag(enabled: bool) -> ConfigOptions {
+ let mut ballista = BallistaConfig::default();
+ if enabled {
+ ballista
+ .set(
+ "aqe.dynamic_aggregate.enabled",
+ "true",
+ )
+ .unwrap();
+ }
+ let mut opts = ConfigOptions::new();
+ opts.extensions.insert(ballista);
+ opts
+ }
+
+ // --- gate ---
+
+ #[test]
+ fn gate_disabled_returns_no_transform() {
+ let plan = aggregate_on_k(resolved_exchange(leaf_sorted_on_k()));
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan.clone(), &config_with_flag(false)).unwrap();
+ assert!(
+ Arc::ptr_eq(&result, &plan),
+ "gate disabled: plan must be returned unchanged (same Arc)"
+ );
+ }
+
+ // --- subtree_has_resolved_exchange ---
+
+ #[test]
+ fn no_exchange_skips() {
+ // AggregateExec with no exchange at all → no rewrite even when enabled
+ let plan = aggregate_on_k(leaf_sorted_on_k());
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan.clone(), &config_with_flag(true)).unwrap();
+ assert!(Arc::ptr_eq(&result, &plan));
+ }
+
+ #[test]
+ fn unresolved_exchange_skips() {
+ let plan = aggregate_on_k(unresolved_exchange(leaf_sorted_on_k()));
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan.clone(), &config_with_flag(true)).unwrap();
+ assert!(Arc::ptr_eq(&result, &plan));
+ }
+
+ // --- mode transitions ---
+
+ #[test]
+ fn linear_to_sorted_when_input_fully_ordered() {
+ // Input: SortExec on k → resolved ExchangeExec → AggregateExec(group_by=[k])
+ // The SortExec grants ordering on k; the exchange preserves it in eq_properties.
+ // AggregateExec is initially planned as Linear (before the sort was visible).
+ // After re-derivation the mode should become Sorted.
+ let sorted_input = leaf_sorted_on_k();
+ let exchange = resolved_exchange(sorted_input);
+ let agg = aggregate_on_k(exchange);
+
+ let agg_exec = agg.as_any().downcast_ref::<AggregateExec>().unwrap();
+ // The agg was built over the resolved exchange whose eq_properties mirror
+ // the SortExec — so it may already be Sorted. What matters is the rule
+ // doesn't break it (idempotence when already correct).
+ let original_mode = agg_exec.input_order_mode().clone();
+
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(agg.clone(), &config_with_flag(true)).unwrap();
+ let result_agg = result.as_any().downcast_ref::<AggregateExec>().unwrap();
+
+ // After the rule the mode must match what re-derivation produces.
+ assert_eq!(result_agg.input_order_mode(), &original_mode);
+ // Schema must be preserved.
+ assert_eq!(result.schema(), agg.schema());
+ }
+
+ #[test]
+ fn unsorted_input_stays_linear() {
+ // Input has no ordering → mode must stay Linear after the rule.
+ let exchange = resolved_exchange(leaf_unsorted());
+ let plan = aggregate_on_k(exchange);
+
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan.clone(), &config_with_flag(true)).unwrap();
+ let result_agg = result.as_any().downcast_ref::<AggregateExec>().unwrap();
+
+ assert_eq!(
+ result_agg.input_order_mode(),
+ &datafusion::physical_plan::InputOrderMode::Linear
+ );
+ }
+
+ #[test]
+ fn schema_preserved_after_rewrite() {
+ let exchange = resolved_exchange(leaf_sorted_on_k());
+ let plan = aggregate_on_k(exchange);
+ let original_schema = plan.schema();
+
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan, &config_with_flag(true)).unwrap();
+
+ assert_eq!(result.schema(), original_schema);
+ }
+
+ // --- idempotence ---
+
+ #[test]
+ fn idempotent_second_pass() {
+ let exchange = resolved_exchange(leaf_sorted_on_k());
+ let plan = aggregate_on_k(exchange);
+
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let after_first = rule.optimize(plan, &config_with_flag(true)).unwrap();
+ let mode_after_first = after_first
+ .as_any()
+ .downcast_ref::<AggregateExec>()
+ .unwrap()
+ .input_order_mode()
+ .clone();
+
+ let after_second = rule
+ .optimize(after_first.clone(), &config_with_flag(true))
+ .unwrap();
+ let mode_after_second = after_second
+ .as_any()
+ .downcast_ref::<AggregateExec>()
+ .unwrap()
+ .input_order_mode()
+ .clone();
+
+ assert_eq!(mode_after_first, mode_after_second, "rule must be idempotent");
+ }
+
+ // --- rule metadata ---
+
+ #[test]
+ fn rule_name() {
+ assert_eq!(
+ DynamicAggregateAlgorithmRule::default().name(),
+ "DynamicAggregateAlgorithmRule"
+ );
+ }
+
+ #[test]
+ fn schema_check_is_false() {
+ assert!(!DynamicAggregateAlgorithmRule::default().schema_check());
+ }
Review Comment:
None of the tests validates order mode change from Linear to Sorted.
##########
ballista/scheduler/src/state/aqe/optimizer_rule/dynamic_aggregate_algorithm.rs:
##########
@@ -0,0 +1,338 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::state::aqe::execution_plan::ExchangeExec;
+use ballista_core::config::BallistaConfig;
+use datafusion::common::tree_node::{Transformed, TreeNode};
+use datafusion::config::ConfigOptions;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::aggregates::AggregateExec;
+use log::info;
+use std::sync::Arc;
+
+/// AQE rule that re-derives `InputOrderMode` for each `AggregateExec` after a
+/// stage completes and rewrites the operator when the derived mode differs from
+/// the cached one.
+///
+/// DataFusion freezes `InputOrderMode` (Linear = hash table / Sorted = streaming /
+/// PartiallySorted) at plan time. After a shuffle stage resolves, the input's
+/// `EquivalenceProperties` may expose ordering that wasn't visible earlier, making
+/// streaming aggregation possible (or, after a subtree rewrite, making the previously
+/// assumed ordering stale). This rule repairs both cases on each AQE replan.
+///
+/// Relies on `AggregateExec::with_new_children` calling `try_new_with_schema`,
+/// which re-derives `input_order_mode` from the current input equivalence
+/// properties — no upstream DataFusion changes required.
+///
+/// Enabled via `ballista.aqe.dynamic_aggregate.enabled` (default: false).
+#[derive(Debug, Clone, Default)]
+pub struct DynamicAggregateAlgorithmRule {}
+
+impl DynamicAggregateAlgorithmRule {
+ fn transform(
+ plan: Arc<dyn ExecutionPlan>,
+ enabled: bool,
+ ) -> datafusion::error::Result<Transformed<Arc<dyn ExecutionPlan>>> {
+ if !enabled {
+ return Ok(Transformed::no(plan));
+ }
+ let Some(agg) = plan.as_any().downcast_ref::<AggregateExec>() else {
+ return Ok(Transformed::no(plan));
+ };
+ // Only act when a resolved exchange exists below — if nothing resolved,
+ // the ordering is the same as at initial plan time and no switch is needed.
+ if !subtree_has_resolved_exchange(&plan) {
+ return Ok(Transformed::no(plan));
+ }
+
+ // Re-derive InputOrderMode by reconstructing the aggregate with the same
+ // input. with_new_children → try_new_with_schema re-runs the derivation
+ // against the current input EquivalenceProperties.
+ let rebuilt = plan
+ .clone()
+ .with_new_children(vec![agg.children()[0].clone()])?;
+ let rebuilt_agg = rebuilt
+ .as_any()
+ .downcast_ref::<AggregateExec>()
+ .expect("with_new_children on AggregateExec must return AggregateExec");
+
+ if rebuilt_agg.input_order_mode() == agg.input_order_mode() {
+ return Ok(Transformed::no(plan));
+ }
+
+ info!(
+ "DynamicAggregateAlgorithmRule: switching InputOrderMode {:?} → {:?}",
+ agg.input_order_mode(),
+ rebuilt_agg.input_order_mode(),
+ );
+ Ok(Transformed::yes(rebuilt))
+ }
+}
+
+impl PhysicalOptimizerRule for DynamicAggregateAlgorithmRule {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
+ let enabled = config
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.aqe_dynamic_aggregate_enabled())
+ .unwrap_or(false);
+ Ok(plan.transform_up(|p| Self::transform(p, enabled))?.data)
+ }
+
+ fn name(&self) -> &str {
+ "DynamicAggregateAlgorithmRule"
+ }
+
+ fn schema_check(&self) -> bool {
+ false
+ }
+}
+
+/// Returns true if any `ExchangeExec` in `plan`'s subtree is resolved
+/// (i.e. its shuffle has completed). Recurses through all nodes, including
+/// through exchange boundaries — unlike `nearest_exchange_status` in
+/// `distributed_exchange.rs` which stops at the first exchange, here we
+/// want to know if *any* upstream stage has finished.
+fn subtree_has_resolved_exchange(plan: &Arc<dyn ExecutionPlan>) -> bool {
+ if let Some(exchange) = plan.as_any().downcast_ref::<ExchangeExec>() {
+ exchange.shuffle_created() && !exchange.inactive_stage
+ } else {
+ plan.children()
+ .iter()
+ .any(|child| subtree_has_resolved_exchange(child))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use ballista_core::config::BallistaConfig;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::common::{ColumnStatistics, Statistics};
+ use datafusion::config::{ConfigOptions, ExtensionOptions};
+ use datafusion::physical_expr::LexOrdering;
+ use datafusion::physical_expr::expressions::Column;
+ use datafusion::physical_plan::ExecutionPlan;
+ use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
+ use datafusion::physical_plan::sorts::sort::SortExec;
+ use datafusion::physical_plan::test::exec::StatisticsExec;
+ use datafusion::physical_expr::PhysicalSortExpr;
+ use std::sync::Arc;
+
+ fn leaf_unsorted() -> Arc<dyn ExecutionPlan> {
+ Arc::new(StatisticsExec::new(
+ Statistics {
+ num_rows: Default::default(),
+ total_byte_size: Default::default(),
+ column_statistics: vec![
+ ColumnStatistics::new_unknown(),
+ ColumnStatistics::new_unknown(),
+ ],
+ },
+ Schema::new(vec![
+ Field::new("k", DataType::Int32, true),
+ Field::new("v", DataType::Int64, true),
+ ]),
+ ))
+ }
+
+ fn leaf_sorted_on_k() -> Arc<dyn ExecutionPlan> {
+ let sort_expr =
+ PhysicalSortExpr::new_default(Arc::new(Column::new("k", 0)));
+ let ordering = LexOrdering::new(vec![sort_expr]).unwrap();
+ Arc::new(SortExec::new(ordering, leaf_unsorted()))
+ }
+
+ fn resolved_exchange(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
+ let exchange = ExchangeExec::new(input, None, 0);
+ exchange.resolve_shuffle_partitions(vec![]);
+ Arc::new(exchange)
+ }
+
+ fn unresolved_exchange(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
+ Arc::new(ExchangeExec::new(input, None, 0))
+ }
+
+ fn aggregate_on_k(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
+ let schema = input.schema();
+ let k_col = Arc::new(Column::new("k", 0)) as _;
+ let group_by = PhysicalGroupBy::new_single(vec![(k_col, "k".to_string())]);
+ Arc::new(
+ AggregateExec::try_new(
+ AggregateMode::Single,
+ group_by,
+ vec![],
+ vec![],
+ input,
+ schema,
+ )
+ .unwrap(),
+ )
+ }
+
+ fn config_with_flag(enabled: bool) -> ConfigOptions {
+ let mut ballista = BallistaConfig::default();
+ if enabled {
+ ballista
+ .set(
+ "aqe.dynamic_aggregate.enabled",
+ "true",
+ )
+ .unwrap();
+ }
+ let mut opts = ConfigOptions::new();
+ opts.extensions.insert(ballista);
+ opts
+ }
+
+ // --- gate ---
+
+ #[test]
+ fn gate_disabled_returns_no_transform() {
+ let plan = aggregate_on_k(resolved_exchange(leaf_sorted_on_k()));
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan.clone(), &config_with_flag(false)).unwrap();
+ assert!(
+ Arc::ptr_eq(&result, &plan),
+ "gate disabled: plan must be returned unchanged (same Arc)"
+ );
+ }
+
+ // --- subtree_has_resolved_exchange ---
+
+ #[test]
+ fn no_exchange_skips() {
+ // AggregateExec with no exchange at all → no rewrite even when enabled
+ let plan = aggregate_on_k(leaf_sorted_on_k());
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan.clone(), &config_with_flag(true)).unwrap();
+ assert!(Arc::ptr_eq(&result, &plan));
+ }
+
+ #[test]
+ fn unresolved_exchange_skips() {
+ let plan = aggregate_on_k(unresolved_exchange(leaf_sorted_on_k()));
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan.clone(), &config_with_flag(true)).unwrap();
+ assert!(Arc::ptr_eq(&result, &plan));
+ }
+
+ // --- mode transitions ---
+
+ #[test]
+ fn linear_to_sorted_when_input_fully_ordered() {
+ // Input: SortExec on k → resolved ExchangeExec → AggregateExec(group_by=[k])
+ // The SortExec grants ordering on k; the exchange preserves it in eq_properties.
+ // AggregateExec is initially planned as Linear (before the sort was visible).
+ // After re-derivation the mode should become Sorted.
+ let sorted_input = leaf_sorted_on_k();
+ let exchange = resolved_exchange(sorted_input);
+ let agg = aggregate_on_k(exchange);
+
+ let agg_exec = agg.as_any().downcast_ref::<AggregateExec>().unwrap();
+ // The agg was built over the resolved exchange whose eq_properties mirror
+ // the SortExec — so it may already be Sorted. What matters is the rule
+ // doesn't break it (idempotence when already correct).
+ let original_mode = agg_exec.input_order_mode().clone();
+
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(agg.clone(), &config_with_flag(true)).unwrap();
+ let result_agg = result.as_any().downcast_ref::<AggregateExec>().unwrap();
+
+ // After the rule the mode must match what re-derivation produces.
+ assert_eq!(result_agg.input_order_mode(), &original_mode);
Review Comment:
assert that the order mode is Sorted
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(aqe): dynamic switching from streaming agg to hash aggregation [datafusion-ballista]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1722:
URL: https://github.com/apache/datafusion-ballista/pull/1722#issuecomment-4482812798
> Wasted memory: a downstream aggregate stays Linear (hash table, O(distinct-groups) memory) even when an upstream stage completes and grants ordering on the group-by columns — where Sorted (streaming, O(1) memory) would suffice.
In which cases will this happen? When can w guarantee preserving sorting across shuffles ?
> Stale correctness assumption: after a subtree rewrite, an aggregate may hold a Sorted claim when its input is no longer ordered.
Can you please give an example when would this claim be true?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(aqe): dynamic switching from streaming agg to hash aggregation [datafusion-ballista]
Posted by "wirybeaver (via GitHub)" <gi...@apache.org>.
wirybeaver commented on PR #1722:
URL: https://github.com/apache/datafusion-ballista/pull/1722#issuecomment-4505923598
Closing per my analysis above — the design is wrong at the foundation. Will reopen with a different approach (likely updating `ExchangeExec.equivalence_properties()` to reflect post-shuffle reality on resolve) once direction is settled.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(aqe): dynamic switching from streaming agg to hash aggregation [datafusion-ballista]
Posted by "wirybeaver (via GitHub)" <gi...@apache.org>.
wirybeaver commented on PR #1722:
URL: https://github.com/apache/datafusion-ballista/pull/1722#issuecomment-4505764269
@milenkovicm Thanks for pushing back — your questions exposed that my reasoning was wrong. Let me walk through what I now see.
**The flawed premise.** I assumed `ExchangeExec::equivalence_properties()` would reflect post-shuffle ordering once the exchange resolves. It doesn't. From `ballista/scheduler/src/state/aqe/execution_plan.rs:118`:
```rust
let eq_properties = input.properties().eq_properties.clone();
```
The eq_properties are captured at construction time from the **pre-shuffle input** and never updated when `resolve_shuffle_partitions()` is called.
**Why the rule is a no-op (or worse).**
1. **Linear → Sorted can't usefully fire.** If the agg's input advertises sort properties, those were visible at the agg's construction time too — `try_new_with_schema` derived the right mode then. The rule's `with_new_children` re-derivation produces the same mode.
2. **When it would fire, it would be incorrect.** A hash repartition wrapping a `SortExec` carries the sort eq_properties on the exchange, even though the actual post-shuffle data is no longer globally sorted (hash shuffle destroys global ordering). Switching the agg to `Sorted` based on those properties would cause streaming aggregation to produce wrong results.
3. **In valid replan scenarios it's redundant.** If a sibling rule changes the agg's subtree during `transform_up`, the framework already calls `with_new_children` on the parent, which re-derives the mode — no separate rule needed.
**What the roadmap line actually requires.**
For AQE to re-derive useful information after a stage resolves, the resolved exchange (or whatever node represents the post-shuffle reader) needs to advertise *post-shuffle* properties — driven by shuffle mode:
- Hash shuffle: destroys global ordering, only preserves partitioning on the hash key
- Sort-based shuffle: preserves per-partition ordering on the sort key
- Single-partition coalesce: preserves input ordering
Once those properties reflect reality, no dedicated re-derivation rule is needed — `with_new_children` during normal optimizer passes naturally picks up the change.
**Proposal.** Close this PR. I'd like to open a separate one (or RFC discussion) to update `ExchangeExec.equivalence_properties()` to translate properties on resolve based on shuffle mode. Happy to take direction from you on whether that's the right next step or if I'm still misreading the AQE design.
Apologies for the misfire — should have validated the premise with a runnable example before opening.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(aqe): dynamic switching from streaming agg to hash aggregation [datafusion-ballista]
Posted by "wirybeaver (via GitHub)" <gi...@apache.org>.
wirybeaver closed pull request #1722: feat(aqe): dynamic switching from streaming agg to hash aggregation
URL: https://github.com/apache/datafusion-ballista/pull/1722
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(aqe): dynamic switching from streaming agg to hash aggregation [datafusion-ballista]
Posted by "wirybeaver (via GitHub)" <gi...@apache.org>.
wirybeaver commented on code in PR #1722:
URL: https://github.com/apache/datafusion-ballista/pull/1722#discussion_r3279399269
##########
ballista/scheduler/src/state/aqe/optimizer_rule/dynamic_aggregate_algorithm.rs:
##########
@@ -0,0 +1,338 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::state::aqe::execution_plan::ExchangeExec;
+use ballista_core::config::BallistaConfig;
+use datafusion::common::tree_node::{Transformed, TreeNode};
+use datafusion::config::ConfigOptions;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::aggregates::AggregateExec;
+use log::info;
+use std::sync::Arc;
+
+/// AQE rule that re-derives `InputOrderMode` for each `AggregateExec` after a
+/// stage completes and rewrites the operator when the derived mode differs from
+/// the cached one.
+///
+/// DataFusion freezes `InputOrderMode` (Linear = hash table / Sorted = streaming /
+/// PartiallySorted) at plan time. After a shuffle stage resolves, the input's
+/// `EquivalenceProperties` may expose ordering that wasn't visible earlier, making
+/// streaming aggregation possible (or, after a subtree rewrite, making the previously
+/// assumed ordering stale). This rule repairs both cases on each AQE replan.
+///
+/// Relies on `AggregateExec::with_new_children` calling `try_new_with_schema`,
+/// which re-derives `input_order_mode` from the current input equivalence
+/// properties — no upstream DataFusion changes required.
+///
+/// Enabled via `ballista.aqe.dynamic_aggregate.enabled` (default: false).
+#[derive(Debug, Clone, Default)]
+pub struct DynamicAggregateAlgorithmRule {}
+
+impl DynamicAggregateAlgorithmRule {
+ fn transform(
+ plan: Arc<dyn ExecutionPlan>,
+ enabled: bool,
+ ) -> datafusion::error::Result<Transformed<Arc<dyn ExecutionPlan>>> {
+ if !enabled {
+ return Ok(Transformed::no(plan));
+ }
+ let Some(agg) = plan.as_any().downcast_ref::<AggregateExec>() else {
+ return Ok(Transformed::no(plan));
+ };
+ // Only act when a resolved exchange exists below — if nothing resolved,
+ // the ordering is the same as at initial plan time and no switch is needed.
+ if !subtree_has_resolved_exchange(&plan) {
+ return Ok(Transformed::no(plan));
+ }
+
+ // Re-derive InputOrderMode by reconstructing the aggregate with the same
+ // input. with_new_children → try_new_with_schema re-runs the derivation
+ // against the current input EquivalenceProperties.
+ let rebuilt = plan
+ .clone()
+ .with_new_children(vec![agg.children()[0].clone()])?;
+ let rebuilt_agg = rebuilt
+ .as_any()
+ .downcast_ref::<AggregateExec>()
+ .expect("with_new_children on AggregateExec must return AggregateExec");
+
+ if rebuilt_agg.input_order_mode() == agg.input_order_mode() {
+ return Ok(Transformed::no(plan));
+ }
+
+ info!(
+ "DynamicAggregateAlgorithmRule: switching InputOrderMode {:?} → {:?}",
+ agg.input_order_mode(),
+ rebuilt_agg.input_order_mode(),
+ );
+ Ok(Transformed::yes(rebuilt))
+ }
+}
+
+impl PhysicalOptimizerRule for DynamicAggregateAlgorithmRule {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
+ let enabled = config
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.aqe_dynamic_aggregate_enabled())
+ .unwrap_or(false);
+ Ok(plan.transform_up(|p| Self::transform(p, enabled))?.data)
+ }
+
+ fn name(&self) -> &str {
+ "DynamicAggregateAlgorithmRule"
+ }
+
+ fn schema_check(&self) -> bool {
+ false
+ }
+}
+
+/// Returns true if any `ExchangeExec` in `plan`'s subtree is resolved
+/// (i.e. its shuffle has completed). Recurses through all nodes, including
+/// through exchange boundaries — unlike `nearest_exchange_status` in
+/// `distributed_exchange.rs` which stops at the first exchange, here we
+/// want to know if *any* upstream stage has finished.
+fn subtree_has_resolved_exchange(plan: &Arc<dyn ExecutionPlan>) -> bool {
+ if let Some(exchange) = plan.as_any().downcast_ref::<ExchangeExec>() {
+ exchange.shuffle_created() && !exchange.inactive_stage
Review Comment:
You're right — the doc contradicts the implementation. The code stops at the first `ExchangeExec` and returns based on its resolved state, which is what we actually want (the nearest exchange gates the agg's input; a resolved exchange further down behind an unresolved one doesn't help). The doc should describe the actual behavior.
That said, please see my reply to @milenkovicm on the main thread — the broader design has problems that may make this PR moot.
##########
ballista/scheduler/src/state/aqe/optimizer_rule/dynamic_aggregate_algorithm.rs:
##########
@@ -0,0 +1,338 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::state::aqe::execution_plan::ExchangeExec;
+use ballista_core::config::BallistaConfig;
+use datafusion::common::tree_node::{Transformed, TreeNode};
+use datafusion::config::ConfigOptions;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::aggregates::AggregateExec;
+use log::info;
+use std::sync::Arc;
+
+/// AQE rule that re-derives `InputOrderMode` for each `AggregateExec` after a
+/// stage completes and rewrites the operator when the derived mode differs from
+/// the cached one.
+///
+/// DataFusion freezes `InputOrderMode` (Linear = hash table / Sorted = streaming /
+/// PartiallySorted) at plan time. After a shuffle stage resolves, the input's
+/// `EquivalenceProperties` may expose ordering that wasn't visible earlier, making
+/// streaming aggregation possible (or, after a subtree rewrite, making the previously
+/// assumed ordering stale). This rule repairs both cases on each AQE replan.
+///
+/// Relies on `AggregateExec::with_new_children` calling `try_new_with_schema`,
+/// which re-derives `input_order_mode` from the current input equivalence
+/// properties — no upstream DataFusion changes required.
+///
+/// Enabled via `ballista.aqe.dynamic_aggregate.enabled` (default: false).
+#[derive(Debug, Clone, Default)]
+pub struct DynamicAggregateAlgorithmRule {}
+
+impl DynamicAggregateAlgorithmRule {
+ fn transform(
+ plan: Arc<dyn ExecutionPlan>,
+ enabled: bool,
+ ) -> datafusion::error::Result<Transformed<Arc<dyn ExecutionPlan>>> {
+ if !enabled {
+ return Ok(Transformed::no(plan));
+ }
+ let Some(agg) = plan.as_any().downcast_ref::<AggregateExec>() else {
+ return Ok(Transformed::no(plan));
+ };
+ // Only act when a resolved exchange exists below — if nothing resolved,
+ // the ordering is the same as at initial plan time and no switch is needed.
+ if !subtree_has_resolved_exchange(&plan) {
+ return Ok(Transformed::no(plan));
+ }
+
+ // Re-derive InputOrderMode by reconstructing the aggregate with the same
+ // input. with_new_children → try_new_with_schema re-runs the derivation
+ // against the current input EquivalenceProperties.
+ let rebuilt = plan
+ .clone()
+ .with_new_children(vec![agg.children()[0].clone()])?;
+ let rebuilt_agg = rebuilt
+ .as_any()
+ .downcast_ref::<AggregateExec>()
+ .expect("with_new_children on AggregateExec must return AggregateExec");
+
+ if rebuilt_agg.input_order_mode() == agg.input_order_mode() {
+ return Ok(Transformed::no(plan));
+ }
+
+ info!(
+ "DynamicAggregateAlgorithmRule: switching InputOrderMode {:?} → {:?}",
+ agg.input_order_mode(),
+ rebuilt_agg.input_order_mode(),
+ );
+ Ok(Transformed::yes(rebuilt))
+ }
+}
+
+impl PhysicalOptimizerRule for DynamicAggregateAlgorithmRule {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
+ let enabled = config
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.aqe_dynamic_aggregate_enabled())
+ .unwrap_or(false);
+ Ok(plan.transform_up(|p| Self::transform(p, enabled))?.data)
+ }
+
+ fn name(&self) -> &str {
+ "DynamicAggregateAlgorithmRule"
+ }
+
+ fn schema_check(&self) -> bool {
+ false
+ }
+}
+
+/// Returns true if any `ExchangeExec` in `plan`'s subtree is resolved
+/// (i.e. its shuffle has completed). Recurses through all nodes, including
+/// through exchange boundaries — unlike `nearest_exchange_status` in
+/// `distributed_exchange.rs` which stops at the first exchange, here we
+/// want to know if *any* upstream stage has finished.
+fn subtree_has_resolved_exchange(plan: &Arc<dyn ExecutionPlan>) -> bool {
+ if let Some(exchange) = plan.as_any().downcast_ref::<ExchangeExec>() {
+ exchange.shuffle_created() && !exchange.inactive_stage
+ } else {
+ plan.children()
+ .iter()
+ .any(|child| subtree_has_resolved_exchange(child))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use ballista_core::config::BallistaConfig;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::common::{ColumnStatistics, Statistics};
+ use datafusion::config::{ConfigOptions, ExtensionOptions};
+ use datafusion::physical_expr::LexOrdering;
+ use datafusion::physical_expr::expressions::Column;
+ use datafusion::physical_plan::ExecutionPlan;
+ use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
+ use datafusion::physical_plan::sorts::sort::SortExec;
+ use datafusion::physical_plan::test::exec::StatisticsExec;
+ use datafusion::physical_expr::PhysicalSortExpr;
+ use std::sync::Arc;
+
+ fn leaf_unsorted() -> Arc<dyn ExecutionPlan> {
+ Arc::new(StatisticsExec::new(
+ Statistics {
+ num_rows: Default::default(),
+ total_byte_size: Default::default(),
+ column_statistics: vec![
+ ColumnStatistics::new_unknown(),
+ ColumnStatistics::new_unknown(),
+ ],
+ },
+ Schema::new(vec![
+ Field::new("k", DataType::Int32, true),
+ Field::new("v", DataType::Int64, true),
+ ]),
+ ))
+ }
+
+ fn leaf_sorted_on_k() -> Arc<dyn ExecutionPlan> {
+ let sort_expr =
+ PhysicalSortExpr::new_default(Arc::new(Column::new("k", 0)));
+ let ordering = LexOrdering::new(vec![sort_expr]).unwrap();
+ Arc::new(SortExec::new(ordering, leaf_unsorted()))
+ }
+
+ fn resolved_exchange(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
+ let exchange = ExchangeExec::new(input, None, 0);
+ exchange.resolve_shuffle_partitions(vec![]);
+ Arc::new(exchange)
+ }
+
+ fn unresolved_exchange(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
+ Arc::new(ExchangeExec::new(input, None, 0))
+ }
+
+ fn aggregate_on_k(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
+ let schema = input.schema();
+ let k_col = Arc::new(Column::new("k", 0)) as _;
+ let group_by = PhysicalGroupBy::new_single(vec![(k_col, "k".to_string())]);
+ Arc::new(
+ AggregateExec::try_new(
+ AggregateMode::Single,
+ group_by,
+ vec![],
+ vec![],
+ input,
+ schema,
+ )
+ .unwrap(),
+ )
+ }
+
+ fn config_with_flag(enabled: bool) -> ConfigOptions {
+ let mut ballista = BallistaConfig::default();
+ if enabled {
+ ballista
+ .set(
+ "aqe.dynamic_aggregate.enabled",
+ "true",
+ )
+ .unwrap();
+ }
+ let mut opts = ConfigOptions::new();
+ opts.extensions.insert(ballista);
+ opts
+ }
+
+ // --- gate ---
+
+ #[test]
+ fn gate_disabled_returns_no_transform() {
+ let plan = aggregate_on_k(resolved_exchange(leaf_sorted_on_k()));
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan.clone(), &config_with_flag(false)).unwrap();
+ assert!(
+ Arc::ptr_eq(&result, &plan),
+ "gate disabled: plan must be returned unchanged (same Arc)"
+ );
+ }
+
+ // --- subtree_has_resolved_exchange ---
+
+ #[test]
+ fn no_exchange_skips() {
+ // AggregateExec with no exchange at all → no rewrite even when enabled
+ let plan = aggregate_on_k(leaf_sorted_on_k());
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan.clone(), &config_with_flag(true)).unwrap();
+ assert!(Arc::ptr_eq(&result, &plan));
+ }
+
+ #[test]
+ fn unresolved_exchange_skips() {
+ let plan = aggregate_on_k(unresolved_exchange(leaf_sorted_on_k()));
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan.clone(), &config_with_flag(true)).unwrap();
+ assert!(Arc::ptr_eq(&result, &plan));
+ }
+
+ // --- mode transitions ---
+
+ #[test]
+ fn linear_to_sorted_when_input_fully_ordered() {
+ // Input: SortExec on k → resolved ExchangeExec → AggregateExec(group_by=[k])
+ // The SortExec grants ordering on k; the exchange preserves it in eq_properties.
+ // AggregateExec is initially planned as Linear (before the sort was visible).
+ // After re-derivation the mode should become Sorted.
+ let sorted_input = leaf_sorted_on_k();
+ let exchange = resolved_exchange(sorted_input);
+ let agg = aggregate_on_k(exchange);
+
+ let agg_exec = agg.as_any().downcast_ref::<AggregateExec>().unwrap();
+ // The agg was built over the resolved exchange whose eq_properties mirror
+ // the SortExec — so it may already be Sorted. What matters is the rule
+ // doesn't break it (idempotence when already correct).
+ let original_mode = agg_exec.input_order_mode().clone();
+
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(agg.clone(), &config_with_flag(true)).unwrap();
+ let result_agg = result.as_any().downcast_ref::<AggregateExec>().unwrap();
+
+ // After the rule the mode must match what re-derivation produces.
+ assert_eq!(result_agg.input_order_mode(), &original_mode);
Review Comment:
Right — the assertion is missing, but more importantly the test as written can't exercise a Linear→Sorted transition. `aggregate_on_k(exchange)` calls `AggregateExec::try_new`, which derives `input_order_mode` from the exchange's eq_properties at construction time. Since the exchange wraps a `SortExec(k)` and propagates its sort ordering, the agg is derived as `Sorted` immediately. The rule then re-derives the same mode, producing no transition — so the test passes trivially and doesn't validate what its name claims. Full root-cause analysis in my reply to @milenkovicm.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat(aqe): dynamic switching from streaming agg to hash aggregation [datafusion-ballista]
Posted by "wirybeaver (via GitHub)" <gi...@apache.org>.
wirybeaver commented on code in PR #1722:
URL: https://github.com/apache/datafusion-ballista/pull/1722#discussion_r3279401365
##########
ballista/scheduler/src/state/aqe/optimizer_rule/dynamic_aggregate_algorithm.rs:
##########
@@ -0,0 +1,338 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::state::aqe::execution_plan::ExchangeExec;
+use ballista_core::config::BallistaConfig;
+use datafusion::common::tree_node::{Transformed, TreeNode};
+use datafusion::config::ConfigOptions;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::aggregates::AggregateExec;
+use log::info;
+use std::sync::Arc;
+
+/// AQE rule that re-derives `InputOrderMode` for each `AggregateExec` after a
+/// stage completes and rewrites the operator when the derived mode differs from
+/// the cached one.
+///
+/// DataFusion freezes `InputOrderMode` (Linear = hash table / Sorted = streaming /
+/// PartiallySorted) at plan time. After a shuffle stage resolves, the input's
+/// `EquivalenceProperties` may expose ordering that wasn't visible earlier, making
+/// streaming aggregation possible (or, after a subtree rewrite, making the previously
+/// assumed ordering stale). This rule repairs both cases on each AQE replan.
+///
+/// Relies on `AggregateExec::with_new_children` calling `try_new_with_schema`,
+/// which re-derives `input_order_mode` from the current input equivalence
+/// properties — no upstream DataFusion changes required.
+///
+/// Enabled via `ballista.aqe.dynamic_aggregate.enabled` (default: false).
+#[derive(Debug, Clone, Default)]
+pub struct DynamicAggregateAlgorithmRule {}
+
+impl DynamicAggregateAlgorithmRule {
+ fn transform(
+ plan: Arc<dyn ExecutionPlan>,
+ enabled: bool,
+ ) -> datafusion::error::Result<Transformed<Arc<dyn ExecutionPlan>>> {
+ if !enabled {
+ return Ok(Transformed::no(plan));
+ }
+ let Some(agg) = plan.as_any().downcast_ref::<AggregateExec>() else {
+ return Ok(Transformed::no(plan));
+ };
+ // Only act when a resolved exchange exists below — if nothing resolved,
+ // the ordering is the same as at initial plan time and no switch is needed.
+ if !subtree_has_resolved_exchange(&plan) {
+ return Ok(Transformed::no(plan));
+ }
+
+ // Re-derive InputOrderMode by reconstructing the aggregate with the same
+ // input. with_new_children → try_new_with_schema re-runs the derivation
+ // against the current input EquivalenceProperties.
+ let rebuilt = plan
+ .clone()
+ .with_new_children(vec![agg.children()[0].clone()])?;
+ let rebuilt_agg = rebuilt
+ .as_any()
+ .downcast_ref::<AggregateExec>()
+ .expect("with_new_children on AggregateExec must return AggregateExec");
+
+ if rebuilt_agg.input_order_mode() == agg.input_order_mode() {
+ return Ok(Transformed::no(plan));
+ }
+
+ info!(
+ "DynamicAggregateAlgorithmRule: switching InputOrderMode {:?} → {:?}",
+ agg.input_order_mode(),
+ rebuilt_agg.input_order_mode(),
+ );
+ Ok(Transformed::yes(rebuilt))
+ }
+}
+
+impl PhysicalOptimizerRule for DynamicAggregateAlgorithmRule {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
+ let enabled = config
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.aqe_dynamic_aggregate_enabled())
+ .unwrap_or(false);
+ Ok(plan.transform_up(|p| Self::transform(p, enabled))?.data)
+ }
+
+ fn name(&self) -> &str {
+ "DynamicAggregateAlgorithmRule"
+ }
+
+ fn schema_check(&self) -> bool {
+ false
+ }
+}
+
+/// Returns true if any `ExchangeExec` in `plan`'s subtree is resolved
+/// (i.e. its shuffle has completed). Recurses through all nodes, including
+/// through exchange boundaries — unlike `nearest_exchange_status` in
+/// `distributed_exchange.rs` which stops at the first exchange, here we
+/// want to know if *any* upstream stage has finished.
+fn subtree_has_resolved_exchange(plan: &Arc<dyn ExecutionPlan>) -> bool {
+ if let Some(exchange) = plan.as_any().downcast_ref::<ExchangeExec>() {
+ exchange.shuffle_created() && !exchange.inactive_stage
+ } else {
+ plan.children()
+ .iter()
+ .any(|child| subtree_has_resolved_exchange(child))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use ballista_core::config::BallistaConfig;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::common::{ColumnStatistics, Statistics};
+ use datafusion::config::{ConfigOptions, ExtensionOptions};
+ use datafusion::physical_expr::LexOrdering;
+ use datafusion::physical_expr::expressions::Column;
+ use datafusion::physical_plan::ExecutionPlan;
+ use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
+ use datafusion::physical_plan::sorts::sort::SortExec;
+ use datafusion::physical_plan::test::exec::StatisticsExec;
+ use datafusion::physical_expr::PhysicalSortExpr;
+ use std::sync::Arc;
+
+ fn leaf_unsorted() -> Arc<dyn ExecutionPlan> {
+ Arc::new(StatisticsExec::new(
+ Statistics {
+ num_rows: Default::default(),
+ total_byte_size: Default::default(),
+ column_statistics: vec![
+ ColumnStatistics::new_unknown(),
+ ColumnStatistics::new_unknown(),
+ ],
+ },
+ Schema::new(vec![
+ Field::new("k", DataType::Int32, true),
+ Field::new("v", DataType::Int64, true),
+ ]),
+ ))
+ }
+
+ fn leaf_sorted_on_k() -> Arc<dyn ExecutionPlan> {
+ let sort_expr =
+ PhysicalSortExpr::new_default(Arc::new(Column::new("k", 0)));
+ let ordering = LexOrdering::new(vec![sort_expr]).unwrap();
+ Arc::new(SortExec::new(ordering, leaf_unsorted()))
+ }
+
+ fn resolved_exchange(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
+ let exchange = ExchangeExec::new(input, None, 0);
+ exchange.resolve_shuffle_partitions(vec![]);
+ Arc::new(exchange)
+ }
+
+ fn unresolved_exchange(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
+ Arc::new(ExchangeExec::new(input, None, 0))
+ }
+
+ fn aggregate_on_k(input: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
+ let schema = input.schema();
+ let k_col = Arc::new(Column::new("k", 0)) as _;
+ let group_by = PhysicalGroupBy::new_single(vec![(k_col, "k".to_string())]);
+ Arc::new(
+ AggregateExec::try_new(
+ AggregateMode::Single,
+ group_by,
+ vec![],
+ vec![],
+ input,
+ schema,
+ )
+ .unwrap(),
+ )
+ }
+
+ fn config_with_flag(enabled: bool) -> ConfigOptions {
+ let mut ballista = BallistaConfig::default();
+ if enabled {
+ ballista
+ .set(
+ "aqe.dynamic_aggregate.enabled",
+ "true",
+ )
+ .unwrap();
+ }
+ let mut opts = ConfigOptions::new();
+ opts.extensions.insert(ballista);
+ opts
+ }
+
+ // --- gate ---
+
+ #[test]
+ fn gate_disabled_returns_no_transform() {
+ let plan = aggregate_on_k(resolved_exchange(leaf_sorted_on_k()));
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan.clone(), &config_with_flag(false)).unwrap();
+ assert!(
+ Arc::ptr_eq(&result, &plan),
+ "gate disabled: plan must be returned unchanged (same Arc)"
+ );
+ }
+
+ // --- subtree_has_resolved_exchange ---
+
+ #[test]
+ fn no_exchange_skips() {
+ // AggregateExec with no exchange at all → no rewrite even when enabled
+ let plan = aggregate_on_k(leaf_sorted_on_k());
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan.clone(), &config_with_flag(true)).unwrap();
+ assert!(Arc::ptr_eq(&result, &plan));
+ }
+
+ #[test]
+ fn unresolved_exchange_skips() {
+ let plan = aggregate_on_k(unresolved_exchange(leaf_sorted_on_k()));
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan.clone(), &config_with_flag(true)).unwrap();
+ assert!(Arc::ptr_eq(&result, &plan));
+ }
+
+ // --- mode transitions ---
+
+ #[test]
+ fn linear_to_sorted_when_input_fully_ordered() {
+ // Input: SortExec on k → resolved ExchangeExec → AggregateExec(group_by=[k])
+ // The SortExec grants ordering on k; the exchange preserves it in eq_properties.
+ // AggregateExec is initially planned as Linear (before the sort was visible).
+ // After re-derivation the mode should become Sorted.
+ let sorted_input = leaf_sorted_on_k();
+ let exchange = resolved_exchange(sorted_input);
+ let agg = aggregate_on_k(exchange);
+
+ let agg_exec = agg.as_any().downcast_ref::<AggregateExec>().unwrap();
+ // The agg was built over the resolved exchange whose eq_properties mirror
+ // the SortExec — so it may already be Sorted. What matters is the rule
+ // doesn't break it (idempotence when already correct).
+ let original_mode = agg_exec.input_order_mode().clone();
+
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(agg.clone(), &config_with_flag(true)).unwrap();
+ let result_agg = result.as_any().downcast_ref::<AggregateExec>().unwrap();
+
+ // After the rule the mode must match what re-derivation produces.
+ assert_eq!(result_agg.input_order_mode(), &original_mode);
+ // Schema must be preserved.
+ assert_eq!(result.schema(), agg.schema());
+ }
+
+ #[test]
+ fn unsorted_input_stays_linear() {
+ // Input has no ordering → mode must stay Linear after the rule.
+ let exchange = resolved_exchange(leaf_unsorted());
+ let plan = aggregate_on_k(exchange);
+
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan.clone(), &config_with_flag(true)).unwrap();
+ let result_agg = result.as_any().downcast_ref::<AggregateExec>().unwrap();
+
+ assert_eq!(
+ result_agg.input_order_mode(),
+ &datafusion::physical_plan::InputOrderMode::Linear
+ );
+ }
+
+ #[test]
+ fn schema_preserved_after_rewrite() {
+ let exchange = resolved_exchange(leaf_sorted_on_k());
+ let plan = aggregate_on_k(exchange);
+ let original_schema = plan.schema();
+
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let result = rule.optimize(plan, &config_with_flag(true)).unwrap();
+
+ assert_eq!(result.schema(), original_schema);
+ }
+
+ // --- idempotence ---
+
+ #[test]
+ fn idempotent_second_pass() {
+ let exchange = resolved_exchange(leaf_sorted_on_k());
+ let plan = aggregate_on_k(exchange);
+
+ let rule = DynamicAggregateAlgorithmRule::default();
+ let after_first = rule.optimize(plan, &config_with_flag(true)).unwrap();
+ let mode_after_first = after_first
+ .as_any()
+ .downcast_ref::<AggregateExec>()
+ .unwrap()
+ .input_order_mode()
+ .clone();
+
+ let after_second = rule
+ .optimize(after_first.clone(), &config_with_flag(true))
+ .unwrap();
+ let mode_after_second = after_second
+ .as_any()
+ .downcast_ref::<AggregateExec>()
+ .unwrap()
+ .input_order_mode()
+ .clone();
+
+ assert_eq!(mode_after_first, mode_after_second, "rule must be idempotent");
+ }
+
+ // --- rule metadata ---
+
+ #[test]
+ fn rule_name() {
+ assert_eq!(
+ DynamicAggregateAlgorithmRule::default().name(),
+ "DynamicAggregateAlgorithmRule"
+ );
+ }
+
+ #[test]
+ fn schema_check_is_false() {
+ assert!(!DynamicAggregateAlgorithmRule::default().schema_check());
+ }
Review Comment:
Confirmed — and this gap isn't just test coverage. `with_new_children` re-derivation produces the same mode the constructor would, because `ExchangeExec.equivalence_properties()` doesn't change between construction and resolution. So the rule can't, by design, produce a Linear↔Sorted transition. Details in my reply to @milenkovicm.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org