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/05/09 20:21:57 UTC
(datafusion-ballista) branch main updated: minor: add pending stage indicator for `AdaptiveDatafusionExec` (#1672)
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 aaf8f082f minor: add pending stage indicator for `AdaptiveDatafusionExec` (#1672)
aaf8f082f is described below
commit aaf8f082ff19d88829ac08b1b90bda7c56e1619c
Author: Marko Milenković <mi...@users.noreply.github.com>
AuthorDate: Sat May 9 21:21:52 2026 +0100
minor: add pending stage indicator for `AdaptiveDatafusionExec` (#1672)
---
ballista/scheduler/src/state/aqe/execution_plan.rs | 14 ++++++-------
.../aqe/optimizer_rule/distributed_exchange.rs | 2 +-
.../scheduler/src/state/aqe/test/alter_stages.rs | 24 +++++++++++-----------
.../scheduler/src/state/aqe/test/plan_to_stages.rs | 12 +++++------
4 files changed, 25 insertions(+), 27 deletions(-)
diff --git a/ballista/scheduler/src/state/aqe/execution_plan.rs b/ballista/scheduler/src/state/aqe/execution_plan.rs
index 94fa3ee57..3fc23de49 100644
--- a/ballista/scheduler/src/state/aqe/execution_plan.rs
+++ b/ballista/scheduler/src/state/aqe/execution_plan.rs
@@ -210,7 +210,7 @@ impl DisplayAs for ExchangeExec {
self.stage_id()
.map(|stage_id| format!("{}", stage_id))
.unwrap_or_else(|| "pending".to_string()),
- self.shuffle_partitions.lock().is_some()
+ self.shuffle_created()
)
}
DisplayFormatType::TreeRender => {
@@ -230,11 +230,7 @@ impl DisplayAs for ExchangeExec {
.map(|stage_id| format!("({})", stage_id))
.unwrap_or_else(|| "pending".to_string()),
)?;
- writeln!(
- f,
- "stage_resolved={}",
- self.shuffle_partitions.lock().is_some()
- )
+ writeln!(f, "stage_resolved={}", self.shuffle_created())
}
}
}
@@ -411,12 +407,13 @@ impl DisplayAs for AdaptiveDatafusionExec {
DisplayFormatType::Default | DisplayFormatType::Verbose => {
write!(
f,
- "AdaptiveDatafusionExec: is_final={:?}, plan_id={}, stage_id={}",
+ "AdaptiveDatafusionExec: is_final={:?}, plan_id={}, stage_id={}, stage_resolved={}",
self.is_final,
self.plan_id,
self.stage_id()
.map(|stage_id| format!("{}", stage_id))
.unwrap_or_else(|| "pending".to_string()),
+ self.shuffle_created()
)
}
DisplayFormatType::TreeRender => {
@@ -428,7 +425,8 @@ impl DisplayAs for AdaptiveDatafusionExec {
self.stage_id()
.map(|stage_id| format!("({})", stage_id))
.unwrap_or_else(|| "pending".to_string()),
- )
+ )?;
+ writeln!(f, "stage_resolved={}", self.shuffle_created())
}
}
}
diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs
index ead623f73..739cf460d 100644
--- a/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs
+++ b/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs
@@ -221,7 +221,7 @@ mod tests {
let result = rule.optimize(input, &config()).unwrap();
assert_plan!(result.as_ref(), @ r"
- AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending
+ AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending, stage_resolved=false
CoalescePartitionsExec
ExchangeExec: partitioning=None, plan_id=0, stage_id=pending, stage_resolved=false
StatisticsExec: col_count=1, row_count=Absent
diff --git a/ballista/scheduler/src/state/aqe/test/alter_stages.rs b/ballista/scheduler/src/state/aqe/test/alter_stages.rs
index 465fea359..9d67fa699 100644
--- a/ballista/scheduler/src/state/aqe/test/alter_stages.rs
+++ b/ballista/scheduler/src/state/aqe/test/alter_stages.rs
@@ -118,7 +118,7 @@ async fn should_propagate_empty_stage_and_remove() -> datafusion::error::Result<
AdaptivePlanner::try_new(ctx.state().config(), plan, "test_job".to_string())?;
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending
+ AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending, stage_resolved=false
ProjectionExec: expr=[c0@0 as c0, count(Int64(1))@1 as count(*)]
AggregateExec: mode=FinalPartitioned, gby=[c0@0 as c0], aggr=[count(Int64(1))]
RepartitionExec: partitioning=Hash([c0@0], 2), input_partitions=2
@@ -153,7 +153,7 @@ async fn should_propagate_empty_stage_and_remove() -> datafusion::error::Result<
let stages = planner.runnable_stages()?;
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=true, plan_id=1, stage_id=2
+ AdaptiveDatafusionExec: is_final=true, plan_id=1, stage_id=2, stage_resolved=false
EmptyExec
");
assert!(stages.is_some());
@@ -184,7 +184,7 @@ async fn should_support_join_re_ordering() -> datafusion::error::Result<()> {
AdaptivePlanner::try_new(ctx.state().config(), join, "test_job".to_string())?;
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending
+ AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending, stage_resolved=false
HashJoinExec: mode=Partitioned, join_type=Inner, on=[(big_col@0, big_col@0)]
ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=0, stage_id=pending, stage_resolved=false
MockPartitionedScan: num_partitions=2, statistics=[Rows=Exact(262144), Bytes=Exact(2097152), [(Col[0]:)]]
@@ -196,7 +196,7 @@ async fn should_support_join_re_ordering() -> datafusion::error::Result<()> {
assert_eq!(2, stages.len());
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending
+ AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending, stage_resolved=false
HashJoinExec: mode=Partitioned, join_type=Inner, on=[(big_col@0, big_col@0)]
ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=0, stage_id=0, stage_resolved=false
MockPartitionedScan: num_partitions=2, statistics=[Rows=Exact(262144), Bytes=Exact(2097152), [(Col[0]:)]]
@@ -212,7 +212,7 @@ async fn should_support_join_re_ordering() -> datafusion::error::Result<()> {
// join ordering changes as build side is bigger than probe side
// after exchange statistic updated.
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending
+ AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending, stage_resolved=false
ProjectionExec: expr=[big_col@1 as big_col, big_col@0 as big_col]
HashJoinExec: mode=Partitioned, join_type=Inner, on=[(big_col@0, big_col@0)]
ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=1, stage_id=1, stage_resolved=true
@@ -227,7 +227,7 @@ async fn should_support_join_re_ordering() -> datafusion::error::Result<()> {
assert_eq!(1, stages.len());
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=true, plan_id=2, stage_id=2
+ AdaptiveDatafusionExec: is_final=true, plan_id=2, stage_id=2, stage_resolved=false
ProjectionExec: expr=[big_col@1 as big_col, big_col@0 as big_col]
HashJoinExec: mode=Partitioned, join_type=Inner, on=[(big_col@0, big_col@0)]
ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=1, stage_id=1, stage_resolved=true
@@ -241,7 +241,7 @@ async fn should_support_join_re_ordering() -> datafusion::error::Result<()> {
planner.finalise_stage_internal(2, small_statistics_exchange())?;
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=true, plan_id=2, stage_id=2
+ AdaptiveDatafusionExec: is_final=true, plan_id=2, stage_id=2, stage_resolved=true
ProjectionExec: expr=[big_col@1 as big_col, big_col@0 as big_col]
HashJoinExec: mode=Partitioned, join_type=Inner, on=[(big_col@0, big_col@0)]
ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=1, stage_id=1, stage_resolved=true
@@ -285,7 +285,7 @@ async fn should_support_cross_join() -> datafusion::error::Result<()> {
//
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending
+ AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending, stage_resolved=false
ProjectionExec: expr=[big_col@1 as big_col, big_col@0 as big_col]
CrossJoinExec
CoalescePartitionsExec
@@ -309,7 +309,7 @@ async fn should_support_cross_join() -> datafusion::error::Result<()> {
planner.finalise_stage_internal(0, small_statistics_exchange())?;
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending
+ AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending, stage_resolved=false
ProjectionExec: expr=[big_col@1 as big_col, big_col@0 as big_col]
CrossJoinExec
CoalescePartitionsExec
@@ -339,7 +339,7 @@ async fn should_support_cross_join() -> datafusion::error::Result<()> {
planner.finalise_stage_internal(1, big_statistics_exchange())?;
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=true, plan_id=1, stage_id=1
+ AdaptiveDatafusionExec: is_final=true, plan_id=1, stage_id=1, stage_resolved=true
ProjectionExec: expr=[big_col@1 as big_col, big_col@0 as big_col]
CrossJoinExec
CoalescePartitionsExec
@@ -388,7 +388,7 @@ async fn should_cancel_the_stage() -> datafusion::error::Result<()> {
assert_eq!(0, cancellable.len());
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending
+ AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending, stage_resolved=false
HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a@0)], projection=[a@1, b@2, c@3]
ExchangeExec: partitioning=Hash([a@0], 2), plan_id=0, stage_id=0, stage_resolved=false
FilterExec: b@1 = 42, projection=[a@0]
@@ -400,7 +400,7 @@ async fn should_cancel_the_stage() -> datafusion::error::Result<()> {
planner.finalise_stage_internal(0, mock_partitions_with_statistics_no_data())?;
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending
+ AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending, stage_resolved=false
EmptyExec
");
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 d093ebdc2..6bccca975 100644
--- a/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs
+++ b/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs
@@ -44,7 +44,7 @@ async fn should_add_exchanges() -> datafusion::error::Result<()> {
AdaptivePlanner::try_new(ctx.state().config(), plan, "test_job".to_string())?;
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending
+ AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending, stage_resolved=false
ProjectionExec: expr=[min(t.a)@1 as c0, max(t.b)@2 as c1, c@0 as c2]
AggregateExec: mode=FinalPartitioned, gby=[c@0 as c], aggr=[min(t.a), max(t.b)]
ExchangeExec: partitioning=Hash([c@0], 2), plan_id=0, stage_id=pending, stage_resolved=false
@@ -70,7 +70,7 @@ async fn should_split_plan_into_runnable_stages_internal() -> datafusion::error:
AdaptivePlanner::try_new(ctx.state().config(), plan, "test_job".to_string())?;
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending
+ AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending, stage_resolved=false
ProjectionExec: expr=[min(t.a)@1 as c0, max(t.b)@2 as c1, c@0 as c2]
AggregateExec: mode=FinalPartitioned, gby=[c@0 as c], aggr=[min(t.a), max(t.b)]
ExchangeExec: partitioning=Hash([c@0], 2), plan_id=0, stage_id=pending, stage_resolved=false
@@ -90,7 +90,7 @@ async fn should_split_plan_into_runnable_stages_internal() -> datafusion::error:
let stages = planner.identify_runnable_stages()?.unwrap();
assert_eq!(1, stages.len());
assert_plan!(stages.first().unwrap().as_ref(), @ r"
- AdaptiveDatafusionExec: is_final=true, plan_id=1, stage_id=1
+ AdaptiveDatafusionExec: is_final=true, plan_id=1, stage_id=1, stage_resolved=false
ProjectionExec: expr=[min(t.a)@1 as c0, max(t.b)@2 as c1, c@0 as c2]
AggregateExec: mode=FinalPartitioned, gby=[c@0 as c], aggr=[min(t.a), max(t.b)]
ExchangeExec: partitioning=Hash([c@0], 2), plan_id=0, stage_id=0, stage_resolved=true
@@ -119,7 +119,7 @@ async fn should_split_plan_into_stages() -> datafusion::error::Result<()> {
AdaptivePlanner::try_new(ctx.state().config(), plan, "test_job".to_string())?;
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending
+ AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending, stage_resolved=false
ProjectionExec: expr=[min(t.a)@1 as c0, max(t.b)@2 as c1, c@0 as c2]
AggregateExec: mode=FinalPartitioned, gby=[c@0 as c], aggr=[min(t.a), max(t.b)]
ExchangeExec: partitioning=Hash([c@0], 2), plan_id=0, stage_id=pending, stage_resolved=false
@@ -177,7 +177,7 @@ async fn should_create_initial_plan() -> datafusion::error::Result<()> {
// plan has only two exchanges after initial planning
// other stages will be added as stages get resolved
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending
+ AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending, stage_resolved=false
ProjectionExec: expr=[sum(t0.c0)@1 as sum(t0.c0)]
AggregateExec: mode=FinalPartitioned, gby=[c0@0 as c0], aggr=[sum(t0.c0)]
RepartitionExec: partitioning=Hash([c0@0], 2), input_partitions=2
@@ -421,7 +421,7 @@ async fn should_ignore_inactive_stages() -> datafusion::error::Result<()> {
)?;
assert_plan!(planner.current_plan(), @ r"
- AdaptiveDatafusionExec: is_final=false, plan_id=0, stage_id=pending
+ AdaptiveDatafusionExec: is_final=false, plan_id=0, stage_id=pending, stage_resolved=false
ExchangeExec: partitioning=None, plan_id=0, stage_id=pending, stage_resolved=false
CooperativeExec
StatisticsExec: col_count=1, row_count=Absent
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org