andygrove commented on code in PR #2765:
URL:
https://github.com/apache/datafusion-ballista/pull/2765#discussion_r4482234714
##########
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:
Both are gone. The replan for a late repeat now lives in
`identify_runnable_stages`, so no rule needs to signal the planner. Every pass
resolves at least one pending exchange, so the loop is bounded by the plan's
exchange count, and going past that returns `internal_err!` rather than exiting
quietly.
##########
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:
Done, it's the `!claimed.insert(exchange.stage_handle())` arm right after
`inactive_stage`, with a sized set.
##########
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:
Renamed to `ballista.planner.reuse_exchange.enabled`, with
`BALLISTA_REUSE_EXCHANGE_ENABLED` and `reuse_exchange_enabled()`. The keying
lives in its own module (`state/aqe/reuse_exchange.rs`) and doesn't depend on
the adaptive planner, so it should be reusable for pipelined mode later.
##########
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:
`with_own_stage` is gone with the redesign, so only `share_stage_of`
remains, and it resets both.
##########
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:
Agreed, I dropped those two and kept
`each_reader_of_a_shared_stage_coalesces_on_its_own`.
##########
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:
Merged into one rstest over the query and `reuse`. The client test now has a
`reader_rules` arm that turns on coalesce and skew join, so per-reader
decisions on a shared stage get exercised end to end.
--
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]
