You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/07/22 12:40:39 UTC
(datafusion-ballista) branch branch-54 updated: fix: keep EmptyExec partition count intact across stage serialization (#2062) (#2132)
This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch branch-54
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/branch-54 by this push:
new aa3a4d5fd fix: keep EmptyExec partition count intact across stage serialization (#2062) (#2132)
aa3a4d5fd is described below
commit aa3a4d5fd052c2242e48383c66095d341cf63a7a
Author: Andy Grove <ag...@apache.org>
AuthorDate: Wed Jul 22 06:40:34 2026 -0600
fix: keep EmptyExec partition count intact across stage serialization (#2062) (#2132)
---
ballista/core/src/serde/mod.rs | 39 +++++++++++
ballista/scheduler/src/planner.rs | 75 ++++++++++++++++++++++
.../scheduler/src/state/aqe/test/alter_stages.rs | 6 +-
3 files changed, 118 insertions(+), 2 deletions(-)
diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs
index 0206afe5e..7150499ea 100644
--- a/ballista/core/src/serde/mod.rs
+++ b/ballista/core/src/serde/mod.rs
@@ -859,6 +859,45 @@ mod test {
);
}
+ // `datafusion-proto` encodes an `EmptyExec` as its schema alone, so a
+ // multi-partition `EmptyExec` decodes with a single partition (apache/datafusion
+ // #23642). `make_empty_exec_serde_safe` in the scheduler works around this by
+ // rewriting such nodes into a round-robin `RepartitionExec` before a stage plan
+ // goes on the wire.
+ //
+ // This test pins the upstream behaviour that makes the workaround necessary: when
+ // it starts failing, DataFusion preserves the partition count and the workaround
+ // can be deleted.
+ #[tokio::test]
+ async fn empty_exec_partition_count_is_lost_by_datafusion_proto() {
+ use datafusion::physical_plan::ExecutionPlanProperties;
+ use datafusion::physical_plan::empty::EmptyExec;
+ use datafusion_proto::physical_plan::{
+ AsExecutionPlan, DefaultPhysicalExtensionCodec,
+ };
+
+ let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
+ let plan: Arc<dyn ExecutionPlan> =
+ Arc::new(EmptyExec::new(schema).with_partitions(4));
+ assert_eq!(
+ plan.output_partitioning().partition_count(),
+ 4,
+ "precondition: EmptyExec reports 4 partitions"
+ );
+
+ let codec = DefaultPhysicalExtensionCodec {};
+ let proto = PhysicalPlanNode::try_from_physical_plan(plan, &codec).unwrap();
+ let ctx = SessionContext::new().task_ctx();
+ let decoded = proto.try_into_physical_plan(&ctx, &codec).unwrap();
+
+ assert_eq!(
+ decoded.output_partitioning().partition_count(),
+ 1,
+ "EmptyExec partition count is dropped on the wire; if this now decodes as \
+ 4, apache/datafusion#23642 is fixed and make_empty_exec_serde_safe can go"
+ );
+ }
+
#[tokio::test]
async fn test_unresolved_shuffle_exec_roundtrip() {
let schema = create_test_schema();
diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs
index 97cc1a7ac..f6993b722 100644
--- a/ballista/scheduler/src/planner.rs
+++ b/ballista/scheduler/src/planner.rs
@@ -33,10 +33,12 @@ use ballista_core::{
serde::scheduler::PartitionLocation,
};
use datafusion::arrow::datatypes::DataType;
+use datafusion::common::tree_node::{Transformed, TreeNode};
use datafusion::config::ConfigOptions;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_optimizer::enforce_sorting::EnforceSorting;
use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
+use datafusion::physical_plan::empty::EmptyExec;
use datafusion::physical_plan::joins::{
HashJoinExec, HashJoinExecBuilder, PartitionMode, SortMergeJoinExec,
};
@@ -672,6 +674,42 @@ pub fn rollback_resolved_shuffles(
Ok(with_new_children_if_necessary(stage, new_children)?)
}
+/// Rewrites every multi-partition [`EmptyExec`] into a round-robin [`RepartitionExec`]
+/// over a single-partition [`EmptyExec`], which is an equivalent plan whose partition
+/// count survives serialization.
+///
+/// `datafusion-proto` encodes an `EmptyExec` as its schema alone, so a multi-partition
+/// `EmptyExec` decodes on the executor with a single partition
+/// (<https://github.com/apache/datafusion/issues/23642>). A stage is sized from the
+/// scheduler-side partition count, so every task above partition 0 would then fail with
+/// `EmptyExec invalid partition N (expected less than 1)`. A `RepartitionExec` carries
+/// its partitioning on the wire, so the count arrives intact.
+///
+/// This runs at the point the stage plan is built for the wire, after all physical
+/// optimizer rules, so no later rule can collapse the rewrite. It can be removed once
+/// the partition count survives `EmptyExec` serialization upstream.
+fn make_empty_exec_serde_safe(
+ plan: Arc<dyn ExecutionPlan>,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(plan
+ .transform_up(|node| {
+ let Some(empty) = node.downcast_ref::<EmptyExec>() else {
+ return Ok(Transformed::no(node));
+ };
+ let partition_count =
+ empty.properties().output_partitioning().partition_count();
+ if partition_count <= 1 {
+ return Ok(Transformed::no(node));
+ }
+ let single: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(empty.schema()));
+ Ok(Transformed::yes(Arc::new(RepartitionExec::try_new(
+ single,
+ Partitioning::RoundRobinBatch(partition_count),
+ )?)))
+ })?
+ .data)
+}
+
pub(crate) fn create_shuffle_writer_with_config(
job_id: &JobId,
stage_id: usize,
@@ -679,6 +717,8 @@ pub(crate) fn create_shuffle_writer_with_config(
partitioning: Option<Partitioning>,
config: &ConfigOptions,
) -> Result<Arc<dyn ShuffleWriter>> {
+ let plan = make_empty_exec_serde_safe(plan)?;
+
// Check if sort-based shuffle is enabled
let ballista_config = config
.extensions
@@ -748,6 +788,41 @@ mod test {
use std::sync::Arc;
use uuid::Uuid;
+ #[test]
+ fn multi_partition_empty_exec_is_rewritten_for_the_wire() {
+ use super::make_empty_exec_serde_safe;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::physical_plan::empty::EmptyExec;
+
+ let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
+ let plan: Arc<dyn ExecutionPlan> =
+ Arc::new(EmptyExec::new(schema).with_partitions(4));
+
+ let rewritten = make_empty_exec_serde_safe(plan).unwrap();
+
+ assert_eq!(
+ displayable(rewritten.as_ref()).indent(false).to_string(),
+ "RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1\n EmptyExec\n"
+ );
+ }
+
+ #[test]
+ fn single_partition_empty_exec_is_left_alone() {
+ use super::make_empty_exec_serde_safe;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::physical_plan::empty::EmptyExec;
+
+ let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
+ let plan: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(schema));
+
+ let rewritten = make_empty_exec_serde_safe(plan).unwrap();
+
+ assert_eq!(
+ displayable(rewritten.as_ref()).indent(false).to_string(),
+ "EmptyExec\n"
+ );
+ }
+
macro_rules! downcast_exec {
($exec: expr, $ty: ty) => {
($exec.as_ref() as &dyn ExecutionPlan)
diff --git a/ballista/scheduler/src/state/aqe/test/alter_stages.rs b/ballista/scheduler/src/state/aqe/test/alter_stages.rs
index 7757cddd7..1bd7e5d65 100644
--- a/ballista/scheduler/src/state/aqe/test/alter_stages.rs
+++ b/ballista/scheduler/src/state/aqe/test/alter_stages.rs
@@ -72,7 +72,8 @@ async fn should_propagate_empty_stage() -> datafusion::error::Result<()> {
assert_eq!(1, stages.len());
assert_plan!(stages.first().unwrap().plan.as_ref(), @ r"
ShuffleWriterExec: partitioning: None
- EmptyExec
+ RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1
+ EmptyExec
");
planner.finalise_stage_internal(1, mock_partitions_with_statistics_no_data())?;
@@ -146,7 +147,8 @@ async fn should_propagate_empty_stage_and_remove() -> datafusion::error::Result<
assert_eq!(1, stages.len());
assert_plan!(stages.first().unwrap().plan.as_ref(), @ r"
ShuffleWriterExec: partitioning: None
- EmptyExec
+ RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1
+ EmptyExec
");
planner.finalise_stage_internal(1, mock_partitions_with_statistics_no_data())?;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org