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 46e58c662..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 } @@ -591,8 +609,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 +635,7 @@ impl AdaptivePlanner { plan_id_generator: Arc, ) -> Vec { vec![ + Arc::new(FilterPushdown::new()), Arc::new(DelayJoinSelectionRule::new(plan_id_generator)), Arc::new(ChaosCreatingRule::default()), ]