fix: preserve EmptyExec and PlaceholderRowExec partition count across proto round-trip#23643
Conversation
EmptyExecNode and PlaceholderRowExecNode encoded only a schema, so the count set by with_partitions was silently dropped on a physical plan round-trip and the decoded plan reported a single partition. Add a partitions field to both messages, write it on encode, and apply it on decode. Plans encoded before the field existed carry no value, which decodes as 0 and maps back to the previous default of 1.
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
Adding the partitions field to these generated protobuf structs breaks downstream code that constructs them with an exhaustive struct literal, so document the change and migration in the 55.0.0 upgrade guide.
|
The semver check is reporting
This is a real, though unavoidable, break: the finding is inherent to adding a field to a prost-generated message, and the field is what fixes the bug. It only affects downstream code that constructs these two structs with an exhaustive struct literal; anyone using Marking these Per the API health policy, I've addressed it by:
Happy to reconsider if reviewers would rather see this deferred to a different release. |
| self | ||
| } | ||
|
|
||
| /// Number of partitions this plan produces |
There was a problem hiding this comment.
Can we add a note that this method is being made public only for our proto serialization and we don't guarantee any sort of API stability, etc? I know it's improper, we can always decide later to respect public API contracts for this method. But I think it's best to do what we can to err on the side of caution now because we'll probably be removing it again soon.
There was a problem hiding this comment.
I removed these methods. They weren't needed since we already have output_partitioning that the tests can use. Thanks for catching that.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #23643 +/- ##
=======================================
Coverage ? 80.64%
=======================================
Files ? 1086
Lines ? 366132
Branches ? 366132
=======================================
Hits ? 295278
Misses ? 53259
Partials ? 17595 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Yes, I'd be happy to work on that |
…sors EmptyExec and PlaceholderRowExec mirror their private `partitions` field into PlanProperties as UnknownPartitioning(n), so the encoder can read the count through the existing ExecutionPlan trait rather than through newly added public accessors. Drop `EmptyExec::partitions()` and `PlaceholderRowExec::partitions()`, which now have no callers, and read the count via `properties().output_partitioning().partition_count()` instead. This keeps the fix free of any new public API in datafusion-physical-plan. The proto `partitions` field and its semver impact are unaffected.
adriangb
left a comment
There was a problem hiding this comment.
Nice fix removing the new pub methods!
Which issue does this PR close?
Rationale for this change
EmptyExecNodeandPlaceholderRowExecNodeencoded only a schema, so the partition count set byEmptyExec::with_partitions(n)was silently lost across a physical-plan round-trip: a plan reportingnpartitions before serialization reported1after.This is silent data loss rather than an error — the decoded plan is well-formed but describes a different plan than the one encoded. It bites any setup that plans in one process and executes in another. In Ballista, an optimizer rule collapses provably-empty sub-plans into an
EmptyExecinheriting the replaced node's partition count; the scheduler sizes the stage's task count from that count, ships the plan to an executor, and every task above partition 0 fails with:What changes are included in this PR?
uint32 partitionsfield toEmptyExecNodeandPlaceholderRowExecNodeindatafusion.proto, and regenerate the prost/pbjson code viaregen.sh.partitionson encode and apply it viawith_partitions(...)on decode.Both
EmptyExecandPlaceholderRowExecmirror their privatepartitionsfield intoPlanPropertiesasUnknownPartitioning(n), so the encoder reads the count viaproperties().output_partitioning().partition_count()on the existingExecutionPlantrait. No new public API is added todatafusion-physical-plan.The wire format stays compatible in both directions. Plans encoded before this field existed carry no value for it, which prost surfaces as
0; decode maps that to the previous default of1. Plans encoded after this change add a field that older readers skip.Note for #23501, which migrates these nodes to the
try_to_proto/try_from_protopattern: that issue specifies "schema only" and a byte-for-byte identical wire format, which would reintroduce this bug. Thepartitionsfield should be folded into that rewrite.Are these changes tested?
Yes, three new tests in
roundtrip_physical_plan.rs: partition-count round-trips forEmptyExecandPlaceholderRowExec, plus one decodingpartitions: 0nodes directly to pin the backward-compatibility mapping. The round-trip tests were confirmed to fail against the unfixed decoder with exactly the reported symptom (4 partitions decoding to 1).Are there any user-facing changes?
No API changes to
datafusion-physical-plan. Encoded plans gain a new protobuf field, which is backward and forward compatible as described above. The generatedEmptyExecNodeandPlaceholderRowExecNodestructs gain apartitionsfield, which breaks downstream code constructing them with an exhaustive struct literal; this is documented in the 55.0.0 upgrade guide.