refactor(proto): put Partitioning / sort-expression serde on the types - #24003
refactor(proto): put Partitioning / sort-expression serde on the types#24003adriangb wants to merge 3 commits into
Conversation
github-actions
Bot
added
physical-expr
adriangb
changed the title
There was a problem hiding this comment.
Pull request overview
Note
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Refactors protobuf (de)serialization for Partitioning and flat PhysicalSortExprNode orderings to live on/near the owning types, so plan-level hooks and datafusion-proto can reuse shared conversion logic.
Changes:
- Adds
try_to_proto/try_from_protoonPartitioningandPhysicalSortExpr, plus sequence helpers for sort expressions. - Updates multiple physical plan nodes and
datafusion-prototo route ordering/partitioning serde through the new shared APIs. - Introduces expression-level encode/decode context bridges (
expr_ctx) on execution plan proto contexts.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| datafusion/proto/src/physical_plan/to_proto.rs | Switches partitioning serialization to Partitioning::try_to_proto via an expression-level encode ctx. |
| datafusion/proto/src/physical_plan/from_proto.rs | Switches partitioning parsing to Partitioning::try_from_proto via an expression-level decode ctx. |
| datafusion/physical-plan/src/windows/proto.rs | Replaces hand-rolled window order_by sort-expr encoding/decoding with shared helpers. |
| datafusion/physical-plan/src/repartition/mod.rs | Routes RepartitionExec partitioning serde through Partitioning::{try_to_proto, try_from_proto}. |
| datafusion/physical-plan/src/proto.rs | Adds expr_ctx() adapters and implements PhysicalExprEncode/Decode for plan proto contexts. |
| datafusion/physical-plan/src/joins/symmetric_hash_join.rs | Replaces manual sort-expr serde with shared sort-expr sequence helpers. |
| datafusion/physical-plan/src/aggregates/mod.rs | Replaces manual aggregate ordering requirement serde with shared sort-expr helpers. |
| datafusion/physical-expr/src/partitioning.rs | Adds Partitioning protobuf conversions + tests; adds partition_count helper for safe decode. |
| datafusion/physical-expr-common/src/sort_expr.rs | Adds PhysicalSortExpr protobuf conversions and sequence encode/decode helpers. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| let partition_method = match self { | ||
| Partitioning::RoundRobinBatch(n) => { | ||
| protobuf::partitioning::PartitionMethod::RoundRobin(*n as u64) | ||
| } | ||
| Partitioning::Hash(exprs, n) => { | ||
| protobuf::partitioning::PartitionMethod::Hash( | ||
| protobuf::PhysicalHashRepartition { | ||
| hash_expr: ctx.encode_children_expressions(exprs)?, | ||
| partition_count: *n as u64, | ||
| }, | ||
| ) | ||
| } |
There was a problem hiding this comment.
Fixed. Encode now goes through a wire_partition_count helper (u64::try_from, internal error on overflow), mirroring the partition_count narrowing on the decode side, so an out-of-range count is an error in both directions rather than a silent truncation. Applied to all three sites (round-robin, hash, unknown).
| Partitioning::UnknownPartitioning(n) => { | ||
| protobuf::partitioning::PartitionMethod::Unknown(*n as u64) | ||
| } | ||
| }; |
There was a problem hiding this comment.
Same as the sibling comment above — both encode sites now use the checked wire_partition_count helper.
| /// Protobuf conversions for [`Partitioning`]. | ||
| /// | ||
| /// Child expressions (hash keys, range orderings) and `ScalarValue` split | ||
| /// points are (de)serialized through the expression-level context, so this is | ||
| /// the single copy of the partitioning wire format: `RepartitionExec`, | ||
| /// `FileScanConfig` and `datafusion-proto`'s central serializer all route | ||
| /// through it. |
There was a problem hiding this comment.
Reworded — FileScanConfig is now described as a next-step consumer (#23683), not a current one: the doc says RepartitionExec and datafusion-proto's central serializer route through this, and the remaining per-plan migrations are meant to do the same rather than grow another copy.
| /// An expression-level encode context backed by this plan context. | ||
| /// | ||
| /// Lets a plan hand `ctx` to expression-level conversions that own their own | ||
| /// wire logic — e.g. | ||
| /// [`Partitioning::try_to_proto`](datafusion_physical_expr::Partitioning::try_to_proto) | ||
| /// and | ||
| /// [`PhysicalSortExpr::try_to_proto`](datafusion_physical_expr::PhysicalSortExpr::try_to_proto). |
There was a problem hiding this comment.
No change needed here: PhysicalSortExpr is re-exported from datafusion_physical_expr (physical-expr/src/lib.rs), so the link resolves. Verified with RUSTDOCFLAGS="-D rustdoc::broken_intra_doc_links" cargo doc --no-deps -p datafusion-physical-plan --features proto — clean.
| /// Lets [`ExecutionPlanEncodeCtx`] back an [`PhysicalExprEncodeCtx`], so | ||
| /// expression-level conversions can be reused from plan hooks. |
There was a problem hiding this comment.
Fixed — now reads "back a PhysicalExprEncodeCtx".
| /// Lets [`ExecutionPlanDecodeCtx`] back a [`PhysicalExprDecodeCtx`], so | ||
| /// expression-level conversions can be reused from plan hooks. |
There was a problem hiding this comment.
This one already reads "back a PhysicalExprDecodeCtx" — fixed the encode-side twin, which did say "an".
| input_schema, | ||
| &decoder, | ||
| ); | ||
| partitioning |
There was a problem hiding this comment.
Could we make it delegate to the new shared decoder, or use checked conversion, so direct callers do not panic on oversized counts on 32-bit targets?
There was a problem hiding this comment.
Good catch, done. parse_protobuf_hash_partitioning now delegates to the shared Partitioning::try_from_proto. I also made the symmetric fix on the encode side: Partitioning::try_to_proto now widens through a checked wire_partition_count helper instead of as u6
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #24003 +/- ##
=======================================
Coverage 80.85% 80.85%
=======================================
Files 1099 1099
Lines 374338 374306 -32
Branches 374338 374306 -32
=======================================
- Hits 302671 302649 -22
Misses 53621 53621
+ Partials 18046 18036 -10 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
99ea991 to
ee3f054
Compare
`Partitioning`'s protobuf conversion existed in three places: inline in
`RepartitionExec::try_to_proto` / `try_from_proto`, and in
`datafusion-proto`'s `serialize_partitioning` /
`parse_protobuf_partitioning`. Each new plan or data source that needs to
encode an output partitioning has so far copied it again.
Put the single copy next to the types that own it, taking the
expression-level context (`datafusion-physical-expr{,-common}` already
carry the `proto` feature):
- `PhysicalSortExpr::try_to_proto` / `try_from_proto` (physical-expr-common)
- `Partitioning::try_to_proto` / `try_from_proto` (physical-expr)
So plan hooks can reach them, `ExecutionPlanEncodeCtx` /
`ExecutionPlanDecodeCtx` now back the expression-level contexts and hand
one out via `expr_ctx()`; any plan hook can pass its ctx to
expression-level conversions from here on. `RepartitionExec` and
`datafusion-proto`'s central serializer both route through the type
methods, which retires `serialize_range_partitioning`,
`serialize_range_split_point`, `parse_protobuf_range_partitioning` and
`parse_protobuf_range_split_point`.
The wire format is unchanged. Behavior differences: out-of-range
partition counts now error instead of wrapping or panicking on an
`unwrap`, and a missing sort-expression child reports which field is
missing rather than "Unexpected empty physical expression".
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`PhysicalSortExpr::try_to_proto` / `try_from_proto` convert one sort expression, but every caller holds a sequence of them, so each one hand-rolls the same map/collect over the flat `PhysicalSortExprNode` list: `AggregateExec`'s ordering requirement, `SymmetricHashJoinExec`'s left/right sort expressions, the window expressions, range partitioning, and (in flight) `FileScanConfig`'s output orderings and the file sinks' required ordering. Add `sort_exprs_try_to_proto` / `sort_exprs_try_from_proto` next to `PhysicalSortExpr` and route those callers through them. The encode side takes anything that borrows a `PhysicalSortExpr`, so a `LexOrdering`, a `&[PhysicalSortExpr]` and a `LexRequirement` mapped through `PhysicalSortExpr::from` all work. The decode side returns the expressions rather than a `LexOrdering`, because callers disagree on what an empty list means: "no ordering declared" for a scan, an error for an operator that requires one. Decode now reports the missing inner expression uniformly as "PhysicalSortExpr is missing required field 'expr'", replacing three bespoke messages. No wire-format change.
- `parse_protobuf_hash_partitioning` delegates to the shared `Partitioning::try_from_proto` instead of keeping its own copy of the hash wire format, so its `try_into().unwrap()` on the partition count can no longer panic for direct callers on a 32-bit target. Covered by a new unit test, since the function has no in-tree callers left. - Encode validates the partition count with `u64::try_from` rather than `as u64`, mirroring the decode side instead of truncating silently. - Reword the `Partitioning` proto docs: `FileScanConfig` is an intended consumer, not a current one. - Grammar: "a `PhysicalExprEncodeCtx`". Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ee3f054 to
4a85a4d
Compare
Which issue does this PR close?
try_to_proto/try_from_protomigration). Precursor for Proto: add DataSource/FileSource try_to_proto hook + port FileScanConfig serde #23497 / Add DataSource/FileSource proto hooks and FileScanConfig serde #23683 / Proto: add DataSink try_to_proto hook + port FileSinkConfig serde #23498 and for the remaining per-plan migrations that need to encode an output partitioning or an ordering.Rationale for this change
Partitioning's protobuf conversion currently exists in three places:RepartitionExec::try_to_proto,RepartitionExec::try_from_proto,datafusion-proto'sserialize_partitioning/parse_protobuf_partitioning.Every plan or data source migrated to the hooks that has an output
partitioning has so far copied it again — #23683 is about to add a fourth
copy for
FileScanConfig.The flat
PhysicalSortExprNodeencoding is the same story, one level down andmore widespread:
AggregateExec's ordering requirement,SymmetricHashJoinExec'sleft and right sort expressions, the window expressions, and range partitioning
each hand-roll the same
map/collectoverPhysicalSortExprNode { expr, asc, nulls_first }, on both the encode and thedecode side. #23683 (
output_ordering) and #23752 / #23781 (the sinks' requiredordering) are about to add two more.
There is no reason for this logic to live in the callers: it converts a
Partitioning(and thePhysicalSortExprs andScalarValues inside it), andneeds nothing from the plan level beyond the ability to encode a child
expression.
What changes are included in this PR?
Put the single copy next to the types that own it, taking the expression-level
context (
datafusion-physical-exprand-commonalready carry theprotofeature):
PhysicalSortExpr::try_to_proto/try_from_proto(physical-expr-common)sort_exprs_try_to_proto/sort_exprs_try_from_proto(physical-expr-common),the sequence form every caller actually needs
Partitioning::try_to_proto/try_from_proto(physical-expr)So plan hooks can reach them,
ExecutionPlanEncodeCtx/ExecutionPlanDecodeCtxnow back the expression-level contexts (
PhysicalExprEncode/PhysicalExprDecode) and hand one out viaexpr_ctx(). That bridge is usefulbeyond partitioning: from here on any plan hook can pass its ctx straight to an
expression-level conversion, which is the shape the rest of the migration wants.
The sequence encoder is generic over
Borrow<PhysicalSortExpr>, so one functionserves a
LexOrdering, a&[PhysicalSortExpr], and aLexRequirementmappedthrough
PhysicalSortExpr::from. The decoder returns the expressions ratherthan a
LexOrdering, because callers disagree on what an empty list means: "noordering declared" for a scan, an error for an operator that requires one.
Routed through the new methods:
RepartitionExecanddatafusion-proto's central serializer, which alsoretires
serialize_range_partitioning,serialize_range_split_point,parse_protobuf_range_partitioningandparse_protobuf_range_split_point,AggregateExec's ordering requirement,SymmetricHashJoinExec's left/rightsort expressions, and the window expressions — encode and decode each.
Net: −380 / +458 lines with the new tests included; production code shrinks. The
next operator that needs partitioning or ordering serde writes one line instead
of sixty.
Are these changes tested?
Yes.
LexRequirementinput, encode-error propagation, missing inner expression),using the existing
proto_test_utilstubs.path: all four partitioning variants (
datafusion-proto'sroundtrip_physical_plantests exercise the central serializer,RepartitionExec's own hook tests exercise the plan path), plus aggregateORDER BY, windowORDER BYand symmetric-hash-join sort expressions for theordering helpers.
datafusion-proto: 209 passed / 0 failed.(
cargo test --profile ci --workspace --lib --tests --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption).The only failures are 8 backtrace-symbolization tests in
datafusion-commonand 4
datafusion-clitests that hard-code a repo-relativeparquet-testing/path; both are artifacts of running from a linked worktree on macOS, in crates
this PR does not touch, and the
datafusion-clifour pass once that pathresolves.
cargo fmtandci/scripts/rust_clippy.shclean.Are there any user-facing changes?
The protobuf wire format is unchanged.
Additive API:
Partitioning::try_to_proto/try_from_proto,PhysicalSortExpr::try_to_proto/try_from_proto,sort_exprs_try_to_proto/sort_exprs_try_from_proto(feature
proto).ExecutionPlanEncodeCtx::expr_ctx()/ExecutionPlanDecodeCtx::expr_ctx(schema).Two behavior differences, both in error paths:
(
as usize) or panicking (try_into().unwrap()inparse_protobuf_hash_partitioning).(
PhysicalSortExpr is missing required field 'expr') instead ofUnexpected empty physical expression, and the same message now replaces thethree bespoke ones in
AggregateExec,SymmetricHashJoinExecand the windowexpressions.
Four private helpers in
datafusion-protoare removed(
serialize_range_partitioning,serialize_range_split_point,parse_protobuf_range_partitioning,parse_protobuf_range_split_point); thepublic
serialize_partitioning/parse_protobuf_partitioningkeep theirsignatures and behavior.