diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index 0206afe5e8..7150499ea2 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 = + 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 97cc1a7ac9..f6993b7227 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 +/// (). 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, +) -> Result> { + Ok(plan + .transform_up(|node| { + let Some(empty) = node.downcast_ref::() 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 = 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, config: &ConfigOptions, ) -> Result> { + 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 = + 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 = 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 7757cddd79..1bd7e5d65e 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())?;