adriangb commented on code in PR #24631: URL: https://github.com/apache/datafusion/pull/24631#discussion_r4011232340
########## datafusion/physical-plan/src/proto/registry.rs: ########## @@ -0,0 +1,418 @@ +// 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. An extension plan declares a +//! globally unique name via [`ExtensionExecutionPlan::PLAN_NAME`], stamps it on +//! the wire with +//! [`ExecutionPlanEncodeCtx::encode_extension`](super::ExecutionPlanEncodeCtx::encode_extension), +//! and registers its decoder on the session with +//! [`ExecutionPlanRegistryExt::register_execution_plan`]. Decoding then selects +//! the decoder by name, independent of registration order. +//! +//! The registry is *session scoped*, matching the `FunctionRegistry` +//! precedent: it lives in the [`SessionConfig`] extension map, 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_execution::config::SessionConfig; +use datafusion_proto_models::protobuf::PhysicalPlanNode; + +use crate::ExecutionPlan; +use crate::proto::ExecutionPlanDecodeCtx; + +/// How the registry stores a decoder internally: a function pointer to the +/// monomorphized [`ExtensionExecutionPlan::try_from_proto`]. +/// +/// Deliberately private. [`ExtensionExecutionPlan`] 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>>; + +/// An extension [`ExecutionPlan`] that serializes itself, without a +/// `PhysicalExtensionCodec`. +/// +/// Implement this alongside +/// [`ExecutionPlan::try_to_proto`], then +/// register the type on 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, ExecutionPlanRegistryExt, +/// # ExtensionExecutionPlan, +/// # }; +/// # 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>> { +/// # Ok(Some(ctx.encode_extension(self, vec![])?)) +/// # } +/// # } +/// impl ExtensionExecutionPlan for MyExec { +/// const PLAN_NAME: &'static str = "my-crate.MyExec"; +/// +/// fn try_from_proto( +/// node: &PhysicalPlanNode, +/// ctx: &ExecutionPlanDecodeCtx<'_>, +/// ) -> Result<Arc<dyn ExecutionPlan>> { +/// let parts = ctx.decode_extension::<Self>(node)?; +/// // ... rebuild `MyExec` from the payload and children ... +/// # unimplemented!() +/// } +/// } +/// +/// let mut config = SessionConfig::new(); +/// config.register_execution_plan::<MyExec>()?; +/// # Ok::<(), datafusion_common::DataFusionError>(()) +/// ``` +pub trait ExtensionExecutionPlan: ExecutionPlan + Sized { + /// A globally unique name for this plan type, used as the wire + /// discriminator. + /// + /// Namespace it with the owning crate (`"my-crate.MyExec"`) so that a + /// collision between two independent crates surfaces as a registration + /// error rather than as a wrong decode. + const PLAN_NAME: &'static str; + + /// Reconstruct the plan from the `PhysicalPlanNode` written by + /// [`ExecutionPlan::try_to_proto`]. + /// + /// Use + /// [`ExecutionPlanDecodeCtx::decode_extension`](super::ExecutionPlanDecodeCtx::decode_extension) + /// to unwrap the payload and decode the children. + fn try_from_proto( + node: &PhysicalPlanNode, + ctx: &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. +/// +/// Usually not constructed directly: use +/// [`ExecutionPlanRegistryExt::register_execution_plan`] on a [`SessionConfig`], +/// which creates, updates and stores the registry for you. +#[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 [`ExtensionExecutionPlan::PLAN_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: ExtensionExecutionPlan>(&mut self) -> Result<()> { + self.insert( + T::PLAN_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)) + } + + /// Whether a decoder is registered under `name`. + pub fn contains(&self, name: &str) -> bool { + self.decoders.contains_key(name) + } + + /// Every registered name, in arbitrary order. + pub fn names(&self) -> impl Iterator<Item = &str> { Review Comment: Lets make sure any public method is actually needed or justified. We should not guess at public APIs, prefer to wait until they are needed. ########## datafusion/physical-plan/src/proto/registry.rs: ########## @@ -0,0 +1,418 @@ +// 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. An extension plan declares a +//! globally unique name via [`ExtensionExecutionPlan::PLAN_NAME`], stamps it on +//! the wire with +//! [`ExecutionPlanEncodeCtx::encode_extension`](super::ExecutionPlanEncodeCtx::encode_extension), +//! and registers its decoder on the session with +//! [`ExecutionPlanRegistryExt::register_execution_plan`]. Decoding then selects +//! the decoder by name, independent of registration order. +//! +//! The registry is *session scoped*, matching the `FunctionRegistry` +//! precedent: it lives in the [`SessionConfig`] extension map, 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_execution::config::SessionConfig; +use datafusion_proto_models::protobuf::PhysicalPlanNode; + +use crate::ExecutionPlan; +use crate::proto::ExecutionPlanDecodeCtx; + +/// How the registry stores a decoder internally: a function pointer to the +/// monomorphized [`ExtensionExecutionPlan::try_from_proto`]. +/// +/// Deliberately private. [`ExtensionExecutionPlan`] 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>>; + +/// An extension [`ExecutionPlan`] that serializes itself, without a +/// `PhysicalExtensionCodec`. +/// +/// Implement this alongside +/// [`ExecutionPlan::try_to_proto`], then +/// register the type on 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, ExecutionPlanRegistryExt, +/// # ExtensionExecutionPlan, +/// # }; +/// # 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>> { +/// # Ok(Some(ctx.encode_extension(self, vec![])?)) +/// # } +/// # } +/// impl ExtensionExecutionPlan for MyExec { +/// const PLAN_NAME: &'static str = "my-crate.MyExec"; +/// +/// fn try_from_proto( +/// node: &PhysicalPlanNode, +/// ctx: &ExecutionPlanDecodeCtx<'_>, +/// ) -> Result<Arc<dyn ExecutionPlan>> { +/// let parts = ctx.decode_extension::<Self>(node)?; +/// // ... rebuild `MyExec` from the payload and children ... +/// # unimplemented!() +/// } +/// } +/// +/// let mut config = SessionConfig::new(); +/// config.register_execution_plan::<MyExec>()?; +/// # Ok::<(), datafusion_common::DataFusionError>(()) +/// ``` +pub trait ExtensionExecutionPlan: ExecutionPlan + Sized { + /// A globally unique name for this plan type, used as the wire + /// discriminator. + /// + /// Namespace it with the owning crate (`"my-crate.MyExec"`) so that a + /// collision between two independent crates surfaces as a registration + /// error rather than as a wrong decode. + const PLAN_NAME: &'static str; + + /// Reconstruct the plan from the `PhysicalPlanNode` written by + /// [`ExecutionPlan::try_to_proto`]. + /// + /// Use + /// [`ExecutionPlanDecodeCtx::decode_extension`](super::ExecutionPlanDecodeCtx::decode_extension) + /// to unwrap the payload and decode the children. + fn try_from_proto( + node: &PhysicalPlanNode, + ctx: &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. +/// +/// Usually not constructed directly: use +/// [`ExecutionPlanRegistryExt::register_execution_plan`] on a [`SessionConfig`], +/// which creates, updates and stores the registry for you. +#[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 [`ExtensionExecutionPlan::PLAN_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: ExtensionExecutionPlan>(&mut self) -> Result<()> { + self.insert( + T::PLAN_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)) + } + + /// Whether a decoder is registered under `name`. + pub fn contains(&self, name: &str) -> bool { + self.decoders.contains_key(name) + } + + /// Every registered name, in arbitrary order. + pub fn names(&self) -> impl Iterator<Item = &str> { + self.decoders.keys().map(String::as_str) + } +} + +/// Session-scoped registration of extension [`ExecutionPlan`] decoders. +/// +/// Implemented for [`SessionConfig`], which carries the registry into every +/// `TaskContext` derived from the session. Register on the session that will +/// *decode* the plan — for a distributed engine, that means every worker as +/// well as the coordinator. +/// +/// Registration is additive: each call merges into the registry the config +/// already carries, so a library registering its own plans cannot silently +/// drop another's. (Attaching a hand-built [`ExecutionPlanRegistry`] with +/// `SessionConfig::with_extension` *replaces* whatever was there, which is why +/// this trait offers no setter.) See [`ExtensionExecutionPlan`] for a full +/// example. +pub trait ExecutionPlanRegistryExt { Review Comment: Is this necessary? Could we for now have people do something like: ```rust let mut registry = ...; // register plans let mut builder = SessionConfigBUilder::default(); builder = builder.with_execution_plan_registry(registry); ... ``` -- 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]
