comphead commented on code in PR #2765:
URL:
https://github.com/apache/datafusion-ballista/pull/2765#discussion_r4469975303
##########
ballista/core/src/config.rs:
##########
@@ -87,6 +87,10 @@ pub const BALLISTA_CLIENT_IO_RETRY_WAIT_TIME_MS: &str =
"ballista.client.io_retry_wait_time_ms";
/// Enables adaptive query planning
pub const BALLISTA_ADAPTIVE_PLANNER_ENABLED: &str =
"ballista.planner.adaptive.enabled";
+/// Lets the adaptive planner run a repeated subplan's stage once and have
+/// every consumer read the same shuffle output (Spark's `ReuseExchange`).
+pub const BALLISTA_ADAPTIVE_REUSE_EXCHANGE: &str =
+ "ballista.planner.adaptive.reuse_exchange";
Review Comment:
Naming: every other AQE rule toggle follows
`ballista.planner.<rule>.enabled`, for example `coalesce`, `skew_join`,
`propagate_empty`, `parallel_window` and `adaptive_join`.
`ballista.planner.reuse_exchange.enabled`, with
`BALLISTA_REUSE_EXCHANGE_ENABLED` and `reuse_exchange_enabled()`, would match.
Renaming is cheap now and harder after a release.
##########
ballista/scheduler/src/state/aqe/optimizer_rule/reuse_exchange.rs:
##########
@@ -0,0 +1,468 @@
+// 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.
+
+//! Runs a repeated subplan's stage once (Spark's `ReuseExchange`).
+//!
+//! A query that reads one subplan twice, such as a common table expression
+//! joined to itself or TPC-H q15's revenue view, plans two identical
+//! [`ExchangeExec`]s. Without this rule each becomes its own stage and the
+//! subplan runs twice. This rule finds exchanges whose subtrees are identical
+//! and makes every one of them a reader of a single stage: the stage runs
+//! once, and each consumer reads the same shuffle output within the job.
+//!
+//! # Matching
+//!
+//! An exchange's key is the protobuf encoding of its subtree under the job's
+//! physical extension codec, the encoding that ships stages to executors, so
+//! two subtrees match only when an executor could not tell them apart. Inside
+//! a subtree a nested exchange is encoded by what it writes, its partitioning
+//! and whether it is a broadcast, never by its stage id. A subtree the codec
+//! cannot encode gets no key and is left alone.
+//!
+//! # Sharing
+//!
+//! A duplicate is rewritten with [`ExchangeExec::share_stage_of`]: it takes
+//! the canonical exchange's input and write-side state, and keeps its own
+//! reader-side state (coalesce and skew-join decisions), since each consuming
+//! stage decides those for itself. The canonical exchange of a key is one
+//! whose shuffle is already resolved, else one already running, else the
+//! first met bottom-up. An exchange already running a stage of its own is
+//! never re-pointed at another.
+//!
+//! # What is never merged
+//!
+//! - A subtree with a volatile expression such as `random()`: it encodes like
+//! its twin but has to produce values of its own.
+//! - Range repartitioning and the window rewrites built on it. The scheduler
+//! finishes those stages by reading state off the one consumer of the
+//! stage, so a second consumer is not supported.
+//!
+//! The rule runs last in the replan chain, after an embedder's rules, so keys
+//! cover each stage as it will run.
+
+use std::collections::{HashMap, HashSet};
+use std::fmt::{self, Write};
+use std::hash::{DefaultHasher, Hash, Hasher};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+
+use ballista_core::config::BallistaConfig;
+use ballista_core::execution_plans::{
+ OrderedRangeRepartitionExec, PartitionedBoundedWindowAggExec,
PerPartitionFilterExec,
+ PrefixMergeExec, RangeFilterExec, RuntimeStatsExec,
UnorderedRangeRepartitionExec,
+};
+use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
+use datafusion::config::ConfigOptions;
+use datafusion::error::Result;
+use datafusion::execution::TaskContext;
+use datafusion::logical_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF,
WindowUDF};
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_expr_common::physical_expr::is_volatile;
+use
datafusion::physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::{ExecutionPlan, Partitioning, displayable};
+use datafusion_proto::bytes::physical_plan_to_bytes_with_extension_codec;
+use datafusion_proto::physical_plan::PhysicalExtensionCodec;
+use datafusion_proto::physical_plan::PhysicalProtoConverterExtension;
+use datafusion_proto::physical_plan::to_proto::serialize_partitioning;
+use log::debug;
+use prost::Message;
+
+use crate::state::aqe::execution_plan::ExchangeExec;
+
+/// Makes every exchange in a set of identical exchanges a reader of one
+/// stage. See the module docs.
+#[derive(Debug)]
+pub struct ReuseExchangeRule {
+ codec: Arc<dyn PhysicalExtensionCodec>,
+ read_finished_stage: Arc<AtomicBool>,
+}
+
+impl ReuseExchangeRule {
+ /// A rule that matches subtrees by their encoding under `codec`, the job's
+ /// physical extension codec. It raises `read_finished_stage` whenever it
+ /// points a pending exchange at a stage that has already finished: that
+ /// resolves the exchange after the rules that react to a resolved exchange
+ /// have run, so the caller must run the optimizer chain again.
+ pub fn new(
+ codec: Arc<dyn PhysicalExtensionCodec>,
+ read_finished_stage: Arc<AtomicBool>,
+ ) -> Self {
+ Self {
+ codec,
+ read_finished_stage,
+ }
+ }
+
+ /// The key `node` is matched on, or `None` when `node` is not an exchange
+ /// that may be merged.
+ fn key(&self, node: &Arc<dyn ExecutionPlan>) -> Option<Vec<u8>> {
+ let exchange = node.downcast_ref::<ExchangeExec>()?;
+ if matches!(exchange.partitioning(), Some(Partitioning::Range(_)))
+ || !mergeable(node).ok()?
+ {
+ return None;
+ }
+ let codec = KeyCodec {
+ inner: Arc::clone(&self.codec),
+ };
+ physical_plan_to_bytes_with_extension_codec(Arc::clone(node), &codec)
+ .ok()
+ .map(Vec::from)
+ }
+}
+
+/// How far along an exchange's stage is, ordering the candidates for the
+/// canonical exchange of a key.
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+enum Progress {
+ Pending,
+ Running,
+ Resolved,
+}
+
+fn progress(exchange: &ExchangeExec) -> Progress {
+ if exchange.shuffle_created() {
+ Progress::Resolved
+ } else if exchange.stage_id().is_some() {
+ Progress::Running
+ } else {
+ Progress::Pending
+ }
+}
+
+/// The exchanges of one key: how many there are, and the one furthest along.
+struct Group {
+ exchanges: usize,
+ furthest: Arc<dyn ExecutionPlan>,
+}
+
+impl PhysicalOptimizerRule for ReuseExchangeRule {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let enabled = config
+ .extensions
+ .get::<BallistaConfig>()
+ .map(BallistaConfig::adaptive_reuse_exchange_enabled)
+ .unwrap_or_else(|| {
+ BallistaConfig::default().adaptive_reuse_exchange_enabled()
+ });
+ if !enabled {
+ return Ok(plan);
+ }
+
+ // Encoding a subtree is costly, and most plans repeat nothing. Only
+ // exchanges a cheap fingerprint can't tell apart, or that already
+ // share a stage handle, are keyed.
+ let mut exchanges = Vec::new();
+ fingerprints(&plan, &mut exchanges);
+ let mut fingerprint_counts: HashMap<u64, usize> = HashMap::new();
+ let mut handle_counts: HashMap<*const (), usize> = HashMap::new();
+ for (node, fingerprint) in &exchanges {
+ *fingerprint_counts.entry(*fingerprint).or_default() += 1;
+ *handle_counts
+ .entry(as_exchange(node).stage_handle())
+ .or_default() += 1;
+ }
+ let candidates = exchanges.iter().filter(|(node, fingerprint)| {
+ fingerprint_counts[fingerprint] > 1
+ || handle_counts[&as_exchange(node).stage_handle()] > 1
+ });
+
+ // Key each candidate, and find for each key the exchange furthest
+ // along, which the others will read from. Keys are interned so the
+ // second pass looks them up by exchange instead of encoding again;
+ // an exchange's plan id and stage handle survive the rewrites below.
+ let mut key_ids: HashMap<Vec<u8>, usize> = HashMap::new();
+ let mut groups: Vec<Group> = Vec::new();
+ let mut exchange_keys: HashMap<(usize, *const ()), usize> =
HashMap::new();
+ let mut handle_key: HashMap<*const (), usize> = HashMap::new();
+ // Pending stage handles that exchanges of different subtrees share,
+ // as `to_broadcast` or an embedder's rewrite of one reader can leave
+ // behind. Neither may write for the other.
+ let mut conflicted: HashSet<*const ()> = HashSet::new();
+ for (node, _) in candidates {
+ let Some(key) = self.key(node) else {
Review Comment:
A reader that gets no key is skipped before its handle is recorded, so a
divergence between a keyed reader and an unkeyed one is never detected.
`ExchangeExec::with_new_children` keeps the shared `stage_id` Arc, so any later
rewrite of one reader's subtree (an embedder rule is the case the comment above
mentions) keeps the two readers on one stage. If that rewrite makes the subtree
unkeyable, through a volatile expression or a node the codec can't encode,
nothing splits them. Whichever reader claims the stage first writes the output
for both.
I reproduced this locally. Build a pending shared stage, rewrite one reader
with `with_new_children` over a projection containing `random()`, then run the
rule on `Union(hash, diverged)`. Both readers still `shares_stage_with` each
other afterwards.
Suggest recording the handle of every candidate before the `continue`, for
example with a sentinel group id for unkeyed exchanges, so a pending handle
seen with an unkeyed reader is added to `conflicted` and split.
##########
ballista/scheduler/src/state/aqe/planner.rs:
##########
@@ -135,9 +146,31 @@ impl AdaptivePlanner {
runnable_stage_cache: HashMap::new(),
job_id,
runnable_stage_output: HashMap::new(),
+ read_finished_stage,
})
}
+ /// Runs the optimizer chain over `plan`, and again for as long as a pass
+ /// resolves an exchange by pointing it at a finished stage, as a stage
+ /// finishing does. Bounded, since each pass can only point exchanges at
+ /// stages already finished.
+ fn optimize(
+ planner: &DefaultPhysicalPlanner,
+ session_state: &SessionState,
+ read_finished_stage: &AtomicBool,
+ mut plan: Arc<dyn ExecutionPlan>,
+ ) -> common::Result<Arc<dyn ExecutionPlan>> {
+ const MAX_PASSES: usize = 8;
Review Comment:
The comment says the loop is bounded because each pass can only point
exchanges at stages that already finished. If that holds, `MAX_PASSES` is a
safety net, and reaching it means the invariant broke. Right now it exits
silently with the flag still set, which is the state that let a stage launch
with an undecided join (q25, q64, q74). Returning `internal_err!`, or at least
a `warn!`, when the cap is hit would surface it.
Separately, the `Arc<AtomicBool>` threaded from the rule into the planner
couples a stateless `PhysicalOptimizerRule` to its caller. A rule-agnostic
check could replace it: count the resolved exchanges before and after
`optimize_physical_plan` and run again if the count changed.
`CountResolvedExchanges` in the tests already does exactly this. Stage
completions happen outside the chain, so only a mid-chain resolution can change
the count.
##########
ballista/scheduler/src/state/aqe/optimizer_rule/reuse_exchange.rs:
##########
@@ -0,0 +1,468 @@
+// 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.
+
+//! Runs a repeated subplan's stage once (Spark's `ReuseExchange`).
+//!
+//! A query that reads one subplan twice, such as a common table expression
+//! joined to itself or TPC-H q15's revenue view, plans two identical
+//! [`ExchangeExec`]s. Without this rule each becomes its own stage and the
+//! subplan runs twice. This rule finds exchanges whose subtrees are identical
+//! and makes every one of them a reader of a single stage: the stage runs
+//! once, and each consumer reads the same shuffle output within the job.
+//!
+//! # Matching
+//!
+//! An exchange's key is the protobuf encoding of its subtree under the job's
+//! physical extension codec, the encoding that ships stages to executors, so
+//! two subtrees match only when an executor could not tell them apart. Inside
+//! a subtree a nested exchange is encoded by what it writes, its partitioning
+//! and whether it is a broadcast, never by its stage id. A subtree the codec
+//! cannot encode gets no key and is left alone.
+//!
+//! # Sharing
+//!
+//! A duplicate is rewritten with [`ExchangeExec::share_stage_of`]: it takes
+//! the canonical exchange's input and write-side state, and keeps its own
+//! reader-side state (coalesce and skew-join decisions), since each consuming
+//! stage decides those for itself. The canonical exchange of a key is one
+//! whose shuffle is already resolved, else one already running, else the
+//! first met bottom-up. An exchange already running a stage of its own is
+//! never re-pointed at another.
+//!
+//! # What is never merged
+//!
+//! - A subtree with a volatile expression such as `random()`: it encodes like
+//! its twin but has to produce values of its own.
+//! - Range repartitioning and the window rewrites built on it. The scheduler
+//! finishes those stages by reading state off the one consumer of the
+//! stage, so a second consumer is not supported.
+//!
+//! The rule runs last in the replan chain, after an embedder's rules, so keys
+//! cover each stage as it will run.
+
+use std::collections::{HashMap, HashSet};
+use std::fmt::{self, Write};
+use std::hash::{DefaultHasher, Hash, Hasher};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+
+use ballista_core::config::BallistaConfig;
+use ballista_core::execution_plans::{
+ OrderedRangeRepartitionExec, PartitionedBoundedWindowAggExec,
PerPartitionFilterExec,
+ PrefixMergeExec, RangeFilterExec, RuntimeStatsExec,
UnorderedRangeRepartitionExec,
+};
+use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
+use datafusion::config::ConfigOptions;
+use datafusion::error::Result;
+use datafusion::execution::TaskContext;
+use datafusion::logical_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF,
WindowUDF};
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_expr_common::physical_expr::is_volatile;
+use
datafusion::physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::{ExecutionPlan, Partitioning, displayable};
+use datafusion_proto::bytes::physical_plan_to_bytes_with_extension_codec;
+use datafusion_proto::physical_plan::PhysicalExtensionCodec;
+use datafusion_proto::physical_plan::PhysicalProtoConverterExtension;
+use datafusion_proto::physical_plan::to_proto::serialize_partitioning;
+use log::debug;
+use prost::Message;
+
+use crate::state::aqe::execution_plan::ExchangeExec;
+
+/// Makes every exchange in a set of identical exchanges a reader of one
+/// stage. See the module docs.
+#[derive(Debug)]
+pub struct ReuseExchangeRule {
+ codec: Arc<dyn PhysicalExtensionCodec>,
+ read_finished_stage: Arc<AtomicBool>,
+}
+
+impl ReuseExchangeRule {
+ /// A rule that matches subtrees by their encoding under `codec`, the job's
+ /// physical extension codec. It raises `read_finished_stage` whenever it
+ /// points a pending exchange at a stage that has already finished: that
+ /// resolves the exchange after the rules that react to a resolved exchange
+ /// have run, so the caller must run the optimizer chain again.
+ pub fn new(
+ codec: Arc<dyn PhysicalExtensionCodec>,
+ read_finished_stage: Arc<AtomicBool>,
+ ) -> Self {
+ Self {
+ codec,
+ read_finished_stage,
+ }
+ }
+
+ /// The key `node` is matched on, or `None` when `node` is not an exchange
+ /// that may be merged.
+ fn key(&self, node: &Arc<dyn ExecutionPlan>) -> Option<Vec<u8>> {
+ let exchange = node.downcast_ref::<ExchangeExec>()?;
+ if matches!(exchange.partitioning(), Some(Partitioning::Range(_)))
+ || !mergeable(node).ok()?
Review Comment:
`mergeable` walks the whole subtree for each candidate, and `key` then
encodes that same subtree again. Nested candidates make this quadratic in
nesting depth, and it runs on every replan. `fingerprints` already visits every
node once, bottom-up. Returning `(fingerprint, mergeable)` from it, and
skipping unmergeable exchanges there, would fold the volatility and
single-consumer check into the existing single pass.
##########
ballista/scheduler/src/state/aqe/planner.rs:
##########
@@ -521,6 +569,14 @@ impl AdaptivePlanner {
let mut runnable_stages = Vec::new();
Self::find_runnable_exchanges(&self.plan, &mut runnable_stages);
+ // Every reader of a shared stage (see `ReuseExchangeRule`) is found,
+ // and the stage runs once. A held-back reader doesn't claim it.
+ let mut stages = HashSet::new();
+ runnable_stages.retain(|exec| {
Review Comment:
This adds a second pass over `runnable_stages` ahead of the loop below,
which checks `inactive_stage` again. Folding the dedupe into that loop keeps it
in one pass and in one place, for example `Some(exchange) if
!claimed.insert(exchange.stage_handle()) => continue,` after the
`inactive_stage` arm. `HashSet::with_capacity(runnable_stages.len())` would
bound the set.
##########
ballista/client/tests/reuse_exchange.rs:
##########
@@ -0,0 +1,159 @@
+// 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.
+
+mod common;
+
+/// End-to-end tests of `ballista.planner.adaptive.reuse_exchange`. Queries
+/// that read one subplan more than once run under adaptive query planning
+/// with exchange reuse on and off, and must return the same rows as a run of
+/// the static planner.
+#[cfg(test)]
+#[cfg(feature = "standalone")]
+mod reuse_exchange {
+ use ballista::prelude::SessionContextExt;
+ use ballista_core::config::{
+ BALLISTA_ADAPTIVE_PLANNER_ENABLED, BALLISTA_ADAPTIVE_REUSE_EXCHANGE,
+ BALLISTA_STANDALONE_PARALLELISM,
+ };
+ use ballista_core::extension::SessionConfigExt;
+ use datafusion::arrow::util::pretty::pretty_format_batches;
+ use datafusion::execution::{SessionState, SessionStateBuilder};
+ use datafusion::prelude::{CsvReadOptions, SessionConfig, SessionContext};
+ use rstest::rstest;
+ use std::io::Write;
+ use std::path::Path;
+ use tempfile::TempDir;
+
+ /// Writes an `orders` table split over four files, so its scan has four
+ /// partitions, and a small `customers` table.
+ fn write_tables(dir: &Path) {
+ let orders = dir.join("orders");
+ std::fs::create_dir(&orders).unwrap();
+ for file in 0..4 {
+ let mut out =
+
std::fs::File::create(orders.join(format!("part-{file}.csv"))).unwrap();
+ writeln!(out, "o_id,o_cust,o_amount").unwrap();
+ for i in 0..250 {
+ let id = file * 250 + i;
+ writeln!(out, "{id},{},{}", id % 97, (id * 7) % 1000).unwrap();
+ }
+ }
+ let mut out =
std::fs::File::create(dir.join("customers.csv")).unwrap();
+ writeln!(out, "c_id,c_name,c_region").unwrap();
+ for id in 0..100 {
+ writeln!(out, "{id},customer-{id},region-{}", id % 5).unwrap();
+ }
+ }
+
+ fn state(options: &[(&str, &str)]) -> SessionState {
+ let mut config = SessionConfig::new_with_ballista()
+ .with_target_partitions(4)
+ .set_str(BALLISTA_STANDALONE_PARALLELISM, "4");
+ for (key, value) in options {
+ config = config.set_str(key, value);
+ }
+ SessionStateBuilder::new()
+ .with_config(config)
+ .with_default_features()
+ .build()
+ }
+
+ async fn context(dir: &Path, options: &[(&str, &str)]) -> SessionContext {
+ let ctx = SessionContext::standalone_with_state(state(options))
+ .await
+ .unwrap();
+ ctx.register_csv(
+ "orders",
+ dir.join("orders").to_str().unwrap(),
+ CsvReadOptions::new(),
+ )
+ .await
+ .unwrap();
+ ctx.register_csv(
+ "customers",
+ dir.join("customers.csv").to_str().unwrap(),
+ CsvReadOptions::new(),
+ )
+ .await
+ .unwrap();
+ ctx
+ }
+
+ async fn rows(ctx: &SessionContext, sql: &str) -> Vec<String> {
+ let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap();
+ pretty_format_batches(&batches)
+ .unwrap()
+ .to_string()
+ .lines()
+ .map(str::to_owned)
+ .collect()
+ }
+
+ /// Each query reads the same aggregate, and so the same exchanges, more
+ /// than once.
+ const QUERIES: &[&str] = &[
+ // A common table expression joined to itself.
+ "WITH t AS (SELECT o_cust, count(*) AS n, sum(o_amount) AS total \
+ FROM orders GROUP BY o_cust) \
+ SELECT a.o_cust, a.total, b.n \
+ FROM t a JOIN t b ON a.o_cust = b.o_cust + 1 \
+ ORDER BY a.o_cust",
+ // A union of an aggregate with itself.
+ "WITH t AS (SELECT o_cust, sum(o_amount) AS total \
+ FROM orders GROUP BY o_cust) \
+ SELECT o_cust, total FROM t \
+ UNION ALL SELECT o_cust, total FROM t \
+ ORDER BY o_cust, total",
+ // An aggregate compared with its own average, as TPC-H q15 does.
Review Comment:
Nit: the query compares with `max(total)`, not the average.
##########
ballista/scheduler/src/state/aqe/test/reuse_exchange.rs:
##########
@@ -0,0 +1,322 @@
+// 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.
+
+//! A subplan read more than once runs its stage once (`ReuseExchangeRule`).
+
+use crate::state::aqe::AdaptiveExecutionGraph;
+use crate::state::aqe::execution_plan::ExchangeExec;
+use crate::state::aqe::optimizer_rule::ReuseExchangeRule;
+use crate::state::aqe::planner::AdaptivePlanner;
+use crate::state::aqe::test::{
+ mock_batch, mock_context, mock_partitions_with_statistics,
+};
+use crate::state::execution_graph::ExecutionGraph;
+use crate::state::execution_stage::ExecutionStage;
+use crate::test_utils::{
+ mock_executor, revive_graph_and_complete_next_stage_with_executor,
+};
+use ballista_core::config::BallistaConfig;
+use ballista_core::extension::SessionConfigExt;
+use ballista_core::serde::BallistaPhysicalExtensionCodec;
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
+use datafusion::config::ConfigOptions;
+use datafusion::execution::SessionStateBuilder;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::empty::EmptyExec;
+use datafusion::physical_plan::union::UnionExec;
+use datafusion::physical_plan::{Partitioning, PhysicalExpr,
expressions::Column};
+use datafusion::prelude::SessionConfig;
+use datafusion::prelude::SessionContext;
+use parking_lot::Mutex;
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+
+/// A common table expression read twice by a union.
+const UNION: &str = "with v as (select c, sum(a) as s from t group by c) \
+ select c, s from v union all select c, s from v";
+/// A common table expression joined to itself.
+const SELF_JOIN: &str = "with v as (select c, sum(a) as s from t group by c) \
+ select x.c, x.s, y.s from v x join v y on x.c = y.c";
+/// An aggregate compared with its own maximum, as TPC-H q15 does.
+const Q15_SHAPE: &str = "with v as (select c, sum(a) as s from t group by c) \
+ select c, s from v where s = (select max(s) from v)";
+
+fn context(reuse: bool) -> SessionContext {
+ let config = mock_context()
+ .copied_config()
+ .with_ballista_adaptive_reuse_exchange_enabled(reuse);
+ let ctx = SessionContext::new_with_config(config);
+ ctx.register_batch("t", mock_batch().unwrap()).unwrap();
+ ctx
+}
+
+async fn planner(ctx: &SessionContext, sql: &str) -> AdaptivePlanner {
+ let plan = ctx
+ .sql(sql)
+ .await
+ .unwrap()
+ .create_physical_plan()
+ .await
+ .unwrap();
+ AdaptivePlanner::try_from_plan(ctx.state().config(), plan,
"job".into()).unwrap()
+}
+
+async fn graph(ctx: &SessionContext, sql: &str) -> AdaptiveExecutionGraph {
+ let plan = ctx.sql(sql).await.unwrap().logical_plan().clone();
+ AdaptiveExecutionGraph::try_new("scheduler", &"job".into(), "", ctx,
&plan, 0, None)
+ .await
+ .unwrap()
+}
+
+/// Runs every stage of `graph` on `executor` and returns how many stages ran.
+fn run_to_completion(graph: &mut AdaptiveExecutionGraph, executor: &str) ->
usize {
+ let executor = mock_executor(executor.to_owned());
+ let mut stages = 0;
+ while !graph.is_successful() {
+ revive_graph_and_complete_next_stage_with_executor(graph,
&executor).unwrap();
+ stages += 1;
+ }
+ stages
+}
+
+#[tokio::test]
+async fn a_repeated_aggregate_runs_once() {
Review Comment:
`a_repeated_aggregate_runs_once` and `reuse_can_be_switched_off` run the
same three queries and differ only in the flag and the expected count. One
`#[rstest]` with `#[values(true, false)] reuse` asserting `if reuse { 1 } else
{ 2 }` would replace both.
`a_job_reading_an_aggregate_twice_runs_one_stage_fewer` already compares the
two arms at graph level, too.
Missing coverage: `ballista.planner.coalesce.enabled` defaults to `false`,
so no test exercises per-reader coalesce or skew decisions on a shared stage
end to end. A coalesce-enabled (and skew-enabled) arm in the standalone test
would cover the reader-side state that `share_stage_of` deliberately keeps
separate.
##########
ballista/scheduler/src/state/aqe/execution_plan/exchange.rs:
##########
@@ -815,3 +867,63 @@ mod column_statistics_tests {
}
}
}
+
+#[cfg(test)]
+mod shared_stage_tests {
+ //! Two readers of one stage, as exchange reuse makes them.
+ use super::*;
+ use crate::state::aqe::test::mock_partitions_with_statistics;
+ use ballista_core::execution_plans::PartitionGroup;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::physical_plan::empty::EmptyExec;
+
+ fn exchange(plan_id: usize) -> ExchangeExec {
+ let schema = Arc::new(Schema::new(vec![Field::new("a",
DataType::Int32, true)]));
+ ExchangeExec::new(
+ Arc::new(EmptyExec::new(schema)),
+ Some(Partitioning::UnknownPartitioning(2)),
+ plan_id,
+ )
+ }
+
+ #[test]
+ fn a_reader_of_a_shared_stage_sees_the_stage_resolve() {
Review Comment:
This is covered by
`a_repeat_found_after_the_stage_finished_reads_its_output` and the graph tests,
which assert the same resolution through the rule. Likewise,
`an_exchange_with_its_own_stage_shares_nothing` is covered by
`readers_of_a_pending_stage_that_diverged_are_split`.
`each_reader_of_a_shared_stage_coalesces_on_its_own` is the one with unique
coverage here.
##########
ballista/scheduler/src/state/aqe/optimizer_rule/reuse_exchange.rs:
##########
@@ -0,0 +1,468 @@
+// 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.
+
+//! Runs a repeated subplan's stage once (Spark's `ReuseExchange`).
+//!
+//! A query that reads one subplan twice, such as a common table expression
+//! joined to itself or TPC-H q15's revenue view, plans two identical
+//! [`ExchangeExec`]s. Without this rule each becomes its own stage and the
+//! subplan runs twice. This rule finds exchanges whose subtrees are identical
+//! and makes every one of them a reader of a single stage: the stage runs
+//! once, and each consumer reads the same shuffle output within the job.
+//!
+//! # Matching
+//!
+//! An exchange's key is the protobuf encoding of its subtree under the job's
+//! physical extension codec, the encoding that ships stages to executors, so
+//! two subtrees match only when an executor could not tell them apart. Inside
+//! a subtree a nested exchange is encoded by what it writes, its partitioning
+//! and whether it is a broadcast, never by its stage id. A subtree the codec
+//! cannot encode gets no key and is left alone.
+//!
+//! # Sharing
+//!
+//! A duplicate is rewritten with [`ExchangeExec::share_stage_of`]: it takes
+//! the canonical exchange's input and write-side state, and keeps its own
+//! reader-side state (coalesce and skew-join decisions), since each consuming
+//! stage decides those for itself. The canonical exchange of a key is one
+//! whose shuffle is already resolved, else one already running, else the
+//! first met bottom-up. An exchange already running a stage of its own is
+//! never re-pointed at another.
+//!
+//! # What is never merged
+//!
+//! - A subtree with a volatile expression such as `random()`: it encodes like
+//! its twin but has to produce values of its own.
+//! - Range repartitioning and the window rewrites built on it. The scheduler
+//! finishes those stages by reading state off the one consumer of the
+//! stage, so a second consumer is not supported.
+//!
+//! The rule runs last in the replan chain, after an embedder's rules, so keys
+//! cover each stage as it will run.
+
+use std::collections::{HashMap, HashSet};
+use std::fmt::{self, Write};
+use std::hash::{DefaultHasher, Hash, Hasher};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+
+use ballista_core::config::BallistaConfig;
+use ballista_core::execution_plans::{
+ OrderedRangeRepartitionExec, PartitionedBoundedWindowAggExec,
PerPartitionFilterExec,
+ PrefixMergeExec, RangeFilterExec, RuntimeStatsExec,
UnorderedRangeRepartitionExec,
+};
+use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
+use datafusion::config::ConfigOptions;
+use datafusion::error::Result;
+use datafusion::execution::TaskContext;
+use datafusion::logical_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF,
WindowUDF};
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_expr_common::physical_expr::is_volatile;
+use
datafusion::physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::{ExecutionPlan, Partitioning, displayable};
+use datafusion_proto::bytes::physical_plan_to_bytes_with_extension_codec;
+use datafusion_proto::physical_plan::PhysicalExtensionCodec;
+use datafusion_proto::physical_plan::PhysicalProtoConverterExtension;
+use datafusion_proto::physical_plan::to_proto::serialize_partitioning;
+use log::debug;
+use prost::Message;
+
+use crate::state::aqe::execution_plan::ExchangeExec;
+
+/// Makes every exchange in a set of identical exchanges a reader of one
+/// stage. See the module docs.
+#[derive(Debug)]
+pub struct ReuseExchangeRule {
+ codec: Arc<dyn PhysicalExtensionCodec>,
+ read_finished_stage: Arc<AtomicBool>,
+}
+
+impl ReuseExchangeRule {
+ /// A rule that matches subtrees by their encoding under `codec`, the job's
+ /// physical extension codec. It raises `read_finished_stage` whenever it
+ /// points a pending exchange at a stage that has already finished: that
+ /// resolves the exchange after the rules that react to a resolved exchange
+ /// have run, so the caller must run the optimizer chain again.
+ pub fn new(
+ codec: Arc<dyn PhysicalExtensionCodec>,
+ read_finished_stage: Arc<AtomicBool>,
+ ) -> Self {
+ Self {
+ codec,
+ read_finished_stage,
+ }
+ }
+
+ /// The key `node` is matched on, or `None` when `node` is not an exchange
+ /// that may be merged.
+ fn key(&self, node: &Arc<dyn ExecutionPlan>) -> Option<Vec<u8>> {
+ let exchange = node.downcast_ref::<ExchangeExec>()?;
+ if matches!(exchange.partitioning(), Some(Partitioning::Range(_)))
+ || !mergeable(node).ok()?
+ {
+ return None;
+ }
+ let codec = KeyCodec {
+ inner: Arc::clone(&self.codec),
+ };
+ physical_plan_to_bytes_with_extension_codec(Arc::clone(node), &codec)
+ .ok()
+ .map(Vec::from)
+ }
+}
+
+/// How far along an exchange's stage is, ordering the candidates for the
+/// canonical exchange of a key.
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+enum Progress {
+ Pending,
+ Running,
+ Resolved,
+}
+
+fn progress(exchange: &ExchangeExec) -> Progress {
+ if exchange.shuffle_created() {
+ Progress::Resolved
+ } else if exchange.stage_id().is_some() {
+ Progress::Running
+ } else {
+ Progress::Pending
+ }
+}
+
+/// The exchanges of one key: how many there are, and the one furthest along.
+struct Group {
+ exchanges: usize,
+ furthest: Arc<dyn ExecutionPlan>,
+}
+
+impl PhysicalOptimizerRule for ReuseExchangeRule {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let enabled = config
+ .extensions
+ .get::<BallistaConfig>()
+ .map(BallistaConfig::adaptive_reuse_exchange_enabled)
+ .unwrap_or_else(|| {
+ BallistaConfig::default().adaptive_reuse_exchange_enabled()
+ });
+ if !enabled {
+ return Ok(plan);
+ }
+
+ // Encoding a subtree is costly, and most plans repeat nothing. Only
+ // exchanges a cheap fingerprint can't tell apart, or that already
+ // share a stage handle, are keyed.
+ let mut exchanges = Vec::new();
+ fingerprints(&plan, &mut exchanges);
+ let mut fingerprint_counts: HashMap<u64, usize> = HashMap::new();
Review Comment:
Nit: both maps hold at most `exchanges.len()` entries, so
`HashMap::with_capacity(exchanges.len())` avoids rehashing. The same bound
applies to `key_ids`, `groups`, `exchange_keys` and `handle_key` below.
##########
ballista/scheduler/src/state/aqe/optimizer_rule/reuse_exchange.rs:
##########
@@ -0,0 +1,468 @@
+// 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.
+
+//! Runs a repeated subplan's stage once (Spark's `ReuseExchange`).
+//!
+//! A query that reads one subplan twice, such as a common table expression
+//! joined to itself or TPC-H q15's revenue view, plans two identical
+//! [`ExchangeExec`]s. Without this rule each becomes its own stage and the
+//! subplan runs twice. This rule finds exchanges whose subtrees are identical
+//! and makes every one of them a reader of a single stage: the stage runs
+//! once, and each consumer reads the same shuffle output within the job.
+//!
+//! # Matching
+//!
+//! An exchange's key is the protobuf encoding of its subtree under the job's
+//! physical extension codec, the encoding that ships stages to executors, so
+//! two subtrees match only when an executor could not tell them apart. Inside
+//! a subtree a nested exchange is encoded by what it writes, its partitioning
+//! and whether it is a broadcast, never by its stage id. A subtree the codec
+//! cannot encode gets no key and is left alone.
+//!
+//! # Sharing
+//!
+//! A duplicate is rewritten with [`ExchangeExec::share_stage_of`]: it takes
+//! the canonical exchange's input and write-side state, and keeps its own
+//! reader-side state (coalesce and skew-join decisions), since each consuming
+//! stage decides those for itself. The canonical exchange of a key is one
+//! whose shuffle is already resolved, else one already running, else the
+//! first met bottom-up. An exchange already running a stage of its own is
+//! never re-pointed at another.
+//!
+//! # What is never merged
+//!
+//! - A subtree with a volatile expression such as `random()`: it encodes like
+//! its twin but has to produce values of its own.
+//! - Range repartitioning and the window rewrites built on it. The scheduler
+//! finishes those stages by reading state off the one consumer of the
+//! stage, so a second consumer is not supported.
+//!
+//! The rule runs last in the replan chain, after an embedder's rules, so keys
+//! cover each stage as it will run.
+
+use std::collections::{HashMap, HashSet};
+use std::fmt::{self, Write};
+use std::hash::{DefaultHasher, Hash, Hasher};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+
+use ballista_core::config::BallistaConfig;
+use ballista_core::execution_plans::{
+ OrderedRangeRepartitionExec, PartitionedBoundedWindowAggExec,
PerPartitionFilterExec,
+ PrefixMergeExec, RangeFilterExec, RuntimeStatsExec,
UnorderedRangeRepartitionExec,
+};
+use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
+use datafusion::config::ConfigOptions;
+use datafusion::error::Result;
+use datafusion::execution::TaskContext;
+use datafusion::logical_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF,
WindowUDF};
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_expr_common::physical_expr::is_volatile;
+use
datafusion::physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::{ExecutionPlan, Partitioning, displayable};
+use datafusion_proto::bytes::physical_plan_to_bytes_with_extension_codec;
+use datafusion_proto::physical_plan::PhysicalExtensionCodec;
+use datafusion_proto::physical_plan::PhysicalProtoConverterExtension;
+use datafusion_proto::physical_plan::to_proto::serialize_partitioning;
+use log::debug;
+use prost::Message;
+
+use crate::state::aqe::execution_plan::ExchangeExec;
+
+/// Makes every exchange in a set of identical exchanges a reader of one
+/// stage. See the module docs.
+#[derive(Debug)]
+pub struct ReuseExchangeRule {
+ codec: Arc<dyn PhysicalExtensionCodec>,
+ read_finished_stage: Arc<AtomicBool>,
+}
+
+impl ReuseExchangeRule {
+ /// A rule that matches subtrees by their encoding under `codec`, the job's
+ /// physical extension codec. It raises `read_finished_stage` whenever it
+ /// points a pending exchange at a stage that has already finished: that
+ /// resolves the exchange after the rules that react to a resolved exchange
+ /// have run, so the caller must run the optimizer chain again.
+ pub fn new(
+ codec: Arc<dyn PhysicalExtensionCodec>,
+ read_finished_stage: Arc<AtomicBool>,
+ ) -> Self {
+ Self {
+ codec,
+ read_finished_stage,
+ }
+ }
+
+ /// The key `node` is matched on, or `None` when `node` is not an exchange
+ /// that may be merged.
+ fn key(&self, node: &Arc<dyn ExecutionPlan>) -> Option<Vec<u8>> {
+ let exchange = node.downcast_ref::<ExchangeExec>()?;
+ if matches!(exchange.partitioning(), Some(Partitioning::Range(_)))
+ || !mergeable(node).ok()?
+ {
+ return None;
+ }
+ let codec = KeyCodec {
+ inner: Arc::clone(&self.codec),
+ };
+ physical_plan_to_bytes_with_extension_codec(Arc::clone(node), &codec)
+ .ok()
+ .map(Vec::from)
+ }
+}
+
+/// How far along an exchange's stage is, ordering the candidates for the
+/// canonical exchange of a key.
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+enum Progress {
+ Pending,
+ Running,
+ Resolved,
+}
+
+fn progress(exchange: &ExchangeExec) -> Progress {
+ if exchange.shuffle_created() {
+ Progress::Resolved
+ } else if exchange.stage_id().is_some() {
+ Progress::Running
+ } else {
+ Progress::Pending
+ }
+}
+
+/// The exchanges of one key: how many there are, and the one furthest along.
+struct Group {
+ exchanges: usize,
+ furthest: Arc<dyn ExecutionPlan>,
+}
+
+impl PhysicalOptimizerRule for ReuseExchangeRule {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let enabled = config
+ .extensions
+ .get::<BallistaConfig>()
+ .map(BallistaConfig::adaptive_reuse_exchange_enabled)
+ .unwrap_or_else(|| {
+ BallistaConfig::default().adaptive_reuse_exchange_enabled()
+ });
+ if !enabled {
+ return Ok(plan);
+ }
+
+ // Encoding a subtree is costly, and most plans repeat nothing. Only
+ // exchanges a cheap fingerprint can't tell apart, or that already
+ // share a stage handle, are keyed.
+ let mut exchanges = Vec::new();
+ fingerprints(&plan, &mut exchanges);
+ let mut fingerprint_counts: HashMap<u64, usize> = HashMap::new();
+ let mut handle_counts: HashMap<*const (), usize> = HashMap::new();
+ for (node, fingerprint) in &exchanges {
+ *fingerprint_counts.entry(*fingerprint).or_default() += 1;
+ *handle_counts
+ .entry(as_exchange(node).stage_handle())
+ .or_default() += 1;
+ }
+ let candidates = exchanges.iter().filter(|(node, fingerprint)| {
+ fingerprint_counts[fingerprint] > 1
+ || handle_counts[&as_exchange(node).stage_handle()] > 1
+ });
+
+ // Key each candidate, and find for each key the exchange furthest
+ // along, which the others will read from. Keys are interned so the
+ // second pass looks them up by exchange instead of encoding again;
+ // an exchange's plan id and stage handle survive the rewrites below.
+ let mut key_ids: HashMap<Vec<u8>, usize> = HashMap::new();
+ let mut groups: Vec<Group> = Vec::new();
+ let mut exchange_keys: HashMap<(usize, *const ()), usize> =
HashMap::new();
+ let mut handle_key: HashMap<*const (), usize> = HashMap::new();
+ // Pending stage handles that exchanges of different subtrees share,
+ // as `to_broadcast` or an embedder's rewrite of one reader can leave
+ // behind. Neither may write for the other.
+ let mut conflicted: HashSet<*const ()> = HashSet::new();
+ for (node, _) in candidates {
+ let Some(key) = self.key(node) else {
+ continue;
+ };
+ let exchange = as_exchange(node);
+ let next_id = groups.len();
+ let id = *key_ids.entry(key).or_insert(next_id);
+ if id == next_id {
+ groups.push(Group {
+ exchanges: 0,
+ furthest: Arc::clone(node),
+ });
+ }
+ let group = &mut groups[id];
+ group.exchanges += 1;
+ if progress(exchange) > progress(as_exchange(&group.furthest)) {
+ group.furthest = Arc::clone(node);
+ }
+ let handle = exchange.stage_handle();
+ exchange_keys.insert((exchange.plan_id(), handle), id);
+ if handle_key
+ .insert(handle, id)
+ .is_some_and(|other| other != id)
+ && progress(exchange) == Progress::Pending
+ {
+ conflicted.insert(handle);
+ }
+ }
+ if conflicted.is_empty() && groups.iter().all(|group| group.exchanges
< 2) {
+ return Ok(plan);
+ }
+
+ // Bottom-up, so a nested repeat is merged before the subtrees above.
+ let mut canonical: HashMap<usize, Arc<dyn ExecutionPlan>> =
HashMap::new();
+ let transformed = plan.transform_up(|node| {
+ let Some(exchange) = node.downcast_ref::<ExchangeExec>() else {
+ return Ok(Transformed::no(node));
+ };
+ let handle = exchange.stage_handle();
+ let Some(&id) = exchange_keys.get(&(exchange.plan_id(), handle))
else {
+ return Ok(Transformed::no(node));
+ };
+
+ let split = conflicted.contains(&handle);
+ let node: Arc<dyn ExecutionPlan> = if split {
+ Arc::new(exchange.with_own_stage())
+ } else {
+ node
+ };
+ let exchange = as_exchange(&node);
+
+ // The exchange furthest along is read from, even when the walk
+ // meets it later: its stage has started, so its subtree no longer
+ // changes what it writes. If nothing has run yet, the first met
is.
+ let target = Arc::clone(canonical.entry(id).or_insert_with(|| {
+ let furthest = &groups[id].furthest;
+ if progress(as_exchange(furthest)) == Progress::Pending
+ || exchange.shares_stage_with(as_exchange(furthest))
+ {
+ Arc::clone(&node)
+ } else {
+ Arc::clone(furthest)
+ }
+ }));
+ let target = as_exchange(&target);
+
+ // Already a reader of that stage, or running a stage of its own,
+ // which is never re-pointed.
+ if exchange.shares_stage_with(target)
+ || progress(exchange) != Progress::Pending
+ {
+ return Ok(Transformed::new_transformed(node, split));
+ }
+ debug!(
+ "exchange reuse: exchange {} reads the stage of exchange {}",
+ exchange.plan_id(),
+ target.plan_id()
+ );
+ if progress(target) == Progress::Resolved {
+ self.read_finished_stage.store(true, Ordering::Relaxed);
+ }
+ Ok(Transformed::yes(Arc::new(exchange.share_stage_of(target))))
+ })?;
+ Ok(transformed.data)
+ }
+
+ fn name(&self) -> &str {
+ "ReuseExchangeRule"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+/// Whether nothing in `plan` rules out sharing its stage: no volatile
+/// expression and no single-consumer operator.
+fn mergeable(plan: &Arc<dyn ExecutionPlan>) -> Result<bool> {
+ let excluded = plan.exists(|node| {
+ if requires_single_consumer(node.as_ref()) {
+ return Ok(true);
+ }
+ let mut volatile = false;
+ node.apply_expressions(&mut |expr| {
+ volatile = is_volatile(expr);
+ Ok(if volatile {
+ TreeNodeRecursion::Stop
+ } else {
+ TreeNodeRecursion::Continue
+ })
+ })?;
+ Ok(volatile)
+ })?;
+ Ok(!excluded)
+}
+
+/// Operators whose stage the scheduler finishes by reading state off the
+/// stage's single consumer. See the module docs.
+fn requires_single_consumer(plan: &dyn ExecutionPlan) -> bool {
+ plan.is::<OrderedRangeRepartitionExec>()
+ || plan.is::<UnorderedRangeRepartitionExec>()
+ || plan.is::<RangeFilterExec>()
+ || plan.is::<PerPartitionFilterExec>()
+ || plan.is::<PrefixMergeExec>()
+ || plan.is::<PartitionedBoundedWindowAggExec>()
+ || plan.is::<RuntimeStatsExec>()
+}
+
+/// Pushes every exchange in `plan` onto `exchanges`, children first, with a
+/// hash of its subtree that two subtrees with one key always share, though
+/// different subtrees may share it too. Returns the hash of `plan`. Each node
+/// contributes its one-line display, except that an exchange contributes only
+/// what it writes, so stage ids never reach the hash.
+fn fingerprints(
+ plan: &Arc<dyn ExecutionPlan>,
+ exchanges: &mut Vec<(Arc<dyn ExecutionPlan>, u64)>,
+) -> u64 {
+ let mut hasher = DefaultHasher::new();
+ match plan.downcast_ref::<ExchangeExec>() {
+ Some(exchange) => {
+ "ExchangeExec".hash(&mut hasher);
+ exchange.is_broadcast().hash(&mut hasher);
+ exchange
+ .partitioning()
+ .map(|partitioning| partitioning.to_string())
+ .hash(&mut hasher);
+ }
+ None => {
+ let _ = write!(
+ HashWriter(&mut hasher),
+ "{}",
+ displayable(plan.as_ref()).one_line()
+ );
+ }
+ }
+ for child in plan.children() {
+ fingerprints(child, exchanges).hash(&mut hasher);
+ }
+ let fingerprint = hasher.finish();
+ if plan.is::<ExchangeExec>() {
+ exchanges.push((Arc::clone(plan), fingerprint));
+ }
+ fingerprint
+}
+
+/// Feeds formatted text straight into a hasher.
+struct HashWriter<'a>(&'a mut DefaultHasher);
+
+impl fmt::Write for HashWriter<'_> {
+ fn write_str(&mut self, s: &str) -> fmt::Result {
+ self.0.write(s.as_bytes());
+ Ok(())
+ }
+}
+
+fn as_exchange(plan: &Arc<dyn ExecutionPlan>) -> &ExchangeExec {
+ plan.downcast_ref::<ExchangeExec>()
+ .expect("only exchanges are keyed")
+}
+
+/// Encodes a subtree for matching only, never for decoding. An
+/// [`ExchangeExec`] is encoded by what it writes, so a nested exchange's stage
+/// id never reaches the key; everything else goes to the job's codec.
+#[derive(Debug)]
+struct KeyCodec {
+ inner: Arc<dyn PhysicalExtensionCodec>,
+}
+
+/// Marks an exchange in a key, so it can't collide with another node's bytes.
+const EXCHANGE_TAG: &[u8] = b"ballista.aqe.ExchangeExec";
+
+impl PhysicalExtensionCodec for KeyCodec {
+ fn try_decode(
+ &self,
+ _buf: &[u8],
+ _inputs: &[Arc<dyn ExecutionPlan>],
+ _ctx: &TaskContext,
+ _proto_converter: &dyn PhysicalProtoConverterExtension,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ datafusion::common::internal_err!("exchange reuse keys are never
decoded")
+ }
+
+ fn try_encode(
+ &self,
+ node: Arc<dyn ExecutionPlan>,
+ buf: &mut Vec<u8>,
+ proto_converter: &dyn PhysicalProtoConverterExtension,
+ ) -> Result<()> {
+ let Some(exchange) = node.downcast_ref::<ExchangeExec>() else {
+ return self.inner.try_encode(node, buf, proto_converter);
+ };
+ buf.extend_from_slice(EXCHANGE_TAG);
+ buf.push(exchange.is_broadcast() as u8);
+ match exchange.partitioning() {
+ Some(partitioning) => {
+ buf.push(1);
+ serialize_partitioning(partitioning, self, proto_converter)?
+ .encode(buf)
+ .map_err(|e| {
+
datafusion::error::DataFusionError::Internal(e.to_string())
+ })?;
+ }
+ None => buf.push(0),
+ }
+ Ok(())
+ }
+
+ fn try_encode_expr(
+ &self,
+ node: &Arc<dyn PhysicalExpr>,
+ buf: &mut Vec<u8>,
+ ctx: &PhysicalExprEncodeCtx<'_>,
+ ) -> Result<()> {
+ self.inner.try_encode_expr(node, buf, ctx)
+ }
+
+ fn try_encode_higher_order_function(
+ &self,
+ node: &HigherOrderUDF,
+ buf: &mut Vec<u8>,
+ ) -> Result<()> {
+ self.inner.try_encode_higher_order_function(node, buf)
+ }
+
+ fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>>
{
Review Comment:
Keys are never decoded, so the `try_decode_udf`, `try_decode_udaf` and
`try_decode_udwf` forwards are dead code. `PhysicalExtensionCodec` provides
defaults for all three in datafusion-proto 55.2. Only the encode hooks need
forwarding.
##########
ballista/scheduler/src/state/aqe/execution_plan/exchange.rs:
##########
@@ -361,6 +361,58 @@ impl ExchangeExec {
exchange
}
+ /// This exchange as another reader of `canonical`'s stage, for a subplan
+ /// that appears more than once in the job's plan. The stage runs once
+ /// and both read the same shuffle output.
+ ///
+ /// What the stage writes is taken from `canonical`: its input, its
+ /// partitioning, and the shared stage id, shuffle partitions, column
+ /// statistics and range routing. How each consumer reads that output is
+ /// its own: this exchange keeps its `plan_id` and inactive flag, and gets
+ /// fresh coalesce and skew-join slots, since each consuming stage decides
+ /// those for its own alignment group.
+ pub(crate) fn share_stage_of(&self, canonical: &ExchangeExec) -> Self {
+ Self::new_with_details(
+ canonical.input.clone(),
+ canonical.partitioning.clone(),
+ self.plan_id,
+ canonical.stage_id.clone(),
+ canonical.shuffle_partitions.clone(),
+ canonical.column_statistics.clone(),
+ Arc::new(Mutex::new(None)),
+ canonical.range_repartition_routing.clone(),
+ canonical.broadcast,
+ self.inactive_stage,
+ )
+ }
+
+ /// Whether this exchange and `other` are readers of one stage, so that
+ /// resolving either resolves both.
+ pub(crate) fn shares_stage_with(&self, other: &ExchangeExec) -> bool {
+ self.stage_handle() == other.stage_handle()
+ }
+
+ /// Identifies the stage state this exchange shares with its other
+ /// readers, for grouping exchanges by stage.
+ pub(crate) fn stage_handle(&self) -> *const () {
+ Arc::as_ptr(&self.stage_id) as *const ()
+ }
+
+ /// This exchange with stage state of its own, unshared with any other
+ /// exchange, and no stage assigned yet.
+ pub(crate) fn with_own_stage(&self) -> Self {
+ let input = self.input.clone();
+ let partitioning = self.partitioning.clone();
+ let exchange = if self.broadcast {
+ Self::new_broadcast(input, partitioning, self.plan_id)
+ } else {
+ Self::new(input, partitioning, self.plan_id)
+ };
+ let mut exchange = exchange.with_inactive_stage(self.inactive_stage);
+ exchange.skew_join = self.skew_join.clone();
Review Comment:
Question: `with_own_stage` keeps `skew_join` but drops `coalesce` (it gets a
fresh slot from `new` or `new_broadcast`). `share_stage_of` resets both. Is
keeping one and dropping the other intended for a split pending stage? If both
are decided per consuming stage, resetting both would match `share_stage_of`.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
