adriangb commented on code in PR #24631: URL: https://github.com/apache/datafusion/pull/24631#discussion_r4039069541
########## datafusion/physical-plan/src/proto/registry.rs: ########## @@ -0,0 +1,374 @@ +// 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 decoding 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. An extension plan +//! implements [`ExtensionPlanFromProto`], which pairs the name it is +//! dispatched by with the constructor that rebuilds it: its `try_to_proto` +//! stamps [`NAME`](ExtensionPlanFromProto::NAME) on the +//! `PhysicalExtensionNode` it writes, and [`register_execution_plan`] puts its +//! decoder in the [`ProtoDecoderRegistry`] that the decoding session carries. +//! Decoding then selects the decoder by name, independent of registration +//! order. +//! +//! Only extension plans implement the trait. A built-in has a +//! `PhysicalPlanType` variant of its own, is decoded by an inherent +//! `try_from_proto` dispatched from that variant, and has no wire name to be +//! registered under — so [`register_execution_plan`] will not accept one. +//! +//! This module is only the `ExecutionPlan` face of the registry. The store +//! itself is shared with every other extension kind, keyed by the trait as +//! well as the name, so one session carries one registry. +//! +//! 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::sync::Arc; + +use datafusion_common::Result; +use datafusion_proto_models::ProtoDecoderRegistry; +use datafusion_proto_models::protobuf::PhysicalPlanNode; +use datafusion_proto_models::protobuf::physical_plan_node::PhysicalPlanType; + +use crate::ExecutionPlan; +use crate::proto::ExecutionPlanDecodeCtx; + +/// The wire name of an extension [`ExecutionPlan`], and the key it is +/// registered and dispatched by. +/// +/// Only extension plans implement this. A built-in plan has a +/// `PhysicalPlanType` variant of its own and is dispatched by that variant, so +/// it needs no name and cannot be registered — [`register_execution_plan`] +/// takes this trait as its bound, so `register_execution_plan::<FilterExec>()` +/// does not compile. +/// +/// One impl block says what the plan is called and how it decodes; its +/// `try_to_proto` on [`ExecutionPlan`] writes the matching node: +/// +/// ``` Review Comment: This is way too much fluff. ########## 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: Agreed — better than a `name(&self)` helper: the author never writes the name at all. ```rust fn try_to_proto(&self, ctx: &ExecutionPlanEncodeCtx<'_>) -> Result<Option<PhysicalPlanNode>> { Ok(Some(ctx.extension_node::<Self>(my_payload_bytes(self)?, self.children())?)) } ``` `extension_node` builds the whole node and stamps `Self::NAME` from the same const the registry is keyed by. That also answers your point below: `decode` reads the name off the node instead of taking it. ########## 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 think the systems might need to co-exist at least a bit longer. At beast we could deprecate in df56. But we probably want to sort out FFI first. ########## 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: Agreed — done. It is a free function over the shared registry now, and it reads the name off the node: ```rust pub fn decode_execution_plan( registry: &ProtoDecoderRegistry, node: &PhysicalPlanNode, ctx: &ExecutionPlanDecodeCtx<'_>, ) -> Option<Result<Arc<dyn ExecutionPlan>>> ``` `None` now also covers "not an extension node" and "carries no name", which simplified the caller in `datafusion-proto` as well — it no longer has to check for the name before calling. ########## 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: Done: `decode_children` now sits directly after `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() + .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: Good catch. The old code built a fresh `Configuration` error and put the codec's error inside a string, so a codec that failed for an unrelated reason (a short read, an I/O fault) came back as a config error and the caller lost the kind. It now adds the hint as context on the codec's own error, with `DataFusionError::context`, so the kind survives: ```rust codec_error.context(format!( "No decoder is registered for the extension ExecutionPlan '{plan_name}'. ..." )) ``` New test `the_codecs_own_error_survives_the_missing_registration_hint` uses a codec that fails with `ResourcesExhausted` and asserts that `err.find_root()` is still `ResourcesExhausted`, that the codec's own message survives, and that the hint is attached. ########## 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: Two things here. **A default registry in `SessionConfig`** cannot be done: `SessionConfig` is in `datafusion-execution`, below `datafusion-physical-plan`, so it cannot name the type, and `get_extension::<T>` needs the type. That layering is also why the registry rides in the extension map at all. **Multiple registries:** there is now exactly one, `ProtoDecoderRegistry`, shared by every extension kind and keyed by `(trait, name)` — so plans and expressions register into the same object and cannot collide. The application that owns the session builds it; a library exposes `pub fn register(&mut ProtoDecoderRegistry)` and never calls `with_extension` itself. Documented on the type and in the upgrade guide. ########## 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 is right, and I confirmed it in the generated code. `datafusion/proto-models/src/generated/pbjson.rs` ends the field match with: ```rust _ => Err(serde::de::Error::unknown_field(value, FIELDS)), ``` 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. The upgrade guide and the commit message now say this, with the advice to upgrade readers before opting a plan in. ########## 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: Good catch — fixed at the root: built-ins no longer implement the decode trait at all. `ExecutionPlanFromProto` is gone. Extension plans implement a single `ExtensionPlanFromProto` (wire name + constructor), and `register_execution_plan` takes that as its bound: ``` error[E0277]: the trait bound `FilterExec: ExtensionPlanFromProto` is not satisfied ``` Built-ins keep the inherent `try_from_proto` they have had since 55.0.0. That removed the 29-file commit and its import break entirely. A `compile_fail` doctest pins this; I verified it fails on the bound and not on something incidental. -- 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]
