You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by mi...@apache.org on 2026/08/10 20:30:04 UTC
(datafusion-ballista) branch main updated: fix(scheduler): key AQE shuffle output by job_id (#2278)
This is an automated email from the ASF dual-hosted git repository.
milenkovicm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new 2ffc29bec fix(scheduler): key AQE shuffle output by job_id (#2278)
2ffc29bec is described below
commit 2ffc29bec141d4f18418e65bdebb93a5f78da3e0
Author: Ville Brofeldt <33...@users.noreply.github.com>
AuthorDate: Mon Aug 10 13:29:59 2026 -0700
fix(scheduler): key AQE shuffle output by job_id (#2278)
---
ballista/scheduler/src/state/aqe/mod.rs | 2 +-
ballista/scheduler/src/state/aqe/planner.rs | 31 +++---
.../src/state/aqe/test/broadcast_thresholds.rs | 2 +-
.../scheduler/src/state/aqe/test/join_selection.rs | 14 +--
.../scheduler/src/state/aqe/test/plan_to_stages.rs | 111 +++++++++++++--------
5 files changed, 92 insertions(+), 68 deletions(-)
diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs
index dac848221..acebf50eb 100644
--- a/ballista/scheduler/src/state/aqe/mod.rs
+++ b/ballista/scheduler/src/state/aqe/mod.rs
@@ -142,7 +142,7 @@ impl AdaptiveExecutionGraph {
let session_id = ctx.session_id();
let mut planner =
- AdaptivePlanner::try_new(ctx, logical_plan, job_name.to_owned()).await?;
+ AdaptivePlanner::try_new(ctx, logical_plan, job_id.clone()).await?;
let logical_plan = Some(logical_plan.display_indent().to_string());
diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs
index 6dc08dab6..ed3e29f50 100644
--- a/ballista/scheduler/src/state/aqe/planner.rs
+++ b/ballista/scheduler/src/state/aqe/planner.rs
@@ -70,8 +70,8 @@ pub struct AdaptivePlanner {
pub(crate) plan: Arc<dyn ExecutionPlan>,
/// caches current runnable stages
runnable_stage_cache: HashMap<usize, Arc<dyn ExecutionPlan>>,
- /// job name
- job_name: String,
+ /// Unique identifier for the job, used to key this job's shuffle output.
+ job_id: JobId,
runnable_stage_output: HashMap<usize, StageOutput>,
}
@@ -79,7 +79,7 @@ pub struct AdaptivePlanner {
impl Debug for AdaptivePlanner {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AdaptivePlanner")
- .field("job_name", &self.job_name)
+ .field("job_id", &self.job_id)
.finish()
}
}
@@ -90,7 +90,7 @@ impl AdaptivePlanner {
/// # Arguments:
/// * `state_builder` - Session state builder,
/// * `plan` - The physical execution plan for the job.
- /// * `job_name` - The name of the job.
+ /// * `job_id` - The unique identifier for the job.
/// * `physical_optimizer_rules` - A list of physical optimizer rules to apply.
///
/// # Returns
@@ -98,7 +98,7 @@ impl AdaptivePlanner {
pub fn try_new_with_optimizers(
state_builder: SessionStateBuilder,
plan: Arc<dyn ExecutionPlan>,
- job_name: String,
+ job_id: JobId,
physical_optimizer_rules: Vec<PhysicalOptimizerRuleRef>,
) -> common::Result<Self> {
let session_state =
@@ -112,7 +112,7 @@ impl AdaptivePlanner {
physical_planner: planner.into(),
plan,
runnable_stage_cache: HashMap::new(),
- job_name,
+ job_id,
runnable_stage_output: HashMap::new(),
})
}
@@ -122,7 +122,7 @@ impl AdaptivePlanner {
///
/// * `session_config` - The session configuration for the job.
/// * `plan` - The physical execution plan for the job.
- /// * `job_name` - The name of the job.
+ /// * `job_id` - The unique identifier for the job.
///
/// # Returns
/// A new instance of `AdaptivePlanner` or an error if the initialization fails.
@@ -130,7 +130,7 @@ impl AdaptivePlanner {
pub fn try_from_plan(
session_config: &SessionConfig,
plan: Arc<dyn ExecutionPlan>,
- job_name: String,
+ job_id: JobId,
) -> common::Result<Self> {
let plan_id_generator = Arc::new(AtomicUsize::new(0));
let state_builder = SessionStateBuilder::new_with_default_features()
@@ -138,7 +138,7 @@ impl AdaptivePlanner {
Self::try_new_with_optimizers(
state_builder,
plan,
- job_name,
+ job_id,
Self::default_optimizers(plan_id_generator),
)
}
@@ -149,14 +149,14 @@ impl AdaptivePlanner {
///
/// * `ctx` - The session context
/// * `logical_plan` - The logical plan for the job.
- /// * `job_name` - The name of the job.
+ /// * `job_id` - The unique identifier for the job.
///
/// # Returns
/// A new instance of `AdaptivePlanner` or an error if the initialization fails.
pub async fn try_new(
ctx: &SessionContext,
logical_plan: &LogicalPlan,
- job_name: String,
+ job_id: JobId,
) -> common::Result<Self> {
// session state with very limited set of optimizers.
// this optimizer set will be executed only once, before
@@ -174,9 +174,6 @@ impl AdaptivePlanner {
.create_physical_plan(logical_plan)
.await?;
- // Note: the signature requires a JobId, but we are passing a JobName. The below is a
- // dirty fix but this seems like a bug or a design flaw.
- let job_id: JobId = job_name.clone().into();
let plan = handle_explain_plan(&job_id, ctx, logical_plan, plan)
.await
.map_err(|e| DataFusionError::Execution(e.to_string()))?;
@@ -185,7 +182,7 @@ impl AdaptivePlanner {
Self::try_new_with_optimizers(
state_builder,
plan,
- job_name,
+ job_id,
Self::default_optimizers(plan_id_generator),
)
}
@@ -417,9 +414,7 @@ impl AdaptivePlanner {
// that would arise if the rule walked the entire residual
// plan in `default_optimizers()`.
let plan = CoalescePartitionsRule.optimize(plan, config)?;
- // adapt_to_ballista takes an job_id, we are passing a job_name. Need to transform to fix compiler.
- let job_id = self.job_name.clone().into();
- BallistaAdapter::adapt_to_ballista(plan, &job_id, config)
+ BallistaAdapter::adapt_to_ballista(plan, &self.job_id, config)
.map(|w| (w.plan.stage_id(), w))
})
.collect::<common::Result<(HashSet<_>, Vec<_>)>>()?;
diff --git a/ballista/scheduler/src/state/aqe/test/broadcast_thresholds.rs b/ballista/scheduler/src/state/aqe/test/broadcast_thresholds.rs
index dfa0fac75..411955284 100644
--- a/ballista/scheduler/src/state/aqe/test/broadcast_thresholds.rs
+++ b/ballista/scheduler/src/state/aqe/test/broadcast_thresholds.rs
@@ -83,7 +83,7 @@ fn register(ctx: &SessionContext, name: &str, schema: Arc<Schema>, stats: Statis
/// Resolves the join and reports whether any side ended up broadcast.
async fn plan_broadcasts(ctx: &SessionContext, sql: &str) -> (bool, String) {
let lp = ctx.sql(sql).await.unwrap().into_optimized_plan().unwrap();
- let planner = AdaptivePlanner::try_new(ctx, &lp, "test_job".to_owned())
+ let planner = AdaptivePlanner::try_new(ctx, &lp, "test_job".into())
.await
.unwrap();
let plan = displayable(planner.current_plan()).indent(true).to_string();
diff --git a/ballista/scheduler/src/state/aqe/test/join_selection.rs b/ballista/scheduler/src/state/aqe/test/join_selection.rs
index 428cb7c6f..9bf234778 100644
--- a/ballista/scheduler/src/state/aqe/test/join_selection.rs
+++ b/ballista/scheduler/src/state/aqe/test/join_selection.rs
@@ -154,7 +154,7 @@ async fn test_hash_join_two_tables_coalesce() -> datafusion::common::Result<()>
.into_optimized_plan()
.unwrap();
- let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned()).await?;
+ let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".into()).await?;
let plan = planner.current_plan();
// TODO: we could probably push datasource to read single partition
assert_plan!(plan, @ r"
@@ -185,7 +185,7 @@ async fn test_left_join_not_collected_left() -> datafusion::common::Result<()> {
.into_optimized_plan()
.unwrap();
- let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned()).await?;
+ let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".into()).await?;
let plan = planner.current_plan();
assert_plan!(plan, @ r"
AdaptiveDatafusionExec: is_final=false, plan_id=3, stage_id=pending, stage_resolved=false
@@ -211,7 +211,7 @@ async fn test_hash_join_two_tables_repartition() -> datafusion::common::Result<(
.into_optimized_plan()
.unwrap();
- let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned()).await?;
+ let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".into()).await?;
let plan = planner.current_plan();
assert_plan!(plan, @ r"
@@ -272,7 +272,7 @@ async fn test_sort_merge_join_two_tables_repartition() -> datafusion::common::Re
.into_optimized_plan()
.unwrap();
- let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned()).await?;
+ let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".into()).await?;
let plan = planner.current_plan();
assert_plan!(plan, @ r"
@@ -338,7 +338,7 @@ async fn test_hash_join_three_tables_collect_left() -> datafusion::common::Resul
.unwrap()
.into_optimized_plan()
.unwrap();
- let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned()).await?;
+ let mut planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".into()).await?;
let plan = planner.current_plan();
assert_plan!(plan, @ r"
@@ -409,7 +409,7 @@ async fn test_hash_join_three_tables_repartition() -> datafusion::common::Result
.unwrap()
.into_optimized_plan()
.unwrap();
- let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned()).await?;
+ let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".into()).await?;
let plan = planner.current_plan();
assert_plan!(plan, @ r"
@@ -440,7 +440,7 @@ async fn test_sort_merge_join_three_tables_repartition() -> datafusion::common::
.unwrap()
.into_optimized_plan()
.unwrap();
- let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".to_owned()).await?;
+ let planner = AdaptivePlanner::try_new(&ctx, &lp, "test_job".into()).await?;
let plan = planner.current_plan();
assert_plan!(plan, @ r"
diff --git a/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs b/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs
index 993996af4..d744050cb 100644
--- a/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs
+++ b/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs
@@ -16,12 +16,16 @@
// under the License.
use crate::assert_plan;
+use crate::state::aqe::AdaptiveExecutionGraph;
use crate::state::aqe::execution_plan::ExchangeExec;
use crate::state::aqe::planner::AdaptivePlanner;
use crate::state::aqe::test::{
mock_batch, mock_context, mock_memory_table, mock_partitions_with_statistics,
};
+use crate::state::execution_graph::ExecutionGraph;
+use ballista_core::JobId;
use ballista_core::execution_plans::SortShuffleWriterExec;
+use ballista_core::serde::protobuf::job_status::Status;
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::common::ColumnStatistics;
use datafusion::physical_plan::ExecutionPlan;
@@ -40,11 +44,8 @@ async fn should_add_exchanges() -> datafusion::error::Result<()> {
"#;
let plan = ctx.sql(q).await?.create_physical_plan().await?;
- let planner = AdaptivePlanner::try_from_plan(
- ctx.state().config(),
- plan,
- "test_job".to_owned(),
- )?;
+ let planner =
+ AdaptivePlanner::try_from_plan(ctx.state().config(), plan, "test_job".into())?;
assert_plan!(planner.current_plan(), @ r"
AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending, stage_resolved=false
@@ -69,11 +70,8 @@ async fn should_split_plan_into_runnable_stages_internal() -> datafusion::error:
"#;
let plan = ctx.sql(q).await?.create_physical_plan().await?;
- let mut planner = AdaptivePlanner::try_from_plan(
- ctx.state().config(),
- plan,
- "test_job".to_owned(),
- )?;
+ let mut planner =
+ AdaptivePlanner::try_from_plan(ctx.state().config(), plan, "test_job".into())?;
assert_plan!(planner.current_plan(), @ r"
AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending, stage_resolved=false
@@ -121,11 +119,8 @@ async fn should_split_plan_into_stages() -> datafusion::error::Result<()> {
"#;
let plan = ctx.sql(q).await?.create_physical_plan().await?;
- let mut planner = AdaptivePlanner::try_from_plan(
- ctx.state().config(),
- plan,
- "test_job".to_owned(),
- )?;
+ let mut planner =
+ AdaptivePlanner::try_from_plan(ctx.state().config(), plan, "test_job".into())?;
assert_plan!(planner.current_plan(), @ r"
AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending, stage_resolved=false
@@ -143,6 +138,7 @@ async fn should_split_plan_into_stages() -> datafusion::error::Result<()> {
AggregateExec: mode=Partial, gby=[c@2 as c], aggr=[min(t.a), max(t.b)]
DataSourceExec: partitions=1, partition_sizes=[1]
");
+ assert_eq!(stages.first().unwrap().plan.job_id().as_str(), "test_job");
planner.finalise_stage_internal(0, mock_partitions_with_statistics())?;
@@ -154,6 +150,7 @@ async fn should_split_plan_into_stages() -> datafusion::error::Result<()> {
AggregateExec: mode=FinalPartitioned, gby=[c@0 as c], aggr=[min(t.a), max(t.b)]
ShuffleReaderExec: upstream_stage: 0, partitioning: Hash([c@0], 2)
");
+ assert_eq!(stages.first().unwrap().plan.job_id().as_str(), "test_job");
planner.finalise_stage_internal(1, mock_partitions_with_statistics())?;
let stages = planner.runnable_stages()?;
@@ -180,11 +177,8 @@ async fn should_create_initial_plan() -> datafusion::error::Result<()> {
let plan = ctx.sql(q).await?.create_physical_plan().await?;
- let planner = AdaptivePlanner::try_from_plan(
- ctx.state().config(),
- plan,
- "test_job".to_owned(),
- )?;
+ let planner =
+ AdaptivePlanner::try_from_plan(ctx.state().config(), plan, "test_job".into())?;
// plan has only two exchanges after initial planning
// other stages will be added as stages get resolved
@@ -230,11 +224,8 @@ async fn should_split_stages_resolve_right_branch() -> datafusion::error::Result
let plan = ctx.sql(q).await?.create_physical_plan().await?;
- let mut planner = AdaptivePlanner::try_from_plan(
- ctx.state().config(),
- plan,
- "test_job".to_owned(),
- )?;
+ let mut planner =
+ AdaptivePlanner::try_from_plan(ctx.state().config(), plan, "test_job".into())?;
let runnable_stages = planner.identify_runnable_stages()?.unwrap();
assert_eq!(2, runnable_stages.len());
@@ -300,11 +291,8 @@ async fn should_split_stages_resolve_left_branch() -> datafusion::error::Result<
let plan = ctx.sql(q).await?.create_physical_plan().await?;
- let mut planner = AdaptivePlanner::try_from_plan(
- ctx.state().config(),
- plan,
- "test_job".to_owned(),
- )?;
+ let mut planner =
+ AdaptivePlanner::try_from_plan(ctx.state().config(), plan, "test_job".into())?;
let runnable_stages = planner.identify_runnable_stages()?.unwrap();
assert_eq!(2, runnable_stages.len());
@@ -375,11 +363,8 @@ async fn should_split_stages_resolve_both() -> datafusion::error::Result<()> {
let plan = ctx.sql(q).await?.create_physical_plan().await?;
- let mut planner = AdaptivePlanner::try_from_plan(
- ctx.state().config(),
- plan,
- "test_job".to_owned(),
- )?;
+ let mut planner =
+ AdaptivePlanner::try_from_plan(ctx.state().config(), plan, "test_job".into())?;
let runnable_stages = planner.identify_runnable_stages()?.unwrap();
assert_eq!(2, runnable_stages.len());
@@ -438,7 +423,7 @@ async fn should_ignore_inactive_stages() -> datafusion::error::Result<()> {
let mut planner = AdaptivePlanner::try_from_plan(
ctx.state().config(),
exchange_exec,
- "test_job".to_owned(),
+ "test_job".into(),
)?;
assert_plan!(planner.current_plan(), @ r"
@@ -464,11 +449,8 @@ async fn should_use_sort_shuffle_by_default() -> datafusion::error::Result<()> {
"#;
let plan = ctx.sql(q).await?.create_physical_plan().await?;
- let mut planner = AdaptivePlanner::try_from_plan(
- ctx.state().config(),
- plan,
- "test_job".to_owned(),
- )?;
+ let mut planner =
+ AdaptivePlanner::try_from_plan(ctx.state().config(), plan, "test_job".into())?;
let stages = planner.runnable_stages()?.unwrap();
assert_eq!(1, stages.len());
@@ -483,3 +465,50 @@ async fn should_use_sort_shuffle_by_default() -> datafusion::error::Result<()> {
Ok(())
}
+
+#[tokio::test]
+async fn newly_built_adaptive_graph_has_expected_state()
+-> ballista_core::error::Result<()> {
+ let ctx = mock_context();
+ ctx.register_batch("t", mock_batch()?)?;
+
+ let q = r#"
+ select min(a) as c0, max(b) as c1, c as c2 from t group by c
+ "#;
+ let logical_plan = ctx.sql(q).await?.logical_plan().clone();
+ let job_id: JobId = "job-id-xyz".into();
+
+ let mut graph = AdaptiveExecutionGraph::try_new(
+ "scheduler-1",
+ &job_id,
+ "my job name",
+ &ctx,
+ &logical_plan,
+ 7,
+ )
+ .await?;
+
+ // Identity comes straight from the constructor arguments.
+ assert_eq!(graph.job_id().as_str(), "job-id-xyz");
+ assert_eq!(graph.job_name(), "my job name");
+ assert_eq!(graph.session_id(), ctx.session_id().as_str());
+ assert!(graph.logical_plan().is_some());
+
+ // A freshly built graph is a running job that has done no work yet.
+ assert!(matches!(
+ graph.status().status.as_ref(),
+ Some(Status::Running(_))
+ ));
+ assert!(graph.start_time() > 0);
+ assert_eq!(graph.end_time(), 0);
+ assert_eq!(graph.completed_stages(), 0);
+ assert!(!graph.is_successful());
+ assert!(graph.output_locations().is_empty());
+
+ // The plan was split into stages that expose runnable tasks once revived.
+ assert!(graph.revive());
+ assert!(!graph.running_stages().is_empty());
+ assert!(graph.available_tasks() > 0);
+
+ Ok(())
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org