fix: [branch-54] keep EmptyExec partition count intact across stage serialization (backport of #2062)

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();

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,13 +674,51 @@ 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,
     plan: Arc<dyn ExecutionPlan>,
     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)

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())?;