Allow extension ExecutionPlans to decode via a per-type registry instead of PhysicalExtensionCodec

Is your feature request related to a problem or challenge?

#23494 moved built-in ExecutionPlan serialization onto per-type try_to_proto / try_from_proto hooks. The result is genuinely nicer to implement than the central downcast_ref chain: encode and decode live next to the plan, and private state stays private.

Third-party plans can't fully use it. ExecutionPlan::try_to_proto is callable by them, but there is no way back — decoding an extension node still routes through PhysicalExtensionCodec::try_decode.

The blocker is the wire type:

pub struct PhysicalExtensionNode {
    node: Vec<u8>,
    inputs: Vec<PhysicalPlanNode>,
}

There is no type discriminator, so the codec is the discriminator. ComposedPhysicalExtensionCodec copes by writing the position of the encoding codec into the payload bytes (DataEncoderTuple { encoder_position, blob }), and on encode it tries each codec in sequence until one accepts the plan. So the decoding side must 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' codecs is undetectable. datafusion-python and datafusion-distributed have each built their own workaround for this (a per-codec id layer, and a user_codec.rs composition file).

Describe the solution you'd like

Add an optional name to the wire type plus a per-type decoder registry:

optional string plan_name = 3; on PhysicalExtensionNode — additive and proto3-compatible: old writers omit it, old readers ignore it.
A session-scoped registry mapping that name to a decoder fn.
Decode rule: name present and registered → registry; otherwise → the existing codec chain, unchanged.

Per-plan opt-in with a codec fallback, which is the same migration shape #23494 already used for encode.

Three things make this cheaper than it sounds:

The decoder signature already exists and is uniform. try_from_proto is a plain inherent fn, not a trait method, so there's no Self-return object-safety problem:
```
pub fn try_from_proto(
    node: &PhysicalPlanNode,
    ctx: &ExecutionPlanDecodeCtx<'_>,
) -> Result<Arc<dyn ExecutionPlan>>
```
The registry is HashMap<String, fn(&PhysicalPlanNode, &ExecutionPlanDecodeCtx) -> Result<Arc<dyn ExecutionPlan>>>.
ExecutionPlanDecodeCtx already exposes task_ctx(), so session-dependent extension plans can decode through this path, not just self-contained ones.
It mirrors a policy already in the codebase. UDF decode already implements "payload → codec; else registry → codec fallback" (datafusion/proto/src/physical_plan/mod.rs). This is the same policy one layer up.

The encode side needs one helper — something like ctx.encode_extension(name, bytes, children) — otherwise every extension author hand-rolls the Extension wrapper and some will forget to recurse into inputs.

Registration would look like register_execution_plan::<MyExec>() on the session, sitting alongside the existing FunctionRegistry.

Worked example: datafusion-distributed

datafusion-distributed ships six extension plans and serializes every query plan across the network, so it exercises this path hard. Concretely, on DataFusion 55:

src/codec/distributed_codec.rs is 833 lines. It holds six downcast_ref arms in try_encode and a matching six-arm match in try_decode, for NetworkShuffleExec, NetworkCoalesceExec, NetworkBroadcastExec, BroadcastExec, ChildrenIsolatorUnionExec and SamplerExec. This is exactly the central dispatch chain [EPIC] Port ExecutionPlan serialization to try_to_proto / try_from_proto hooks #23494 set out to remove, just re-created downstream — every extension project rebuilds it.
src/codec/user_codec.rs is 33 lines that exist purely for composition. Its only job is accumulating a Vec<Arc<dyn PhysicalExtensionCodec>> into a ComposedPhysicalExtensionCodec so the library's own codec and the end user's codec can coexist. It exists only because codecs don't compose by name. A name-keyed registry deletes the file and makes collisions detectable at registration instead of resolving by ordering.
The public API gets smaller. DistributedExt::with_distributed_user_codec(MyCodec) — which users must remember to call on both the coordinator and every worker — becomes register_execution_plan::<MyExec>(), with no separate codec type to author at all.
The docs get shorter. The project's "distributing custom ExecutionPlans" guide spends the first of its three sections teaching users to write a PhysicalExtensionCodec before they can distribute anything.

Two details from that project worth noting as design validation:

Its plans are session-dependent: NetworkShuffleExec reconstructs a worker connection pool out of the TaskContext at decode time. ExecutionPlanDecodeCtx::task_ctx() already covers this, which is good evidence the registry path is viable for real-world extension plans.
It implements no expr or UDF codec methods at all, so a plan-only registry would take it completely off PhysicalExtensionCodec — no partial migration, no keeping a codec around for the leftovers.

Describe alternatives you've considered

Keep codecs and migrate only encode to try_to_proto. This is possible today, but it splits encode and decode across different files for the same plan while keeping the registration burden identical. Strictly worse than either endpoint.
A global static registry. Rejected: session-scoped matches the FunctionRegistry precedent and stays testable and multi-tenant-safe.

Additional context

Registered names should probably be namespaced (e.g. datafusion-distributed.NetworkShuffleExec) so collisions surface at registration rather than as a mis-decode.

This does not retire PhysicalExtensionCodec — extension PhysicalExprs and UDF payloads still need it, and the codec fallback stays for unmigrated plans regardless. Companion issue for the expression side: #24626

refactor: make the decoder storage private and `ExtensionPlanParts` a…

Just sumarizing here the API proposed in #24631:

Before

This is how today users can register their own codecs, and do a roundtrip over a custom plan.

struct MyExec;

struct MyExecCodec;

impl PhysicalExtensionCodec for MyExecCodec {
    fn try_encode(
        &self,
        node: Arc<dyn ExecutionPlan>,
        buf: &mut Vec<u8>,
        proto_converter: &dyn PhysicalProtoConverterExtension,
    ) -> Result<()> {
        todo!() // encode logic
    }

    fn try_decode(
        &self,
        buf: &[u8],
        inputs: &[Arc<dyn ExecutionPlan>],
        ctx: &TaskContext,
        proto_converter: &dyn PhysicalProtoConverterExtension,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        todo!() // decode logic
    }
}

fn main() {
    let session_config = SessionConfig::new();
    let session_state = SessionStateBuilder::new()
        .with_config(session_config)
        .build();
    let task_ctx = session_state.task_ctx();

    let plan = Arc::new(MyExec);

    let codec = ComposedPhysicalExtensionCodec::new(vec![Arc::new(MyExecCodec)]);

    let proto = PhysicalPlanNode::try_from_physical_plan_with_converter(
        plan,
        &codec,
        &DefaultPhysicalProtoConverter {},
    )?;
    let plan = proto.try_into_physical_plan_with_context(
        &PhysicalPlanDecodeContext::new(&task_ctx, &codec),
        &DefaultPhysicalProtoConverter {},
    )?;
}

After

This is the alternative based on a type registry based on the current state of #24631:

struct MyExec;

impl ExecutionPlan for MyExec {
    fn try_to_proto(
        &self,
        ctx: &ExecutionPlanEncodeCtx<'_>,
    ) -> Result<Option<PhysicalPlanNode>> {
        todo!() // encode logic
    }
}

impl ExtensionExecutionPlan for MyExec {
    const PLAN_NAME: &'static str = "my-crate.MyExec";

    fn try_from_proto(
        node: &PhysicalPlanNode,
        ctx: &ExecutionPlanDecodeCtx<'_>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        todo!() // decode logic
    }
}

fn main() {
    let session_config = SessionConfig::new().with_execution_plan::<MyExec>()?;
    let session_state = SessionStateBuilder::new()
        .with_config(session_config)
        .build();
    let task_ctx = session_state.task_ctx();

    let plan = Arc::new(MyExec);

    let proto = PhysicalPlanNode::try_from_physical_plan_with_converter(
        plan,
        &DefaultPhysicalExtensionCodec {},
        &DefaultPhysicalProtoConverter {},
    )?;
    let plan = proto.try_into_physical_plan_with_context(
        &PhysicalPlanDecodeContext::new(&task_ctx, &DefaultPhysicalExtensionCodec {}),
        &DefaultPhysicalProtoConverter {},
    )?;
}

Here are a couple of thoughts, subject to me having understood correctly the proposed API (which might not be the case):

The type-registry based API is slightly nicer to work with, although the impact in the API change does not seem very big. The most important thing probably is that it would allow to consumers to remove big downcast-branch functions into per-plan scoped implementations.
The new API seems to be coupled to protobuf, but the PhysicalExtensionCodec allows plain bytes, which opens the door for arbitrary encodings. In my selfish opinion this is not a big deal, but I wonder if there are people in the community relying on non-protobuf based encodings.
The concurrent presence of the two APIs for providing serde capabilities to plans (PhysicalExtensionCodec + type registry) is probably not very nice from a public API standpoint. My guess is that it'd make sense to deprecate PhysicalExtensionCodec if we want to pursue the type registry based API.

As a user of the PhysicalExtensionCodec + ComposedPhysicalExtensionCodec, I have to admit that it's actually not that bad, and it gets the job done with no measurable performance impact. I'd even say that a successor of ComposedPhysicalExtensionCodec that stores a registry of types rather than a vec of codecs to try sequentially has good options for bringing most of the benefits of this proposed API with a very scoped change.

This API change resolves this issue as well: #24331. I ran into this error in our production use case and had to basically take an educated guess as to what was happening. The current codec implementation today is brittle/confusing in two ways:

Codec order matters
When you encounter an error, it's hard to tell if you used the wrong codec or if there was a real decoding failure

feat: add an optional `plan_name` discriminator to `PhysicalExtension…

feat: add a session-scoped decode registry for extension `ExecutionPl…

feat: decode extension `ExecutionPlan`s through the session registry

docs: document the extension `ExecutionPlan` decode registry

refactor: make the decoder storage private and `ExtensionPlanParts` a…

Thanks Gabriel and Jay.

An update on what the API looks like now in the PRs:

impl ExecutionPlan for MyExec {
    fn try_to_proto(&self, ctx: &ExecutionPlanEncodeCtx<'_>) -> Result<Option<PhysicalPlanNode>> {
        Ok(Some(PhysicalPlanNode {
            physical_plan_type: Some(PhysicalPlanType::Extension(PhysicalExtensionNode {
                // Made up function that knows how to go from MyExec -> bytes.
                // Any encoding, as today; it doesn't have to be protobuf.
                node: encode_my_exec(self)?,
                inputs: ctx.encode_children(self.children())?,
                plan_name: Some(Self::NAME.to_string()),
            })),
        }))
    }
}

impl ExecutionPlanFromProto for MyExec {
    const NAME: &'static str = "my-crate.MyExec";

    fn try_from_proto(node: &PhysicalPlanNode, ctx: &ExecutionPlanDecodeCtx<'_>) -> Result<Arc<dyn ExecutionPlan>> {
        // The registry hands over the whole node, so the plan matches its own variant.
        let extension = expect_plan_variant!(node, PhysicalPlanType::Extension, "Extension");
        let children = ctx.decode_children(&extension.inputs)?;
        let [input] = children.try_into().map_err(|_| plan_datafusion_err!("{} expects one input", Self::NAME))?;
        // Made up function that knows how to go from bytes + children + session -> MyExec.
        // `ctx.task_ctx()` gives the decoding session, for a plan that rebuilds
        // session-scoped state (a worker pool, a channel resolver) at decode time.
        let exec = decode_my_exec(extension.node.as_slice(), input, ctx.task_ctx())?;
        Ok(Arc::new(exec))
    }
}

let mut registry = ExecutionPlanRegistry::new();
registry.register::<MyExec>()?;
let config = SessionConfig::new().with_extension(Arc::new(registry));

Regarding your three points:

Protobuf coupling. The payload stays Vec<u8>. Only the envelope is protobuf, and it already is today: a codec's bytes ride inside PhysicalExtensionNode.node either way. The only addition to the envelope is an optional string. A non-protobuf payload works on both paths.
Two APIs at once. Agreed, this is worth thinking about. The PR does not deprecate PhysicalExtensionCodec; it is kept around for three reasons, all of which can be resolved in time:
- Any plans or expressions not migrated yet (gradual rollout / backwards compatibility).
- UDF/UDAF/UDWF payloads (a function is an instance, not a type). UDFs are already resolved by name through the session's FunctionRegistry before the codec is consulted, so the codec only matters for UDFs that carry instance state. The end state is a self-serialization hook on the UDF trait plus a per-name decoder, keyed through the FunctionRegistry we already have. Not designed yet.
- FFI_PhysicalExtensionCodec, which datafusion-python exposes as public API. FFI_ExecutionPlan has no serialization slot today, so foreign plans always go through the codec. The registry stores decoders privately, so the storage can later become Arc<dyn Fn> with a register_decoder(name, ..) for foreign decoders, and PhysicalPlanNode already crosses the ABI as prost bytes. No breaking change needed to get there.

What the registry makes redundant is ComposedPhysicalExtensionCodec, whose reason to exist is composition. I would rather deprecate that once both plans and expressions can go through a registry, as a separate decision. Might have to wait for UDFs to be resolved as well.

A type-keyed successor to ComposedPhysicalExtensionCodec is close to what this PR is, with one difference: the discriminator is on the wire (plan_name), not inside a codec-private wrapper. That is what lets a reader distinguish "no decoder is registered for X" from "the decoder for X failed", which resolves the confusion described in #24331. It also means no codec type has to exist per plan at all. We could still do a name-keyed composed codec if there is an extended migration period, but I hope that is not necessary.

@jayshrivastava

@jayshrivastava I've marked the PRs (#24631, #24628) as ready - could you give them a high level read?

Left comments on #24631. Will add #24628 to my queue 🫡. Thanks for working on this - it will be nice to see.

feat: allow extension `ExecutionPlan`s to decode via a per-type regis…

Navigation Menu

Repository navigation

Allow extension `ExecutionPlan`s to decode via a per-type registry instead of `PhysicalExtensionCodec` #24625

Is your feature request related to a problem or challenge?

Describe the solution you'd like

Worked example: datafusion-distributed

Describe alternatives you've considered

Additional context

Before

After

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Allow extension ExecutionPlans to decode via a per-type registry instead of PhysicalExtensionCodec #24625

Description

Is your feature request related to a problem or challenge?

Describe the solution you'd like

Worked example: datafusion-distributed

Describe alternatives you've considered

Additional context

Activity

gabotechs commented on Sep 7, 2026

Before

After

jayshrivastava commented on Sep 9, 2026

adriangb commented on Sep 15, 2026

adriangb commented on Sep 15, 2026

jayshrivastava commented on Sep 16, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Allow extension `ExecutionPlan`s to decode via a per-type registry instead of `PhysicalExtensionCodec` #24625