jayshrivastava commented on code in PR #24631:
URL: https://github.com/apache/datafusion/pull/24631#discussion_r4030893711
##########
datafusion/proto/src/physical_plan/mod.rs:
##########
@@ -1393,20 +1399,58 @@ pub trait PhysicalPlanNodeExt: Sized {
ctx: &PhysicalPlanDecodeContext<'_>,
proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<Arc<dyn ExecutionPlan>> {
+ // Lookup-order policy, mirroring the one function decode already uses
a
+ // layer down (payload -> codec; else registry -> codec fallback): a
+ // plan that named itself on the wire and is registered on this session
+ // decodes itself, no codec involved. Anything else — an unnamed node
+ // from a codec-encoded writer, or a name this session does not know —
+ // takes the codec chain exactly as before.
+ let registry = ctx
+ .task_ctx()
+ .session_config()
+ .get_extension::<ExecutionPlanRegistry>();
+ if let Some(plan_name) = extension.plan_name.as_deref()
+ && let Some(registry) = registry.as_ref()
+ {
+ let plan_decoder = ConverterPlanDecoder {
+ ctx,
+ proto_converter,
+ };
+ // The decoder receives the whole node, like every built-in
+ // `try_from_proto`, and decodes its own children through the ctx.
+ // `None` here means no decoder claims the name; a decode *failure*
+ // is returned as-is rather than falling through to the codec.
+ if let Some(decoded) = registry.decode(
+ plan_name,
+ self.node(),
+ &ExecutionPlanDecodeCtx::new(&plan_decoder),
+ ) {
+ return decoded;
+ }
+ }
+
let inputs: Vec<Arc<dyn ExecutionPlan>> = extension
.inputs
.iter()
.map(|i| proto_converter.proto_to_execution_plan(i, ctx))
.collect::<Result<_>>()?;
- let extension_node = ctx.codec().try_decode(
- extension.node.as_slice(),
- &inputs,
- ctx.task_ctx(),
- proto_converter,
- )?;
-
- Ok(extension_node)
+ ctx.codec()
+ .try_decode(
+ extension.node.as_slice(),
+ &inputs,
+ ctx.task_ctx(),
+ proto_converter,
+ )
+ .map_err(|e| match extension.plan_name.as_deref() {
+ // The writer named the plan but this session has no decoder
for
+ // it: say so, rather than leaving only the codec's
"unsupported
+ // plan" error to explain a missing registration.
+ Some(plan_name) => {
+ unregistered_extension_plan_err(plan_name,
registry.as_deref(), &e)
Review Comment:
This looks like it classifies any codec error as a missing registration
error. Probably need a fix here.
##########
datafusion/physical-plan/src/proto/mod.rs:
##########
@@ -343,6 +359,14 @@ impl<'a> ExecutionPlanDecodeCtx<'a> {
self.decoder.decode_udwf(name, payload)
}
+ /// Deserialize a slice of child plans.
Review Comment:
nit: place this closer to `decode_child`
##########
datafusion/proto/src/physical_plan/mod.rs:
##########
@@ -1393,20 +1399,58 @@ pub trait PhysicalPlanNodeExt: Sized {
ctx: &PhysicalPlanDecodeContext<'_>,
proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<Arc<dyn ExecutionPlan>> {
+ // Lookup-order policy, mirroring the one function decode already uses
a
+ // layer down (payload -> codec; else registry -> codec fallback): a
+ // plan that named itself on the wire and is registered on this session
+ // decodes itself, no codec involved. Anything else — an unnamed node
+ // from a codec-encoded writer, or a name this session does not know —
+ // takes the codec chain exactly as before.
+ let registry = ctx
+ .task_ctx()
+ .session_config()
+ .get_extension::<ExecutionPlanRegistry>();
+ if let Some(plan_name) = extension.plan_name.as_deref()
+ && let Some(registry) = registry.as_ref()
+ {
+ let plan_decoder = ConverterPlanDecoder {
+ ctx,
+ proto_converter,
+ };
+ // The decoder receives the whole node, like every built-in
+ // `try_from_proto`, and decodes its own children through the ctx.
+ // `None` here means no decoder claims the name; a decode *failure*
+ // is returned as-is rather than falling through to the codec.
+ if let Some(decoded) = registry.decode(
+ plan_name,
+ self.node(),
+ &ExecutionPlanDecodeCtx::new(&plan_decoder),
+ ) {
+ return decoded;
+ }
+ }
+
let inputs: Vec<Arc<dyn ExecutionPlan>> = extension
.inputs
.iter()
.map(|i| proto_converter.proto_to_execution_plan(i, ctx))
.collect::<Result<_>>()?;
- let extension_node = ctx.codec().try_decode(
- extension.node.as_slice(),
- &inputs,
- ctx.task_ctx(),
- proto_converter,
- )?;
-
- Ok(extension_node)
+ ctx.codec()
Review Comment:
I suppose we need this fallback for backwards compatibility in distributed
contexts. Otherwise, it would be nice to remove it.
##########
datafusion/physical-plan/src/proto/mod.rs:
##########
@@ -366,6 +390,97 @@ impl PhysicalExprDecode for ExecutionPlanDecodeCtx<'_> {
}
}
+/// The decode half of [`ExecutionPlan::try_to_proto`]: a plan type that can
+/// be rebuilt from the `PhysicalPlanNode` its `try_to_proto` wrote.
+///
+/// Every self-serializing plan implements this. Built-in plans are dispatched
+/// to their `try_from_proto` by their `PhysicalPlanType` variant, so for them
+/// [`NAME`](Self::NAME) is informational. A third-party plan shares the single
+/// `PhysicalExtensionNode` variant with every other extension, so it is
+/// dispatched by `NAME` instead: its `try_to_proto` writes the name on the
+/// `PhysicalExtensionNode`, and it is registered in an
+/// [`ExecutionPlanRegistry`] attached to the session that will decode it:
+///
+/// ```
+/// # use std::any::Any;
+/// # use std::fmt::Formatter;
+/// # use std::sync::Arc;
+/// # use datafusion_common::Result;
+/// # use datafusion_execution::config::SessionConfig;
+/// # use datafusion_physical_plan::{DisplayAs, DisplayFormatType,
ExecutionPlan, PlanProperties, SendableRecordBatchStream};
+/// # use datafusion_physical_plan::proto::{
+/// # ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx,
ExecutionPlanRegistry,
+/// # ExecutionPlanFromProto,
+/// # };
+/// # use datafusion_proto_models::protobuf::PhysicalPlanNode;
+/// # #[derive(Debug)]
+/// # struct MyExec { properties: Arc<PlanProperties> }
+/// # impl DisplayAs for MyExec {
+/// # fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) ->
std::fmt::Result { write!(f, "MyExec") }
+/// # }
+/// # impl ExecutionPlan for MyExec {
+/// # fn name(&self) -> &str { "MyExec" }
+/// # fn properties(&self) -> &Arc<PlanProperties> { &self.properties }
+/// # fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { vec![] }
+/// # fn apply_expressions(&self, _f: &mut dyn FnMut(&Arc<dyn
datafusion_physical_expr::PhysicalExpr>) ->
Result<datafusion_common::tree_node::TreeNodeRecursion>) ->
Result<datafusion_common::tree_node::TreeNodeRecursion> {
Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue) }
+/// # fn with_new_children(self: Arc<Self>, _: Vec<Arc<dyn
ExecutionPlan>>) -> Result<Arc<dyn ExecutionPlan>> { Ok(self) }
+/// # fn execute(&self, _: usize, _:
Arc<datafusion_execution::TaskContext>) -> Result<SendableRecordBatchStream> {
unimplemented!() }
+/// # fn try_to_proto(&self, ctx: &ExecutionPlanEncodeCtx<'_>) ->
Result<Option<PhysicalPlanNode>> {
+/// # use datafusion_proto_models::protobuf::{PhysicalExtensionNode,
physical_plan_node::PhysicalPlanType};
+/// # Ok(Some(PhysicalPlanNode {
+/// # physical_plan_type:
Some(PhysicalPlanType::Extension(PhysicalExtensionNode {
+/// # node: vec![],
+/// # 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>> {
+/// let extension = datafusion_physical_plan::expect_plan_variant!(
+/// node,
+///
datafusion_proto_models::protobuf::physical_plan_node::PhysicalPlanType::Extension,
+/// "Extension",
+/// );
+/// let children = ctx.decode_children(&extension.inputs)?;
+/// // ... rebuild `MyExec` from `extension.node` (the payload) and
`children` ...
+/// # unimplemented!()
+/// }
+/// }
+///
+/// let mut registry = ExecutionPlanRegistry::new();
+/// registry.register::<MyExec>()?;
+/// let config = SessionConfig::new().with_extension(Arc::new(registry));
Review Comment:
Is there any situation where we want to support multiple registries? I
wonder if the canonical thing to do would be to `get_extension` first and add
to that registry. Otherwise, add your own registry.
##########
docs/source/library-user-guide/upgrading/56.0.0.md:
##########
@@ -357,3 +357,75 @@ let description =
ChildFilterDescription::from_child_with_column_mapping(
&child,
)?;
```
+
+### `try_from_proto` on built-in `ExecutionPlan`s is a trait item
+
+Every built-in plan that serializes itself already had an inherent
`try_from_proto(node, ctx)` with the same signature, dispatched by name from
`datafusion-proto`. That informal contract is now the `ExecutionPlanFromProto`
trait (with a `NAME` per type), and the built-in plans implement it instead.
Bodies and signatures are unchanged.
+
+**Who is affected:** code that calls `FilterExec::try_from_proto(..)` (or any
other built-in's) must import the trait: `use
datafusion::physical_plan::proto::ExecutionPlanFromProto;`.
+
+### Extension `ExecutionPlan`s can decode without a `PhysicalExtensionCodec`
+
+`ExecutionPlan::try_to_proto` let a plan serialize itself, but there was no
way back: decoding an extension node always routed through
`PhysicalExtensionCodec::try_decode`. `PhysicalExtensionNode` carried no type
discriminator, so the codec _was_ the discriminator.
`ComposedPhysicalExtensionCodec` works around that by writing the position of
the encoding codec into the payload, which means the decoding side must
register the same codecs in the same order, and a name collision between two
independent crates is undetectable.
+
+`PhysicalExtensionNode` now has an optional `plan_name`, and a session can map
that name to a decoder. An extension plan implements `ExecutionPlanFromProto`
and is registered on the session that will decode it:
+
+```rust,ignore
+use datafusion::physical_plan::proto::{
+ ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx, ExecutionPlanRegistry,
+ ExecutionPlanFromProto,
+};
+
+impl ExecutionPlan for MyExec {
+ // ...
+ fn try_to_proto(
+ &self,
+ ctx: &ExecutionPlanEncodeCtx<'_>,
+ ) -> Result<Option<PhysicalPlanNode>> {
+ // `Self::NAME` is also what the registry looks the plan up by, so the
+ // two cannot drift apart.
+ Ok(Some(PhysicalPlanNode {
+ physical_plan_type:
Some(PhysicalPlanType::Extension(PhysicalExtensionNode {
+ node: my_payload_bytes(self)?,
+ inputs: ctx.encode_children(self.children())?,
+ plan_name: Some(Self::NAME.to_string()),
Review Comment:
It would be nice to be able to avoid the case where someone accidentally
passes the wrong string.
I don't have many great ideas. Maybe we could make the 3rd parameter this
type where you have to call a `name` helper to get it. But that still requires
the user to call `name(self)`, which I suppose they could mess up.
```rust
type RegistryName = &'static str;
fn name<T: ExecutionPlanFromProto>(_: &T) -> RegistryName {
T::NAME
}
```
##########
datafusion/physical-plan/src/proto/mod.rs:
##########
@@ -366,6 +390,97 @@ impl PhysicalExprDecode for ExecutionPlanDecodeCtx<'_> {
}
}
+/// The decode half of [`ExecutionPlan::try_to_proto`]: a plan type that can
+/// be rebuilt from the `PhysicalPlanNode` its `try_to_proto` wrote.
+///
+/// Every self-serializing plan implements this. Built-in plans are dispatched
+/// to their `try_from_proto` by their `PhysicalPlanType` variant, so for them
+/// [`NAME`](Self::NAME) is informational. A third-party plan shares the single
+/// `PhysicalExtensionNode` variant with every other extension, so it is
+/// dispatched by `NAME` instead: its `try_to_proto` writes the name on the
+/// `PhysicalExtensionNode`, and it is registered in an
+/// [`ExecutionPlanRegistry`] attached to the session that will decode it:
+///
+/// ```
+/// # use std::any::Any;
+/// # use std::fmt::Formatter;
+/// # use std::sync::Arc;
+/// # use datafusion_common::Result;
+/// # use datafusion_execution::config::SessionConfig;
+/// # use datafusion_physical_plan::{DisplayAs, DisplayFormatType,
ExecutionPlan, PlanProperties, SendableRecordBatchStream};
+/// # use datafusion_physical_plan::proto::{
+/// # ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx,
ExecutionPlanRegistry,
+/// # ExecutionPlanFromProto,
+/// # };
+/// # use datafusion_proto_models::protobuf::PhysicalPlanNode;
+/// # #[derive(Debug)]
+/// # struct MyExec { properties: Arc<PlanProperties> }
+/// # impl DisplayAs for MyExec {
+/// # fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) ->
std::fmt::Result { write!(f, "MyExec") }
+/// # }
+/// # impl ExecutionPlan for MyExec {
+/// # fn name(&self) -> &str { "MyExec" }
+/// # fn properties(&self) -> &Arc<PlanProperties> { &self.properties }
+/// # fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { vec![] }
+/// # fn apply_expressions(&self, _f: &mut dyn FnMut(&Arc<dyn
datafusion_physical_expr::PhysicalExpr>) ->
Result<datafusion_common::tree_node::TreeNodeRecursion>) ->
Result<datafusion_common::tree_node::TreeNodeRecursion> {
Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue) }
+/// # fn with_new_children(self: Arc<Self>, _: Vec<Arc<dyn
ExecutionPlan>>) -> Result<Arc<dyn ExecutionPlan>> { Ok(self) }
+/// # fn execute(&self, _: usize, _:
Arc<datafusion_execution::TaskContext>) -> Result<SendableRecordBatchStream> {
unimplemented!() }
+/// # fn try_to_proto(&self, ctx: &ExecutionPlanEncodeCtx<'_>) ->
Result<Option<PhysicalPlanNode>> {
+/// # use datafusion_proto_models::protobuf::{PhysicalExtensionNode,
physical_plan_node::PhysicalPlanType};
+/// # Ok(Some(PhysicalPlanNode {
+/// # physical_plan_type:
Some(PhysicalPlanType::Extension(PhysicalExtensionNode {
+/// # node: vec![],
+/// # 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>> {
+/// let extension = datafusion_physical_plan::expect_plan_variant!(
+/// node,
+///
datafusion_proto_models::protobuf::physical_plan_node::PhysicalPlanType::Extension,
+/// "Extension",
+/// );
+/// let children = ctx.decode_children(&extension.inputs)?;
+/// // ... rebuild `MyExec` from `extension.node` (the payload) and
`children` ...
+/// # unimplemented!()
+/// }
+/// }
+///
+/// let mut registry = ExecutionPlanRegistry::new();
+/// registry.register::<MyExec>()?;
+/// let config = SessionConfig::new().with_extension(Arc::new(registry));
Review Comment:
Or, why not always have a registry already in the default `SessionConfig`
when the proto feature is enabled? That way, one would always `get_extension`
and add their plan node codecs to it.
##########
docs/source/library-user-guide/upgrading/56.0.0.md:
##########
@@ -357,3 +357,75 @@ let description =
ChildFilterDescription::from_child_with_column_mapping(
&child,
)?;
```
+
+### `try_from_proto` on built-in `ExecutionPlan`s is a trait item
+
+Every built-in plan that serializes itself already had an inherent
`try_from_proto(node, ctx)` with the same signature, dispatched by name from
`datafusion-proto`. That informal contract is now the `ExecutionPlanFromProto`
trait (with a `NAME` per type), and the built-in plans implement it instead.
Bodies and signatures are unchanged.
+
+**Who is affected:** code that calls `FilterExec::try_from_proto(..)` (or any
other built-in's) must import the trait: `use
datafusion::physical_plan::proto::ExecutionPlanFromProto;`.
+
+### Extension `ExecutionPlan`s can decode without a `PhysicalExtensionCodec`
+
+`ExecutionPlan::try_to_proto` let a plan serialize itself, but there was no
way back: decoding an extension node always routed through
`PhysicalExtensionCodec::try_decode`. `PhysicalExtensionNode` carried no type
discriminator, so the codec _was_ the discriminator.
`ComposedPhysicalExtensionCodec` works around that by writing the position of
the encoding codec into the payload, which means the decoding side must
register the same codecs in the same order, and a name collision between two
independent crates is undetectable.
+
+`PhysicalExtensionNode` now has an optional `plan_name`, and a session can map
that name to a decoder. An extension plan implements `ExecutionPlanFromProto`
and is registered on the session that will decode it:
+
+```rust,ignore
+use datafusion::physical_plan::proto::{
+ ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx, ExecutionPlanRegistry,
+ ExecutionPlanFromProto,
+};
+
+impl ExecutionPlan for MyExec {
+ // ...
+ fn try_to_proto(
+ &self,
+ ctx: &ExecutionPlanEncodeCtx<'_>,
+ ) -> Result<Option<PhysicalPlanNode>> {
+ // `Self::NAME` is also what the registry looks the plan up by, so the
+ // two cannot drift apart.
+ Ok(Some(PhysicalPlanNode {
+ physical_plan_type:
Some(PhysicalPlanType::Extension(PhysicalExtensionNode {
+ node: my_payload_bytes(self)?,
+ inputs: ctx.encode_children(self.children())?,
+ plan_name: Some(Self::NAME.to_string()),
+ })),
+ }))
+ }
+}
+
+impl ExecutionPlanFromProto for MyExec {
+ // Namespace the name so a collision with another crate is an error at
+ // registration rather than a wrong decode.
+ const NAME: &'static str = "my-crate.MyExec";
+
+ fn try_from_proto(
+ node: &PhysicalPlanNode,
+ ctx: &ExecutionPlanDecodeCtx<'_>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let extension = expect_plan_variant!(node,
PhysicalPlanType::Extension, "Extension");
+ let children = ctx.decode_children(&extension.inputs)?;
+ // `ctx.task_ctx()` is available here, so a plan that rebuilds session
+ // state at decode time (a connection pool, say) can do so.
+ my_plan_from_bytes(&extension.node, children)
+ }
+}
+
+// On every session that decodes the plan — for a distributed engine, the
+// workers as well as the coordinator.
+let mut registry = ExecutionPlanRegistry::new();
+registry.register::<MyExec>()?;
+let config = SessionConfig::new().with_extension(Arc::new(registry));
+```
+
+A registered plan needs no `PhysicalExtensionCodec` at all, and two crates'
plans coexist on one session without a `ComposedPhysicalExtensionCodec`.
+
+**Who is affected:**
+
+- Nobody is required to change anything. This is opt-in per plan type: a
`PhysicalExtensionNode` with no name, or with a name the decoding session does
not know, takes the `PhysicalExtensionCodec` chain exactly as before. The new
field is additive: old writers omit it and old readers ignore it.
Review Comment:
Codex tells me that this is not true for JSON codecs
##########
datafusion/physical-plan/src/proto/registry.rs:
##########
@@ -0,0 +1,226 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Name-keyed decode registry for extension [`ExecutionPlan`]s.
+//!
+//! Built-in plans are dispatched by their `PhysicalPlanType` oneof variant, so
+//! the wire format names them. Extension plans all share the single
+//! `PhysicalExtensionNode` variant, which historically carried no
+//! discriminator: the `PhysicalExtensionCodec` *was* the discriminator.
+//! `ComposedPhysicalExtensionCodec` copes by writing the *position* of the
+//! codec that encoded a payload into the bytes, so the decoding side must
+//! register the same codecs in the same order, and a plan two codecs both
+//! claim resolves to whichever was registered first.
+//!
+//! This module supplies the missing discriminator. [`ExecutionPlanFromProto`]
+//! is the decode contract every self-serializing plan implements — built-ins
+//! included, dispatched by their `PhysicalPlanType` variant. An extension plan
+//! is dispatched by its [`NAME`](ExecutionPlanFromProto::NAME) instead: its
+//! `try_to_proto` stamps the name on the `PhysicalExtensionNode` it writes,
and
+//! it is registered in an [`ExecutionPlanRegistry`] that the decoding session
+//! carries. Decoding then selects the decoder by name, independent of
+//! registration order.
+//!
+//! The registry is *session scoped*, matching the `FunctionRegistry`
+//! precedent: it is attached with `SessionConfig::with_extension`, so it
+//! travels with the session into every [`TaskContext`] without any new
+//! plumbing, stays testable, and stays multi-tenant safe.
+//!
+//! Registration is per plan type and entirely opt-in. A
`PhysicalExtensionNode`
+//! with no name — or with a name no decoder claims — falls back to the
existing
+//! `PhysicalExtensionCodec` chain, unchanged.
+//!
+//! [`TaskContext`]: datafusion_execution::TaskContext
+
+use std::any::{TypeId, type_name};
+use std::collections::HashMap;
+use std::collections::hash_map::Entry;
+use std::sync::Arc;
+
+use datafusion_common::{Result, config_err};
+use datafusion_proto_models::protobuf::PhysicalPlanNode;
+
+use crate::ExecutionPlan;
+use crate::proto::{ExecutionPlanDecodeCtx, ExecutionPlanFromProto};
+
+/// How the registry stores a decoder internally: a function pointer to the
+/// monomorphized [`ExecutionPlanFromProto::try_from_proto`].
+///
+/// Deliberately private. [`ExecutionPlanFromProto`] is the public contract,
and
+/// [`ExecutionPlanRegistry::decode`] is the public way to invoke one, so the
+/// storage can become something else (a `dyn` decoder object, to admit
+/// stateful or closure decoders — an FFI decoder carries a vtable and private
+/// data, which a bare `fn` never can) without a breaking change.
+type ExecutionPlanDecoder =
+ fn(&PhysicalPlanNode, &ExecutionPlanDecodeCtx<'_>) -> Result<Arc<dyn
ExecutionPlan>>;
+
+/// One registered decoder, plus the identity used to make re-registering the
+/// same type idempotent while a genuine name collision is an error.
+#[derive(Debug, Clone, Copy)]
+struct RegisteredPlan {
+ decoder: ExecutionPlanDecoder,
+ type_id: TypeId,
+ type_name: &'static str,
+}
+
+/// A name-keyed set of extension [`ExecutionPlan`] decoders.
+///
+/// Build one, register every extension plan the session must decode, and
+/// attach it with `SessionConfig::with_extension`. A session carries at most
+/// one registry: attaching another replaces it, so compose everything into
+/// one registry first. Register on the session that will *decode* the plan —
+/// for a distributed engine, that means every worker as well as the
+/// coordinator.
+#[derive(Debug, Clone, Default)]
+pub struct ExecutionPlanRegistry {
+ decoders: HashMap<String, RegisteredPlan>,
+}
+
+impl ExecutionPlanRegistry {
+ /// Create an empty registry.
+ pub fn new() -> Self {
+ Self::default()
+ }
+
+ /// Register `T` under its [`ExecutionPlanFromProto::NAME`].
+ ///
+ /// Registering the same type twice is a no-op. Registering a *different*
+ /// type under a name already taken is an error, so collisions surface here
+ /// rather than as a wrong decode later.
+ pub fn register<T: ExecutionPlanFromProto>(&mut self) -> Result<()> {
Review Comment:
Because every built-in (ex. `FilterExec`) implements
`ExecutionPlanFromProto`, `registry.register::<FilterExec>()` succeeds. We
probably want to block that or pre-register all the built ins.
##########
datafusion/proto/src/physical_plan/mod.rs:
##########
@@ -1393,20 +1399,58 @@ pub trait PhysicalPlanNodeExt: Sized {
ctx: &PhysicalPlanDecodeContext<'_>,
proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<Arc<dyn ExecutionPlan>> {
+ // Lookup-order policy, mirroring the one function decode already uses
a
+ // layer down (payload -> codec; else registry -> codec fallback): a
+ // plan that named itself on the wire and is registered on this session
+ // decodes itself, no codec involved. Anything else — an unnamed node
+ // from a codec-encoded writer, or a name this session does not know —
+ // takes the codec chain exactly as before.
+ let registry = ctx
+ .task_ctx()
+ .session_config()
+ .get_extension::<ExecutionPlanRegistry>();
+ if let Some(plan_name) = extension.plan_name.as_deref()
+ && let Some(registry) = registry.as_ref()
+ {
+ let plan_decoder = ConverterPlanDecoder {
+ ctx,
+ proto_converter,
+ };
+ // The decoder receives the whole node, like every built-in
+ // `try_from_proto`, and decodes its own children through the ctx.
+ // `None` here means no decoder claims the name; a decode *failure*
+ // is returned as-is rather than falling through to the codec.
+ if let Some(decoded) = registry.decode(
+ plan_name,
+ self.node(),
+ &ExecutionPlanDecodeCtx::new(&plan_decoder),
+ ) {
+ return decoded;
+ }
+ }
+
let inputs: Vec<Arc<dyn ExecutionPlan>> = extension
.inputs
.iter()
.map(|i| proto_converter.proto_to_execution_plan(i, ctx))
.collect::<Result<_>>()?;
- let extension_node = ctx.codec().try_decode(
- extension.node.as_slice(),
- &inputs,
- ctx.task_ctx(),
- proto_converter,
- )?;
-
- Ok(extension_node)
+ ctx.codec()
Review Comment:
We probably want a ticket + TODO to remove this fallback in df57 assuming
this lands in df56.
##########
datafusion/physical-plan/src/proto/registry.rs:
##########
@@ -0,0 +1,226 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Name-keyed decode registry for extension [`ExecutionPlan`]s.
+//!
+//! Built-in plans are dispatched by their `PhysicalPlanType` oneof variant, so
+//! the wire format names them. Extension plans all share the single
+//! `PhysicalExtensionNode` variant, which historically carried no
+//! discriminator: the `PhysicalExtensionCodec` *was* the discriminator.
+//! `ComposedPhysicalExtensionCodec` copes by writing the *position* of the
+//! codec that encoded a payload into the bytes, so the decoding side must
+//! register the same codecs in the same order, and a plan two codecs both
+//! claim resolves to whichever was registered first.
+//!
+//! This module supplies the missing discriminator. [`ExecutionPlanFromProto`]
+//! is the decode contract every self-serializing plan implements — built-ins
+//! included, dispatched by their `PhysicalPlanType` variant. An extension plan
+//! is dispatched by its [`NAME`](ExecutionPlanFromProto::NAME) instead: its
+//! `try_to_proto` stamps the name on the `PhysicalExtensionNode` it writes,
and
+//! it is registered in an [`ExecutionPlanRegistry`] that the decoding session
+//! carries. Decoding then selects the decoder by name, independent of
+//! registration order.
+//!
+//! The registry is *session scoped*, matching the `FunctionRegistry`
+//! precedent: it is attached with `SessionConfig::with_extension`, so it
+//! travels with the session into every [`TaskContext`] without any new
+//! plumbing, stays testable, and stays multi-tenant safe.
+//!
+//! Registration is per plan type and entirely opt-in. A
`PhysicalExtensionNode`
+//! with no name — or with a name no decoder claims — falls back to the
existing
+//! `PhysicalExtensionCodec` chain, unchanged.
+//!
+//! [`TaskContext`]: datafusion_execution::TaskContext
+
+use std::any::{TypeId, type_name};
+use std::collections::HashMap;
+use std::collections::hash_map::Entry;
+use std::sync::Arc;
+
+use datafusion_common::{Result, config_err};
+use datafusion_proto_models::protobuf::PhysicalPlanNode;
+
+use crate::ExecutionPlan;
+use crate::proto::{ExecutionPlanDecodeCtx, ExecutionPlanFromProto};
+
+/// How the registry stores a decoder internally: a function pointer to the
+/// monomorphized [`ExecutionPlanFromProto::try_from_proto`].
+///
+/// Deliberately private. [`ExecutionPlanFromProto`] is the public contract,
and
+/// [`ExecutionPlanRegistry::decode`] is the public way to invoke one, so the
+/// storage can become something else (a `dyn` decoder object, to admit
+/// stateful or closure decoders — an FFI decoder carries a vtable and private
+/// data, which a bare `fn` never can) without a breaking change.
+type ExecutionPlanDecoder =
+ fn(&PhysicalPlanNode, &ExecutionPlanDecodeCtx<'_>) -> Result<Arc<dyn
ExecutionPlan>>;
+
+/// One registered decoder, plus the identity used to make re-registering the
+/// same type idempotent while a genuine name collision is an error.
+#[derive(Debug, Clone, Copy)]
+struct RegisteredPlan {
+ decoder: ExecutionPlanDecoder,
+ type_id: TypeId,
+ type_name: &'static str,
+}
+
+/// A name-keyed set of extension [`ExecutionPlan`] decoders.
+///
+/// Build one, register every extension plan the session must decode, and
+/// attach it with `SessionConfig::with_extension`. A session carries at most
+/// one registry: attaching another replaces it, so compose everything into
+/// one registry first. Register on the session that will *decode* the plan —
+/// for a distributed engine, that means every worker as well as the
+/// coordinator.
+#[derive(Debug, Clone, Default)]
+pub struct ExecutionPlanRegistry {
+ decoders: HashMap<String, RegisteredPlan>,
+}
+
+impl ExecutionPlanRegistry {
+ /// Create an empty registry.
+ pub fn new() -> Self {
+ Self::default()
+ }
+
+ /// Register `T` under its [`ExecutionPlanFromProto::NAME`].
+ ///
+ /// Registering the same type twice is a no-op. Registering a *different*
+ /// type under a name already taken is an error, so collisions surface here
+ /// rather than as a wrong decode later.
+ pub fn register<T: ExecutionPlanFromProto>(&mut self) -> Result<()> {
+ self.insert(
+ T::NAME,
+ RegisteredPlan {
+ decoder: T::try_from_proto,
+ type_id: TypeId::of::<T>(),
+ type_name: type_name::<T>(),
+ },
+ )
+ }
+
+ fn insert(&mut self, name: impl Into<String>, plan: RegisteredPlan) ->
Result<()> {
+ let name = name.into();
+ if name.is_empty() {
+ return config_err!(
+ "Cannot register the extension ExecutionPlan decoder for {}
under an empty name",
+ plan.type_name
+ );
+ }
+ match self.decoders.entry(name) {
+ Entry::Vacant(entry) => {
+ entry.insert(plan);
+ Ok(())
+ }
+ // Re-registering the same type is a no-op: sessions are often
+ // configured by more than one layer of an application.
+ Entry::Occupied(entry) if entry.get().type_id == plan.type_id =>
Ok(()),
+ Entry::Occupied(entry) => config_err!(
+ "Extension ExecutionPlan name '{}' is already registered by
{}, cannot register {}. \
+ Namespace the name with the owning crate to avoid the
collision.",
+ entry.key(),
+ entry.get().type_name,
+ plan.type_name
+ ),
+ }
+ }
+
+ /// Decode `node` with the decoder registered under `name`.
+ ///
+ /// `None` means "no decoder claims this name" — the caller falls back to
+ /// the `PhysicalExtensionCodec` chain. `Some(Err(..))` means the decoder
+ /// that *does* own the name failed, which is fatal: falling back there
+ /// would let another codec decode the payload wrongly, the very thing the
name
+ /// exists to prevent.
+ pub fn decode(
+ &self,
+ name: &str,
+ node: &PhysicalPlanNode,
+ ctx: &ExecutionPlanDecodeCtx<'_>,
+ ) -> Option<Result<Arc<dyn ExecutionPlan>>> {
+ let plan = self.decoders.get(name)?;
+ Some((plan.decoder)(node, ctx))
+ }
Review Comment:
I don't think we need to pass `name` because the name is already in the
`node` right? That avoids someone passing in a `name` which doesn't match the
`node`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
