Skip to content

Commit e732212

Browse files
adriangbclaude
andauthored
feat: allow extension ExecutionPlans to decode via a per-type registry instead of PhysicalExtensionCodec (#24631)
## Which issue does this PR close? - Closes #24625. ## Rationale for this change ### The user-visible problem A library that ships its own `ExecutionPlan` (datafusion-distributed ships six) has to write a `PhysicalExtensionCodec` to serialize it: a central `downcast_ref` chain for encode and a matching `match` for decode, in a file separate from the plan. Two such libraries on one session also have to compose their codecs by hand through `ComposedPhysicalExtensionCodec`, and their users have to remember to install the composed codec on every process that decodes a plan. That composition is fragile in a way that is easy to miss. `PhysicalExtensionNode` carries no type discriminator, so the codec *is* the discriminator. `ComposedPhysicalExtensionCodec` copes by writing the **position** of the encoding codec into the payload bytes. The decoding side must therefore register the same codecs in the same order, a plan that two codecs both accept resolves to whichever was registered first, and a name collision between two independent crates cannot be detected. datafusion-python and datafusion-distributed have each built their own workaround for this (a codec-id layer and a `user_codec.rs` composition file respectively). ### Where this sits in the bigger picture Physical-plan serialization is moving from central `downcast_ref` chains in `datafusion-proto` to per-type hooks that live next to each type: `try_to_proto` (a trait method, dispatched through the vtable) plus an inherent `try_from_proto` (a constructor, so it cannot be a trait method). #23494 did this for every built-in `ExecutionPlan`; #22418 did it for `PhysicalExpr`. Third-party types could use the encode half of that but not the decode half, because decoding an extension node always routed through `PhysicalExtensionCodec::try_decode`. This PR closes the loop for `ExecutionPlan`; #24628 does the same for `PhysicalExpr` with the same shape and the same names (`Extension*FromProto` pairing the wire name with the constructor, `register_*` / `decode_*` facades over one shared registry, attached through `SessionConfig::with_extension`). **#24628 is stacked on this PR**, which adds the shared registry, and must land second. ### How this relates to `PhysicalExtensionCodec` The codec is **not deprecated** and this PR does not touch `ComposedPhysicalExtensionCodec`. The decode rule is: - A `PhysicalExtensionNode` **with a `plan_name` that the decoding session has registered** decodes through the plan's own `try_from_proto`. A failure there is fatal: falling through to the codec would let it wrongly decode a payload it never wrote, which is the bug the name exists to prevent. - **Anything else** (no name, because a codec wrote it, or a name this session does not know) takes the codec chain exactly as before. If the codec also fails, the missing registration is added as *context* on the codec's own error, and lists what the session does know. The codec's error kind is kept, because a codec can fail for a reason that has nothing to do with a missing registration. The codec keeps three jobs the registry does not take: UDF/UDAF/UDWF payloads (a function is an *instance*, not a type, so a type-keyed registry does not fit), anything not yet migrated, and the released `FFI_PhysicalExtensionCodec` in `datafusion-ffi`, which datafusion-python exposes as public API. What becomes deprecable once plans and expressions are registry-backed is `ComposedPhysicalExtensionCodec`, whose reason to exist was composition; that is a later, separate decision. Making the composed codec itself name-keyed instead of position-keyed (as suggested on the issue) is also orthogonal to this PR and can be done independently. ### How FFI will work Today `FFI_ExecutionPlan` has no serialization slot, so a foreign plan always goes through the codec. This PR does not change that. It is designed so FFI can plug in later with no breaking change: - The registry stores decoders as a **private** `fn` pointer. The public surface is `ExtensionPlanFromProto` and the three facade functions. The storage can grow to carry an FFI vtable and private data, and a lower-level entry point that takes a name and a decoder object can be added, both without breaking anything shipped here. - `PhysicalPlanNode` is a prost message, so it crosses the ABI as bytes, the same way `FFI_ExecutionPlan::partition_statistics` already ships a prost-encoded `Statistics`. Protobuf's field-skipping gives version-skew tolerance that a `#[repr(C)]` struct cannot. - One limit FFI inherits from the codec rather than from this PR: `ForeignPhysicalExtensionCodec::try_decode` discards the caller's `TaskContext` and rebuilds one from the provider's side, so a decoder that reads the *decoding* session's state (the second test below) would not see it across FFI. ## What changes are included in this PR? Three stacked commits, each green on its own. Built-in plan serialization is untouched: they keep the inherent `try_from_proto` they have had since 55.0.0, dispatched from their own `PhysicalPlanType` variant. 1. **`feat: add an optional plan_name discriminator to PhysicalExtensionNode`.** `optional string plan_name = 3;` — additive on the binary wire format: old writers omit it and old readers skip it. The generated JSON codec is stricter, see "user-facing changes" below. Nothing reads it yet. 2. **`feat: add ProtoDecoderRegistry, one decoder store for every extension kind`.** A `(trait, name)`-keyed store of type-erased decoders in `datafusion-proto-models`, with no knowledge of any decode context. Nothing uses it yet. 3. **`feat: decode extension ExecutionPlans through a session-scoped registry`.** - An extension plan implements one trait, `ExtensionPlanFromProto`, which pairs the name it is dispatched by with the constructor that rebuilds it. Built-ins do not implement it — they have a wire variant of their own — so `register_execution_plan::<FilterExec>(..)` does not compile, and **no built-in code or caller changes at all**. Its `try_to_proto` writes itself with `ctx.extension_node::<Self>(payload, self.children())`, which builds the `PhysicalExtensionNode` and stamps `Self::NAME`, so the name on the wire and the registry key cannot drift apart. Its `try_from_proto` receives the whole `PhysicalPlanNode`, matches the `Extension` variant, and decodes `inputs` with `ctx.decode_children`. - `ProtoDecoderRegistry` (new, in `datafusion-proto-models`): build one, `register_execution_plan::<T>(&mut registry)` for every extension plan the session must decode, and attach it with the existing `SessionConfig::with_extension`. Registering the same type twice is a no-op (keyed by `TypeId`); a *different* type under a taken name is an error, so collisions surface at registration. `decode_execution_plan` reads the name off the node rather than taking it as an argument, so a caller cannot pair a node with a name it does not carry. - The registry is **keyed by the trait as well as the name**, and is shared by every extension kind DataFusion can decode. This crate contributes only the `dyn ExecutionPlan` face of it. A session therefore carries **one** registry rather than one per kind, a plan and an expression may use the same name without colliding, and #24628 registers into the same object. It lives in `datafusion-proto-models` because that crate already sits below every crate that owns one of these traits, behind the same `proto` feature; it stores each decoder type-erased and never names a decode context. - A session carries one registry and `with_extension` replaces it, so the **application** that owns the session builds it. A library exposes `pub fn register(registry: &mut ProtoDecoderRegistry) -> Result<()>` — which can register its plans *and* its expressions — and the application composes the libraries it uses. The libraries need to know nothing about each other and the call order does not matter. There is no `SessionConfig` extension trait and no merge-on-register helper: public API is added when a use needs it, not ahead of one. - It is attached through the `SessionConfig` extension map because `datafusion-execution` sits *below* the crates that own these traits and cannot name them; the extension map is the one type-erased hole through that layering, and it reaches every `TaskContext` derived from the session with no new plumbing. - Decode dispatch in `datafusion-proto`, per the rule above, plus the upgrade-guide entry (in 56.0.0, the guide for the next release). ## What is the testing strategy for this PR? 12 unit tests across `datafusion-proto-models` and `datafusion-physical-plan`, 10 integration tests in `datafusion/proto/tests/cases/plans/extensions.rs`, and doctests on the public traits: - `ProtoDecoderRegistry` itself: 6 unit tests, including the same name under two different traits resolving to two different decoders. - Registration: idempotent for the same type, an error for a different type under the same name, empty names rejected, a name registered for another kind invisible to the `ExecutionPlan` face, and a `compile_fail` doctest for `register_execution_plan::<FilterExec>(..)` (verified to fail on the `ExtensionPlanFromProto` bound, not on anything incidental). - A runnable example, `datafusion-examples/examples/proto/extension_plan_registry.rs`: the same two-crate tree `composed_extension_codec.rs` builds, decoded by name with no composed codec, plus the error a session that forgot to register gets. - A plan round-tripping with `DefaultPhysicalExtensionCodec`, which decodes nothing, so whatever survives came from the registry. - A **session-dependent** plan modeled on datafusion-distributed's `NetworkShuffleExec`: its worker pool is rebuilt from the *decoding* session through `ExecutionPlanDecodeCtx::task_ctx()`, and never rides the wire. - Two independent crates' plans nested inside each other, decoding with no `ComposedPhysicalExtensionCodec`. - The registry taking precedence over a codec that claims every payload, and a registered decoder's failure *not* falling back to that codec. - A codec failing for an unrelated reason: the missing-registration hint is added, and the codec's own error kind still reaches `find_root()`. - Back-compat both ways: codec-encoded plans stay anonymous on the wire and decode as before; a name the session does not know (on a session that *does* have a registry with other plans in it) still decodes if a codec understands the payload; a node re-encoded without the field reads back as unnamed. - A bytes-level `encode_to_vec` / `decode` hop, so tag 3 is exercised as wire bytes and not just as a struct field. Verified locally on the final tree and with `cargo check` on every commit of the stack: the `datafusion-physical-expr`, `datafusion-physical-plan` and `datafusion-proto` test suites, the examples crate, CI's clippy script, rustdoc with private items under `-D warnings`, prettier, typos and fmt. ## Are there any user-facing changes? New public API, all opt-in; no behavior change for existing code. One mechanical break, in the 56.0.0 upgrade guide: `PhysicalExtensionNode` gained a field, so exhaustive struct literals must set `plan_name` (`None` reproduces the previous behavior). Built-in plans are untouched — no signature, trait or import change for anyone calling `FilterExec::try_from_proto(..)`. Two things worth a reviewer's attention: - **JSON is not as forgiving as the binary format.** The generated `pbjson` deserializer rejects a field it does not know, so a *named* node written by a new writer cannot be read as JSON by a reader on an older DataFusion. The binary prost path skips the unknown field as usual, and a plan that does not opt in writes no name and is unaffected. If you serialize plans as JSON across versions, upgrade the readers before you opt a plan in. - **Opting a plan in changes what the codec fallback means for that plan.** Once its `try_to_proto` writes its own payload, a reader that fails to register the type hands that payload to the codec it removed, and a codec that still recognizes the plan may decode the new payload as the old message rather than fail. The upgrade guide says to register migrated plans everywhere their plans are read. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_017UzVkTrKmRRc2fKCdMkqBp --------- Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
1 parent 7f00b30 commit e732212

16 files changed

Lines changed: 2145 additions & 12 deletions

File tree

‎datafusion-examples/README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,7 @@ cargo run --example dataframe -- dataframe
171171
| ------------------------ | --------------------------------------------------------------------------------- | ----------------------------------------------------------------------------- |
172172
| composed_extension_codec | [`proto/composed_extension_codec.rs`](examples/proto/composed_extension_codec.rs) | Use multiple extension codecs for serialization/deserialization |
173173
| expression_deduplication | [`proto/expression_deduplication.rs`](examples/proto/expression_deduplication.rs) | Example of expression caching/deduplication using the codec decorator pattern |
174+
| extension_plan_registry | [`proto/extension_plan_registry.rs`](examples/proto/extension_plan_registry.rs) | Decode two crates' extension plans by name, with no composed codec |
174175

175176
## Query Planning Examples
176177

‎datafusion-examples/examples/custom_data_source/adapter_serialization.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -364,6 +364,11 @@ impl PhysicalProtoConverterExtension for AdapterPreservingCodec {
364364
PhysicalExtensionNode {
365365
node: payload_bytes,
366366
inputs: vec![inner_proto],
367+
// This converter recognizes its own nodes by the marker
368+
// inside the payload, so it needs no `plan_name`
369+
// discriminator. Plans that decode through the session
370+
// registry set it to their `PLAN_NAME` instead.
371+
plan_name: None,
367372
},
368373
)),
369374
});
Lines changed: 252 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,252 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
//! See `main.rs` for how to run it.
19+
//!
20+
//! The same problem as `composed_extension_codec.rs`, solved without composing
21+
//! codecs: two plans from two independent crates in one tree.
22+
//!
23+
//! ```text
24+
//! ParentExec (from crate A)
25+
//! ChildExec (from crate B)
26+
//! ```
27+
//!
28+
//! With `ComposedPhysicalExtensionCodec` the decoding side has to register the
29+
//! same codecs in the same *order* the encoder used, because the composed codec
30+
//! writes the position of the encoding codec into the payload. Here each plan
31+
//! names itself on the wire instead, so order does not matter and a name two
32+
//! crates both claim is an error at registration rather than a wrong decode.
33+
34+
use std::fmt::Formatter;
35+
use std::sync::Arc;
36+
37+
use datafusion::common::tree_node::TreeNodeRecursion;
38+
use datafusion::common::{Result, internal_err};
39+
use datafusion::execution::TaskContext;
40+
use datafusion::physical_expr::PhysicalExpr;
41+
use datafusion::physical_plan::proto::{
42+
ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx, ExtensionPlanFromProto,
43+
ProtoDecoderRegistry, register_execution_plan,
44+
};
45+
use datafusion::physical_plan::{
46+
ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
47+
ReplaceChildrenOptions, SendableRecordBatchStream,
48+
};
49+
use datafusion::prelude::{SessionConfig, SessionContext};
50+
use datafusion_proto::physical_plan::{AsExecutionPlan, DefaultPhysicalExtensionCodec};
51+
use datafusion_proto::protobuf::PhysicalPlanNode;
52+
use datafusion_proto::protobuf::physical_plan_node::PhysicalPlanType;
53+
54+
pub fn extension_plan_registry() -> Result<()> {
55+
let plan: Arc<dyn ExecutionPlan> = Arc::new(ParentExec {
56+
input: Arc::new(ChildExec {}),
57+
});
58+
59+
// Each crate exposes a function that fills a registry it is handed; the
60+
// application that owns the session composes them. The two crates need to
61+
// know nothing about each other, and the call order does not matter.
62+
let mut registry = ProtoDecoderRegistry::new();
63+
crate_a::register(&mut registry)?;
64+
crate_b::register(&mut registry)?;
65+
66+
let ctx = SessionContext::new_with_config(
67+
SessionConfig::new().with_extension(Arc::new(registry)),
68+
);
69+
70+
// `DefaultPhysicalExtensionCodec` decodes nothing at all, so whatever
71+
// survives the round trip came from the registry.
72+
let codec = DefaultPhysicalExtensionCodec {};
73+
let node = PhysicalPlanNode::try_from_physical_plan(Arc::clone(&plan), &codec)?;
74+
let decoded = node.try_into_physical_plan(ctx.task_ctx().as_ref(), &codec)?;
75+
76+
assert_eq!(format!("{plan:?}"), format!("{decoded:?}"));
77+
78+
// A session that does not register the plans cannot decode them, and the
79+
// error names what is missing rather than only what the codec said. Children
80+
// are decoded before their parent falls back to the codec, so the innermost
81+
// unregistered plan is the one reported:
82+
//
83+
// No decoder is registered for the extension ExecutionPlan
84+
// 'example-crate-b.ChildExec'. Register the plan in the
85+
// ProtoDecoderRegistry attached to the decoding session's
86+
// SessionConfig, ... Registered extension plans: none
87+
let bare = SessionContext::new();
88+
let err = node
89+
.try_into_physical_plan(bare.task_ctx().as_ref(), &codec)
90+
.expect_err("an unregistered plan must not decode");
91+
assert!(err.to_string().contains("No decoder is registered"));
92+
assert!(err.to_string().contains(ChildExec::NAME));
93+
94+
Ok(())
95+
}
96+
97+
/// What the crate that owns `ParentExec` would export.
98+
mod crate_a {
99+
use super::*;
100+
101+
pub fn register(registry: &mut ProtoDecoderRegistry) -> Result<()> {
102+
register_execution_plan::<ParentExec>(registry)
103+
}
104+
}
105+
106+
/// What the crate that owns `ChildExec` would export.
107+
mod crate_b {
108+
use super::*;
109+
110+
pub fn register(registry: &mut ProtoDecoderRegistry) -> Result<()> {
111+
register_execution_plan::<ChildExec>(registry)
112+
}
113+
}
114+
115+
#[derive(Debug)]
116+
struct ParentExec {
117+
input: Arc<dyn ExecutionPlan>,
118+
}
119+
120+
impl ExtensionPlanFromProto for ParentExec {
121+
// Namespaced with the owning crate, so a collision with another crate is
122+
// an error at registration rather than a wrong decode.
123+
const NAME: &'static str = "example-crate-a.ParentExec";
124+
125+
fn try_from_proto(
126+
node: &PhysicalPlanNode,
127+
ctx: &ExecutionPlanDecodeCtx<'_>,
128+
) -> Result<Arc<dyn ExecutionPlan>> {
129+
let Some(PhysicalPlanType::Extension(extension)) = &node.physical_plan_type
130+
else {
131+
return internal_err!("PhysicalPlanNode is not an Extension");
132+
};
133+
let mut children = ctx.decode_children(&extension.inputs)?;
134+
if children.len() != 1 {
135+
return internal_err!("ParentExec expects exactly one child");
136+
}
137+
Ok(Arc::new(ParentExec {
138+
input: children.remove(0),
139+
}))
140+
}
141+
}
142+
143+
#[derive(Debug)]
144+
struct ChildExec {}
145+
146+
impl ExtensionPlanFromProto for ChildExec {
147+
const NAME: &'static str = "example-crate-b.ChildExec";
148+
149+
fn try_from_proto(
150+
_node: &PhysicalPlanNode,
151+
_ctx: &ExecutionPlanDecodeCtx<'_>,
152+
) -> Result<Arc<dyn ExecutionPlan>> {
153+
Ok(Arc::new(ChildExec {}))
154+
}
155+
}
156+
157+
/// Both plans are serde-only: this example never executes them.
158+
macro_rules! serde_only_body {
159+
($name:ident) => {
160+
fn name(&self) -> &str {
161+
stringify!($name)
162+
}
163+
164+
fn properties(&self) -> &Arc<PlanProperties> {
165+
unreachable!("this example only serializes")
166+
}
167+
168+
fn apply_expressions(
169+
&self,
170+
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
171+
) -> Result<TreeNodeRecursion> {
172+
Ok(TreeNodeRecursion::Continue)
173+
}
174+
175+
fn with_new_children(
176+
self: Arc<Self>,
177+
children: Vec<Arc<dyn ExecutionPlan>>,
178+
) -> Result<Arc<dyn ExecutionPlan>> {
179+
self.replace_children(
180+
children,
181+
ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
182+
)
183+
}
184+
185+
fn execute(
186+
&self,
187+
_partition: usize,
188+
_context: Arc<TaskContext>,
189+
) -> Result<SendableRecordBatchStream> {
190+
internal_err!("{} is a serde-only example plan", stringify!($name))
191+
}
192+
193+
/// `extension_node` stamps `Self::NAME`, so the name on the wire and
194+
/// the registry key cannot drift apart. Neither plan has state of its
195+
/// own, hence the empty payload.
196+
fn try_to_proto(
197+
&self,
198+
ctx: &ExecutionPlanEncodeCtx<'_>,
199+
) -> Result<Option<PhysicalPlanNode>> {
200+
Ok(Some(ctx.extension_node::<Self>(vec![], self.children())?))
201+
}
202+
};
203+
}
204+
205+
impl DisplayAs for ParentExec {
206+
fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
207+
write!(f, "ParentExec")
208+
}
209+
}
210+
211+
impl ExecutionPlan for ParentExec {
212+
serde_only_body!(ParentExec);
213+
214+
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
215+
vec![&self.input]
216+
}
217+
218+
fn replace_children(
219+
self: Arc<Self>,
220+
mut children: Vec<Arc<dyn ExecutionPlan>>,
221+
_: ReplaceChildrenOptions,
222+
) -> Result<Arc<dyn ExecutionPlan>> {
223+
if children.len() != 1 {
224+
return internal_err!("ParentExec expects exactly one child");
225+
}
226+
Ok(Arc::new(ParentExec {
227+
input: children.remove(0),
228+
}))
229+
}
230+
}
231+
232+
impl DisplayAs for ChildExec {
233+
fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
234+
write!(f, "ChildExec")
235+
}
236+
}
237+
238+
impl ExecutionPlan for ChildExec {
239+
serde_only_body!(ChildExec);
240+
241+
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
242+
vec![]
243+
}
244+
245+
fn replace_children(
246+
self: Arc<Self>,
247+
_children: Vec<Arc<dyn ExecutionPlan>>,
248+
_: ReplaceChildrenOptions,
249+
) -> Result<Arc<dyn ExecutionPlan>> {
250+
Ok(self)
251+
}
252+
}

‎datafusion-examples/examples/proto/main.rs‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
//!
2222
//! ## Usage
2323
//! ```bash
24-
//! cargo run --example proto -- [all|composed_extension_codec|expression_deduplication]
24+
//! cargo run --example proto -- [all|composed_extension_codec|expression_deduplication|extension_plan_registry]
2525
//! ```
2626
//!
2727
//! Each subcommand runs a corresponding example:
@@ -32,9 +32,13 @@
3232
//!
3333
//! - `expression_deduplication`
3434
//! (file: expression_deduplication.rs, desc: Example of expression caching/deduplication using the codec decorator pattern)
35+
//!
36+
//! - `extension_plan_registry`
37+
//! (file: extension_plan_registry.rs, desc: Decode two crates' extension plans by name, with no composed codec)
3538
3639
mod composed_extension_codec;
3740
mod expression_deduplication;
41+
mod extension_plan_registry;
3842

3943
use datafusion::error::{DataFusionError, Result};
4044
use strum::{IntoEnumIterator, VariantNames};
@@ -46,6 +50,7 @@ enum ExampleKind {
4650
All,
4751
ComposedExtensionCodec,
4852
ExpressionDeduplication,
53+
ExtensionPlanRegistry,
4954
}
5055

5156
impl ExampleKind {
@@ -69,6 +74,9 @@ impl ExampleKind {
6974
ExampleKind::ExpressionDeduplication => {
7075
expression_deduplication::expression_deduplication()?
7176
}
77+
ExampleKind::ExtensionPlanRegistry => {
78+
extension_plan_registry::extension_plan_registry()?
79+
}
7280
}
7381
Ok(())
7482
}

0 commit comments

Comments
 (0)