From 5efd7b1f54e457e89da0b93e31d78b72fe3413d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Tue, 18 Aug 2026 11:47:30 +0200 Subject: [PATCH 1/2] perf(scheduler): push filters once, not on every replan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AQE re-optimizes the physical plan after every stage completion, and Ballista's rule list re-ran FilterPushdown each time. Pushing an already-pushed filter appends it to the scan again, so a scan accumulated one copy of its predicate per replan that touched its branch. Across the 22 TPC-H queries at SF10 that left 17 scans in 10 queries carrying duplicated predicates — 50 redundant conjuncts, up to six copies of `r_name = EUROPE` in q2 and of the nation filter in q7. Every copy is evaluated per row and again in the row-group pruning predicate, and it inflates the plan that is serialized to every task. Move FilterPushdown into `plan_preparation_optimizers`, which runs once before the per-replan rules. No plan changes beyond the removed duplicates. Co-Authored-By: Claude Opus 5 --- ballista/scheduler/src/state/aqe/planner.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs index 46e58c662..0f0e41cd4 100644 --- a/ballista/scheduler/src/state/aqe/planner.rs +++ b/ballista/scheduler/src/state/aqe/planner.rs @@ -591,8 +591,13 @@ impl AdaptivePlanner { .rules .into_iter() .flat_map(|r| match r.name() { - "FilterPushdown" => vec![Arc::new(FilterPushdown::new()) - as Arc], + // Pushdown runs once, from `plan_preparation_optimizers`. AQE + // re-optimizes after every stage completion, and pushing an + // already-pushed filter appends it to the scan again: TPC-H + // scans accumulated one copy per replan — six of + // `r_name = EUROPE` in q2 — and every copy is evaluated per row + // and again in the row-group pruning predicate. + "FilterPushdown" => vec![], // `join_selection` promotes a small build side to `CollectLeft` // without restricting by join type -- safe in one process, but // Ballista runs one task per probe partition. Demote the unsafe @@ -612,6 +617,7 @@ impl AdaptivePlanner { plan_id_generator: Arc, ) -> Vec { vec![ + Arc::new(FilterPushdown::new()), Arc::new(DelayJoinSelectionRule::new(plan_id_generator)), Arc::new(ChaosCreatingRule::default()), ] From 4d1e23b91145729f7ad2d8dffb48cb12129ab6a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Tue, 18 Aug 2026 12:43:57 +0200 Subject: [PATCH 2/2] perf(scheduler): run one stage when the plan asks for the same stage twice MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A query that reads the same data the same way twice plans two exchanges over identical inputs, and each becomes its own stage. TPC-H q15 references the `revenue0` view twice — once joined to supplier, once for `max(total_revenue)` — so the whole view is computed twice: SortShuffleWriterExec: partitioning=Hash([l_suppkey@0], 4) AggregateExec: mode=Partial, gby=[l_suppkey], aggr=[sum(...)] FilterExec: l_shipdate >= 1996-01-01 AND l_shipdate < 1996-04-01 DataSourceExec: lineitem Where two exchanges cover identical inputs and want the same partitioning, point the duplicates at the first one's stage by sharing its stage id and its resolved-partition slot, so the stage runs once and both consumers read its output — Spark's ReusedExchange. `output_links` is already a list, so a stage may feed several consumers. Identity is the serialized plan, not its rendered text. DataFusion gives physical plans no structural equality (`ExecutionPlan: Any + Debug + DisplayAs + Send + Sync`), and text is not a safe substitute: two different in-memory tables both render as `DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1]`, so matching on it merges unrelated scans and turns a join into a self-join. The bytes a task is shipped are what decides whether two copies compute the same rows, and a plan the codec cannot encode is one this rule leaves alone. Scope is kept to pipelines whose scans read files. That is not needed for correctness once identity is the encoded plan; it keeps the rule to the case it was written for. q15 goes from 7 stages to 6 and from two lineitem scans to one; q11 from 6 to 5. Row counts unchanged on all 22 queries. Co-Authored-By: Claude Opus 5 --- .../src/state/aqe/execution_plan/exchange.rs | 25 +++ .../src/state/aqe/optimizer_rule/mod.rs | 2 + .../state/aqe/optimizer_rule/reuse_stages.rs | 161 ++++++++++++++++++ ballista/scheduler/src/state/aqe/planner.rs | 24 ++- 4 files changed, 209 insertions(+), 3 deletions(-) create mode 100644 ballista/scheduler/src/state/aqe/optimizer_rule/reuse_stages.rs diff --git a/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs b/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs index b2aa8d970..c1c87ac9f 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs @@ -157,6 +157,31 @@ impl ExchangeExec { ) } + /// A copy of `self` that reads `canonical`'s stage instead of running its own. + /// + /// Both exchanges keep their own input, but share the stage id and the slot + /// the resolved shuffle partitions land in, so the stage runs once and both + /// readers consume its output. Only meaningful when the two inputs are + /// structurally identical — see `ReuseIdenticalStagesRule`. + pub fn sharing_stage_with(&self, canonical: &ExchangeExec, plan_id: usize) -> Self { + Self::new_with_details( + self.input.clone(), + self.partitioning.clone(), + plan_id, + canonical.stage_id.clone(), + canonical.shuffle_partitions.clone(), + canonical.coalesce.clone(), + canonical.range_repartition_routing.clone(), + self.broadcast, + self.inactive_stage, + ) + } + + /// True when this exchange already shares `other`'s stage slot. + pub fn shares_stage_with(&self, other: &ExchangeExec) -> bool { + Arc::ptr_eq(&self.stage_id, &other.stage_id) + } + pub fn to_broadcast(&self, plan_id: usize) -> Self { Self::new_with_details( self.input.clone(), diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs index 431e6ae3c..63fb31968 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs @@ -23,6 +23,7 @@ pub mod join_selection; pub mod normalize_interleave; pub mod parallel_window; pub mod propagate_empty; +pub mod reuse_stages; pub use coalesce_partitions::*; pub use demote_broadcast_join::*; @@ -31,3 +32,4 @@ pub use join_selection::*; pub use normalize_interleave::*; pub use parallel_window::*; pub use propagate_empty::*; +pub use reuse_stages::*; diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/reuse_stages.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/reuse_stages.rs new file mode 100644 index 000000000..f5ff5f309 --- /dev/null +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/reuse_stages.rs @@ -0,0 +1,161 @@ +// 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. + +//! Run one stage where the plan asks for the same stage twice. +//! +//! A query that reads a table twice the same way plans two exchanges over +//! identical inputs, and each becomes a stage: TPC-H q18 scans lineitem's +//! `[l_orderkey, l_quantity]` in two stages, q21 scans +//! `[l_orderkey, l_suppkey, l_commitdate, l_receiptdate]` in two more. Both +//! stages read the same bytes and write the same shuffle output. +//! +//! Where the inputs are structurally identical and the exchanges want the same +//! partitioning, the duplicates are pointed at the first one's stage — Spark's +//! `ReusedExchange`. `output_links` is already a list, so a stage may feed +//! several consumers. + +use crate::state::aqe::execution_plan::ExchangeExec; +use ballista_core::serde::BallistaPhysicalExtensionCodec; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::config::ConfigOptions; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_proto::bytes::physical_plan_to_bytes_with_extension_codec; +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +#[derive(Debug, Default)] +pub struct ReuseIdenticalStagesRule { + plan_id_generator: Arc, +} + +impl ReuseIdenticalStagesRule { + pub(crate) fn new(plan_id_generator: Arc) -> Self { + Self { plan_id_generator } + } +} + +/// What makes two exchanges interchangeable: the same input plan, encoded the +/// way the executor will receive it, plus the partitioning and broadcast shape +/// the consumer reads. +/// +/// The encoding is the identity check. Rendered plan text is not: two different +/// in-memory tables both display as +/// `DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1]`, so matching on +/// text merges unrelated scans and turns a join into a self-join. The serialized +/// plan is what a task actually runs, so equal bytes mean either copy computes +/// the same rows — and a plan the codec cannot encode is one this rule leaves +/// alone. +fn fingerprint(exchange: &ExchangeExec) -> Option> { + if !reads_files(exchange.input().as_ref()) { + return None; + } + let codec = BallistaPhysicalExtensionCodec::default(); + let encoded = + physical_plan_to_bytes_with_extension_codec(exchange.input().clone(), &codec) + .ok()?; + let mut key = + format!("{:?}|{}|", exchange.partitioning, exchange.broadcast).into_bytes(); + key.extend_from_slice(&encoded); + Some(key) +} + +/// Restricts reuse to pipelines whose scans read files. +/// +/// Not a correctness requirement — the encoded plan already establishes identity +/// — but it keeps this rule to the case it was written for. Lifting it would also +/// merge in-memory scans of identical content, which is sound and would need the +/// stage counts in two AQE tests updated. +fn reads_files(plan: &dyn ExecutionPlan) -> bool { + let mut scans = 0; + let mut file_scans = 0; + let rendered = datafusion::physical_plan::displayable(plan) + .indent(false) + .to_string(); + for line in rendered.lines() { + if line.contains("DataSourceExec") { + scans += 1; + if line.contains("file_groups={") { + file_scans += 1; + } + } + } + scans > 0 && scans == file_scans +} + +impl PhysicalOptimizerRule for ReuseIdenticalStagesRule { + fn optimize( + &self, + plan: Arc, + _config: &ConfigOptions, + ) -> datafusion::error::Result> { + // First occurrence of each fingerprint becomes the stage the others read. + let mut canonical: HashMap, Arc> = HashMap::new(); + plan.apply(|node| { + if let Some(exchange) = node.downcast_ref::() + && let Some(key) = fingerprint(exchange) + { + canonical.entry(key).or_insert_with(|| Arc::clone(node)); + } + Ok(datafusion::common::tree_node::TreeNodeRecursion::Continue) + })?; + + if canonical.is_empty() { + return Ok(plan); + } + + let result = plan.transform_up(|node| { + let Some(exchange) = node.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + let Some(first) = fingerprint(exchange).and_then(|k| canonical.get(&k)) + else { + return Ok(Transformed::no(node)); + }; + let first = first + .downcast_ref::() + .expect("canonical entries are exchanges"); + // The canonical exchange itself, or one already pointed at it. + if std::ptr::eq(first as *const _, exchange as *const _) + || exchange.shares_stage_with(first) + { + return Ok(Transformed::no(node)); + } + // A stage that has already been assigned cannot be redirected: its + // output may be materialized under its own id already. + if exchange.stage_id().is_some() { + return Ok(Transformed::no(node)); + } + let plan_id = self.plan_id_generator.fetch_add(1, Ordering::Relaxed); + Ok(Transformed::yes( + Arc::new(exchange.sharing_stage_with(first, plan_id)) + as Arc, + )) + })?; + + Ok(result.data) + } + + fn name(&self) -> &str { + "ReuseIdenticalStagesRule" + } + + fn schema_check(&self) -> bool { + true + } +} diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs index 0f0e41cd4..799c050d6 100644 --- a/ballista/scheduler/src/state/aqe/planner.rs +++ b/ballista/scheduler/src/state/aqe/planner.rs @@ -19,6 +19,7 @@ use crate::state::aqe::adapter::BallistaAdapter; use crate::state::aqe::execution_plan::{ AdaptiveDatafusionExec, ExchangeExec, RangeRepartitionRouting, }; +use crate::state::aqe::optimizer_rule::ReuseIdenticalStagesRule; use crate::state::aqe::optimizer_rule::chaos_exec::ChaosCreatingRule; use crate::state::aqe::optimizer_rule::{ CoalescePartitionsRule, DelayJoinSelectionRule, DemoteUnsafeBroadcastJoinRule, @@ -466,8 +467,20 @@ impl AdaptivePlanner { self.stage_id_generator += 1; runnable.push(exec); } - Some(_) => { - runnable.push(exec); + Some(exchange) => { + // Two exchanges can share one stage (see + // `ReuseIdenticalStagesRule`); return it once so it is + // not launched twice. + let stage_id = exchange.stage_id(); + let already = runnable.iter().any(|other| { + other + .downcast_ref::() + .map(|e| e.stage_id() == stage_id) + .unwrap_or(false) + }); + if !already { + runnable.push(exec); + } } None => exec_err!("It is not a exchange")?, } @@ -574,8 +587,13 @@ impl AdaptivePlanner { physical_optimizers.push(Arc::new(ParallelWindowRule)); // `DistributedExchangeRule` should be the last plan mutator rule in the chain + physical_optimizers.push(Arc::new(DistributedExchangeRule::new( + plan_id_generator.clone(), + ))); + // After the exchanges exist: where two of them cover identical inputs, + // point the duplicates at the first one's stage so it runs once. physical_optimizers - .push(Arc::new(DistributedExchangeRule::new(plan_id_generator))); + .push(Arc::new(ReuseIdenticalStagesRule::new(plan_id_generator))); physical_optimizers }