datafusion-proto drops EmptyExec partition count on physical plan round-trip

[andygrove]

Describe the bug

datafusion-proto encodes EmptyExec as its schema only, so the partition count set by EmptyExec::with_partitions(n) is silently lost across a physical-plan round-trip. A plan that reported n partitions before serialization reports 1 after.

EmptyExecNode has a single field:

message EmptyExecNode {
  datafusion_common.Schema schema = 1;
}

and the decoder rebuilds the node with the default partition count:

https://github.com/apache/datafusion/blob/main/datafusion/proto/src/physical_plan/mod.rs — try_into_physical_plan for the Empty arm does Ok(Arc::new(EmptyExec::new(schema))), dropping partitions, while the encoder writes only schema.

This is a silent data-loss bug rather than an error: the decoded plan is well-formed but describes a different plan than the one encoded. It bites any distributed setup that plans on one process and executes on another. In Apache DataFusion Ballista, an optimizer rule collapses provably-empty sub-plans into an EmptyExec that inherits the partition count of the node it replaces; the scheduler sizes the stage's task count from that count, ships the plan to an executor, and every task above partition 0 then fails with:

Internal("Assertion failed: partition < self.partitions: EmptyExec invalid partition 1 (expected less than 1)")

PlaceholderRowExec has the same shape (with_partitions + schema-only proto) and looks like it has the same defect.

To Reproduce

use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::physical_plan::empty::EmptyExec;
use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
use datafusion::prelude::SessionContext;
use datafusion_proto::physical_plan::{AsExecutionPlan, DefaultPhysicalExtensionCodec};
use datafusion_proto::protobuf::PhysicalPlanNode;
use std::sync::Arc;

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

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

// fails: left: 1, right: 4
assert_eq!(decoded.output_partitioning().partition_count(), 4);

Expected behavior

The decoded EmptyExec reports the same partition count as the encoded one — the assertion above passes.

Additional context

Reproduced on 54.0.0; the same code is present on main.

Fixing this means adding a partitions field to EmptyExecNode (and PlaceholderRowExecNode) and calling .with_partitions(...) on decode. A uint32 partitions field defaulting to 0 on absent can be mapped to 1 for backward compatibility with already-encoded plans.

This overlaps #23501 (migrating EmptyExec / PlaceholderRowExec to the try_to_proto / try_from_proto pattern under #23494), which currently specifies "schema only" and "keep the wire format byte-for-byte identical" — that migration would preserve this bug, so it may be worth folding the field addition into that work.

[milenkovicm]

can this be back-ported to 54.1 ?