adriangb commented on code in PR #24631: URL: https://github.com/apache/datafusion/pull/24631#discussion_r4047665334
########## datafusion/physical-plan/src/proto/registry.rs: ########## @@ -0,0 +1,275 @@ +// 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. +//! +//! A built-in plan has a `PhysicalPlanType` variant of its own, so the wire +//! format names it. Extension plans all share `PhysicalExtensionNode`, which +//! carried no discriminator, so the `PhysicalExtensionCodec` *was* the +//! discriminator — and `ComposedPhysicalExtensionCodec` resolves by +//! registration order, which two independent crates cannot agree on. +//! +//! An extension plan implements [`ExtensionPlanFromProto`] instead, and +//! [`register_execution_plan`] puts its decoder in the +//! [`ProtoDecoderRegistry`] the decoding session carries. Lookup is by name, +//! so registration order does not matter. Built-ins cannot be registered; +//! anything unnamed or unregistered still takes the codec chain, unchanged. +//! +//! This is only the `ExecutionPlan` face of the registry: the store is shared +//! with every other extension kind, keyed by the trait as well as the name. +//! `datafusion-examples/examples/proto/extension_plan_registry.rs` is a +//! worked example. + +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 constructor that +/// rebuilds it. +/// +/// Only extension plans implement this, so a built-in cannot be registered by +/// mistake. One impl block carries both halves, and +/// [`ExecutionPlan::try_to_proto`] writes the matching node: +/// +/// ```ignore +/// impl ExecutionPlan for MyExec { +/// fn try_to_proto( +/// &self, +/// ctx: &ExecutionPlanEncodeCtx<'_>, +/// ) -> Result<Option<PhysicalPlanNode>> { +/// // Stamps `Self::NAME`, so the encoded name and the registry key +/// // cannot drift apart. +/// Ok(Some(ctx.extension_node::<Self>(self.payload()?, self.children())?)) +/// } +/// } +/// +/// impl ExtensionPlanFromProto for MyExec { +/// 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)?; +/// my_plan_from_bytes(&extension.node, children) +/// } +/// } +/// ``` +pub trait ExtensionPlanFromProto: ExecutionPlan + Sized { Review Comment: Ahhh, this is a great question, I am glad you asked! I spent quite some time chatting with my agent about this yesterday because it worried me too. Here's my summary. ## Status quo Lets look at the serialization side first, because it already has traits involved. We call `ExecutionPlan::try_to_proto` (in this case `DataSourceExec::try_to_proto`) -> calls `DataSource::try_to_proto` -> returns a `PhysicalPlanNode`. This is partially for historical reasons: we are still using the proto from when `ParquetExecutionPlan` was a thing (`ParquetScanExecNode`). The weird thing here is that the `DataSource` return an `ExecutionPlan` level proto object (`PhysicalPlanNode`). It's a violation of the layers of abstraction. The deserialization side is the same: a central match finds a `PhysicalPlanNode` w/ `PhysicalPlanType::ParquetScan`. So for a combination of historical reasons (not breaking the wire messages) and simplicity of implementation we skipped past this question thus far. ## How to fix it To fully preserve the layers of abstraction, a `DataSourceExec` must be able to serialize itself w/o downcast matching or otherwise knowing which `dyn DataSource` implementation it wraps. That implementation must be able to be a 3rd party crate. There may even be 2 crates involved: crate A provides an `ExecutionPlan` (imagine `DataSourceExec` was a 3rd party crate) and crate B provides an inner trait implementation (imagine if crate A provided `DataSource` the trait and crate B implemented it). The only way to make this work is for each trait layer to have it's own registry and wire message. What DataFusion _needs_ to provide is a place to register these registries (🤯). We *can* also provide some generic container messages to allow preserving an introspectable tree, even if it can't be deserialized. The general rule is: 1. A type serializes only itself. 2. It asks the context to serialize anything behind a `dyn`. If you look at https://github.com/apache/datafusion/pull/24628 this is already implemented to some extent: the registry DataFusion provides is keyed on `(TypeId, name)` i.e. `(ExecutionPlan, "Parquet")` will not collide with `(DataSource, "Parquet")` (or `(PhysicalExpr, "Data")`). If we re-implemented `DataSourceExec` / `DataSource` to match this (which I haven't for the proto wire reasons mentioned above) it would look like this: ```rust impl ExecutionPlan { ... fn try_to_proto(&self, ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { let data_source = ctx.encode::<dyn DataSource>(self.data_source())?; // arbitrary bytes let node = protobuf::DataSourceExecNode { data_source, ... }; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some(PhysicalPlanType::ExecutionPlan(node), })) } } ``` The decode side would look similar in reverse, calling `let data_source = ctx.decode::<dyn DataSource>(&node.data_source)?; // returns Arc<dyn DataSource>`. What if there are more layers (e.g. `FileSource`)? This part is the same, but what the wire container is gets complicated. The `DataSourceExecNode::data_source` field could be arbitrary bytes, and we leave it up to each trait to define it's own fire protocol. It also could be a generic typed proto container: ```proto message ProtoAny { string type_name = 1; // "datafusion.ParquetSource" or "my-crate.MySource" bytes payload = 2; repeated ProtoAny children = 3; } ``` Even if it is arbitrary bytes, it could still use the proto container. This is all optional, can be decided on a case by case basis. But it will need to be type erased on the wire. -- 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]
