adriangb commented on code in PR #26555: URL: https://github.com/apache/datafusion/pull/26555#discussion_r4393404989
########## datafusion/datasource/src/proto/registry.rs: ########## @@ -0,0 +1,416 @@ +// 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 [`DataSource`]s. +//! +//! Custom sources serialize through the existing extension-plan wire format: +//! [`DataSourceExec`] is the registered plan type, and its payload identifies +//! the source decoder. Built-in sources keep their existing wire variants. + +use std::sync::Arc; + +use datafusion_common::{Result, plan_datafusion_err, plan_err}; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::expect_plan_variant; +use datafusion_physical_plan::proto::{ + ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx, ExtensionPlanFromProto, +}; +use datafusion_proto_models::ProtoDecoderRegistry; +use datafusion_proto_models::protobuf::physical_plan_node::PhysicalPlanType; +use datafusion_proto_models::protobuf::{ + PhysicalExtensionDataSourceNode, PhysicalPlanNode, +}; +use prost::Message; + +use crate::source::{DataSource, DataSourceExec}; + +/// The wire name of an extension [`DataSource`] and the constructor that +/// rebuilds it. Only custom sources implement this trait. +/// +/// Implement [`DataSource::try_to_proto`] using [`extension_data_source_node`]. +/// The helper writes a named `DataSourceExec` plan containing the source's +/// [`Self::NAME`] and payload. The source constructor receives that typed +/// payload and can read the decoding session through +/// [`ExecutionPlanDecodeCtx::task_ctx`]. Runtime dependencies should be rebuilt +/// from that session rather than serialized. +pub trait ExtensionDataSourceFromProto: DataSource + Sized { + /// Namespace this name with the owning crate, e.g. `"my_crate.MySource"`. + /// Names are scoped to `dyn DataSource`, independently of plan names. + const NAME: &'static str; + + /// Rebuild the source from its typed extension payload. + /// [`DataSourceExec`] supplies the execution-plan wrapper. + fn try_from_proto( + node: &PhysicalExtensionDataSourceNode, + ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn DataSource>>; Review Comment: The input is `&PhysicalExtensionDataSourceNode`, but implementors only need the bytes. `source_name` is internal plumbing. Pass the payload bytes instead, so the message can change (for example, to add a version field) without breaking implementors: ```suggestion fn try_from_proto( payload: &[u8], ctx: &ExecutionPlanDecodeCtx<'_>, ) -> Result<Arc<dyn DataSource>>; ``` Also add a provided `register_proto` here, as `ExtensionPlanFromProto` (`physical-plan/src/proto/registry.rs:136`) and `ExtensionPhysicalExprFromProto` (`physical-expr/src/proto.rs:141`) do. Then `register_data_source` does not need to be public. ########## datafusion/datasource/src/proto/registry.rs: ########## @@ -0,0 +1,416 @@ +// 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 [`DataSource`]s. +//! +//! Custom sources serialize through the existing extension-plan wire format: +//! [`DataSourceExec`] is the registered plan type, and its payload identifies +//! the source decoder. Built-in sources keep their existing wire variants. + +use std::sync::Arc; + +use datafusion_common::{Result, plan_datafusion_err, plan_err}; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::expect_plan_variant; +use datafusion_physical_plan::proto::{ + ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx, ExtensionPlanFromProto, +}; +use datafusion_proto_models::ProtoDecoderRegistry; +use datafusion_proto_models::protobuf::physical_plan_node::PhysicalPlanType; +use datafusion_proto_models::protobuf::{ + PhysicalExtensionDataSourceNode, PhysicalPlanNode, +}; +use prost::Message; + +use crate::source::{DataSource, DataSourceExec}; + +/// The wire name of an extension [`DataSource`] and the constructor that +/// rebuilds it. Only custom sources implement this trait. +/// +/// Implement [`DataSource::try_to_proto`] using [`extension_data_source_node`]. +/// The helper writes a named `DataSourceExec` plan containing the source's +/// [`Self::NAME`] and payload. The source constructor receives that typed +/// payload and can read the decoding session through +/// [`ExecutionPlanDecodeCtx::task_ctx`]. Runtime dependencies should be rebuilt +/// from that session rather than serialized. +pub trait ExtensionDataSourceFromProto: DataSource + Sized { + /// Namespace this name with the owning crate, e.g. `"my_crate.MySource"`. + /// Names are scoped to `dyn DataSource`, independently of plan names. + const NAME: &'static str; + + /// Rebuild the source from its typed extension payload. + /// [`DataSourceExec`] supplies the execution-plan wrapper. + fn try_from_proto( + node: &PhysicalExtensionDataSourceNode, + ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn DataSource>>; +} + +/// Build an extension plan for source `T`, stamping the wrapper's plan name +/// and the source's registered name. +/// +/// This uses [`ExecutionPlanEncodeCtx::extension_node`] for [`DataSourceExec`], +/// whose payload contains `T::NAME` and the source bytes. The helper lives in +/// this crate because the plan encode context cannot depend on [`DataSource`]. +/// +/// A source implementing [`ExtensionDataSourceFromProto`] can use this in its +/// [`DataSource::try_to_proto`] hook: +/// +/// ```ignore +/// fn try_to_proto( +/// &self, +/// ctx: &ExecutionPlanEncodeCtx<'_>, +/// ) -> Result<Option<PhysicalPlanNode>> { +/// let payload = self.partition_id.to_le_bytes().to_vec(); +/// Ok(Some(extension_data_source_node::<Self>(ctx, payload)?)) +/// } +/// ``` +pub fn extension_data_source_node<T: ExtensionDataSourceFromProto>( + ctx: &ExecutionPlanEncodeCtx<'_>, + payload: Vec<u8>, +) -> Result<PhysicalPlanNode> { + let source = PhysicalExtensionDataSourceNode { + node: payload, + source_name: Some(T::NAME.to_string()), + }; + ctx.extension_node::<DataSourceExec>(source.encode_to_vec(), vec![]) Review Comment: This helper is public only because the source must build the `DataSourceExec` wrapper itself (`DataSource::try_to_proto` returns `PhysicalPlanNode`). With the precursor (#26824) (`DataSource::try_to_proto_source -> DataSourceNode`), `DataSourceExec` builds the wrapper. Then this function can be removed or made private. ########## datafusion/datasource/src/proto/registry.rs: ########## @@ -0,0 +1,416 @@ +// 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 [`DataSource`]s. +//! +//! Custom sources serialize through the existing extension-plan wire format: +//! [`DataSourceExec`] is the registered plan type, and its payload identifies +//! the source decoder. Built-in sources keep their existing wire variants. + +use std::sync::Arc; + +use datafusion_common::{Result, plan_datafusion_err, plan_err}; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::expect_plan_variant; +use datafusion_physical_plan::proto::{ + ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx, ExtensionPlanFromProto, +}; +use datafusion_proto_models::ProtoDecoderRegistry; +use datafusion_proto_models::protobuf::physical_plan_node::PhysicalPlanType; +use datafusion_proto_models::protobuf::{ + PhysicalExtensionDataSourceNode, PhysicalPlanNode, +}; +use prost::Message; + +use crate::source::{DataSource, DataSourceExec}; + +/// The wire name of an extension [`DataSource`] and the constructor that +/// rebuilds it. Only custom sources implement this trait. +/// +/// Implement [`DataSource::try_to_proto`] using [`extension_data_source_node`]. +/// The helper writes a named `DataSourceExec` plan containing the source's +/// [`Self::NAME`] and payload. The source constructor receives that typed +/// payload and can read the decoding session through +/// [`ExecutionPlanDecodeCtx::task_ctx`]. Runtime dependencies should be rebuilt +/// from that session rather than serialized. +pub trait ExtensionDataSourceFromProto: DataSource + Sized { + /// Namespace this name with the owning crate, e.g. `"my_crate.MySource"`. + /// Names are scoped to `dyn DataSource`, independently of plan names. + const NAME: &'static str; + + /// Rebuild the source from its typed extension payload. + /// [`DataSourceExec`] supplies the execution-plan wrapper. + fn try_from_proto( + node: &PhysicalExtensionDataSourceNode, + ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn DataSource>>; +} + +/// Build an extension plan for source `T`, stamping the wrapper's plan name +/// and the source's registered name. +/// +/// This uses [`ExecutionPlanEncodeCtx::extension_node`] for [`DataSourceExec`], +/// whose payload contains `T::NAME` and the source bytes. The helper lives in +/// this crate because the plan encode context cannot depend on [`DataSource`]. +/// +/// A source implementing [`ExtensionDataSourceFromProto`] can use this in its +/// [`DataSource::try_to_proto`] hook: +/// +/// ```ignore +/// fn try_to_proto( +/// &self, +/// ctx: &ExecutionPlanEncodeCtx<'_>, +/// ) -> Result<Option<PhysicalPlanNode>> { +/// let payload = self.partition_id.to_le_bytes().to_vec(); +/// Ok(Some(extension_data_source_node::<Self>(ctx, payload)?)) +/// } +/// ``` +pub fn extension_data_source_node<T: ExtensionDataSourceFromProto>( + ctx: &ExecutionPlanEncodeCtx<'_>, + payload: Vec<u8>, +) -> Result<PhysicalPlanNode> { + let source = PhysicalExtensionDataSourceNode { + node: payload, + source_name: Some(T::NAME.to_string()), + }; + ctx.extension_node::<DataSourceExec>(source.encode_to_vec(), vec![]) +} + +// Private so the decoder representation can evolve without changing the API. +type DataSourceDecoder = fn( + &PhysicalExtensionDataSourceNode, + &ExecutionPlanDecodeCtx<'_>, +) -> Result<Arc<dyn DataSource>>; + +/// Register `T` under its [`ExtensionDataSourceFromProto::NAME`], together +/// with the shared [`DataSourceExec`] extension-plan decoder. +/// +/// Registering the same type twice is a no-op; a different type under a taken +/// name, or an empty name, is an error. Source names are scoped to `dyn DataSource`. +/// Add these decoders to the same [`ProtoDecoderRegistry`] as the application's +/// plans and expressions with +/// [`SessionConfig::update_extension`](datafusion_execution::config::SessionConfig::update_extension), +/// so decoders added by other libraries are preserved. +/// +/// Built-in sources do not implement the extension trait: +/// +/// ```compile_fail +/// use datafusion_datasource::memory::MemorySourceConfig; +/// use datafusion_datasource::proto::{ProtoDecoderRegistry, register_data_source}; +/// let mut registry = ProtoDecoderRegistry::new(); +/// register_data_source::<MemorySourceConfig>(&mut registry).unwrap(); +/// ``` +/// +/// The same imports compile without attempting to register a built-in: +/// +/// ``` +/// use datafusion_datasource::memory::MemorySourceConfig; +/// #[allow(unused_imports)] +/// use datafusion_datasource::proto::{ProtoDecoderRegistry, register_data_source}; +/// let _ = (MemorySourceConfig::try_from_proto, ProtoDecoderRegistry::new()); +/// ``` +pub fn register_data_source<T: ExtensionDataSourceFromProto>( + registry: &mut ProtoDecoderRegistry, +) -> Result<()> { + // The shared wrapper may remain registered if source registration fails. + // This is safe: it still requires a registered source decoder and rejects + // unknown sources rather than decoding their payloads through a plan codec. + DataSourceExec::register_proto(registry)?; + registry.register_decoder::<dyn DataSource, T, DataSourceDecoder>( + T::NAME, + T::try_from_proto, + ) +} Review Comment: Registration is not atomic. If `T::NAME` is empty or taken, `datafusion.DataSourceExec` stays registered (`register_rejects_an_empty_source_name` tests this). The simplest fix is to validate first: ```suggestion if T::NAME.is_empty() { return plan_err!("extension DataSource name must not be empty"); } registry.register_decoder::<dyn DataSource, T, DataSourceDecoder>( T::NAME, T::try_from_proto, )?; DataSourceExec::register_proto(registry) } ``` A collision can still leave the wrapper unregistered, but that is harmless. The precursor removes this problem: no wrapper registration is needed. ########## datafusion/datasource/src/proto/registry.rs: ########## @@ -0,0 +1,416 @@ +// 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 [`DataSource`]s. +//! +//! Custom sources serialize through the existing extension-plan wire format: +//! [`DataSourceExec`] is the registered plan type, and its payload identifies +//! the source decoder. Built-in sources keep their existing wire variants. + +use std::sync::Arc; + +use datafusion_common::{Result, plan_datafusion_err, plan_err}; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::expect_plan_variant; +use datafusion_physical_plan::proto::{ + ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx, ExtensionPlanFromProto, +}; +use datafusion_proto_models::ProtoDecoderRegistry; +use datafusion_proto_models::protobuf::physical_plan_node::PhysicalPlanType; +use datafusion_proto_models::protobuf::{ + PhysicalExtensionDataSourceNode, PhysicalPlanNode, +}; +use prost::Message; + +use crate::source::{DataSource, DataSourceExec}; + +/// The wire name of an extension [`DataSource`] and the constructor that +/// rebuilds it. Only custom sources implement this trait. +/// +/// Implement [`DataSource::try_to_proto`] using [`extension_data_source_node`]. +/// The helper writes a named `DataSourceExec` plan containing the source's +/// [`Self::NAME`] and payload. The source constructor receives that typed +/// payload and can read the decoding session through +/// [`ExecutionPlanDecodeCtx::task_ctx`]. Runtime dependencies should be rebuilt +/// from that session rather than serialized. +pub trait ExtensionDataSourceFromProto: DataSource + Sized { + /// Namespace this name with the owning crate, e.g. `"my_crate.MySource"`. + /// Names are scoped to `dyn DataSource`, independently of plan names. + const NAME: &'static str; + + /// Rebuild the source from its typed extension payload. + /// [`DataSourceExec`] supplies the execution-plan wrapper. + fn try_from_proto( + node: &PhysicalExtensionDataSourceNode, + ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn DataSource>>; +} + +/// Build an extension plan for source `T`, stamping the wrapper's plan name +/// and the source's registered name. +/// +/// This uses [`ExecutionPlanEncodeCtx::extension_node`] for [`DataSourceExec`], +/// whose payload contains `T::NAME` and the source bytes. The helper lives in +/// this crate because the plan encode context cannot depend on [`DataSource`]. +/// +/// A source implementing [`ExtensionDataSourceFromProto`] can use this in its +/// [`DataSource::try_to_proto`] hook: +/// +/// ```ignore +/// fn try_to_proto( +/// &self, +/// ctx: &ExecutionPlanEncodeCtx<'_>, +/// ) -> Result<Option<PhysicalPlanNode>> { +/// let payload = self.partition_id.to_le_bytes().to_vec(); +/// Ok(Some(extension_data_source_node::<Self>(ctx, payload)?)) +/// } +/// ``` +pub fn extension_data_source_node<T: ExtensionDataSourceFromProto>( + ctx: &ExecutionPlanEncodeCtx<'_>, + payload: Vec<u8>, +) -> Result<PhysicalPlanNode> { + let source = PhysicalExtensionDataSourceNode { + node: payload, + source_name: Some(T::NAME.to_string()), + }; + ctx.extension_node::<DataSourceExec>(source.encode_to_vec(), vec![]) +} + +// Private so the decoder representation can evolve without changing the API. +type DataSourceDecoder = fn( + &PhysicalExtensionDataSourceNode, + &ExecutionPlanDecodeCtx<'_>, +) -> Result<Arc<dyn DataSource>>; + +/// Register `T` under its [`ExtensionDataSourceFromProto::NAME`], together +/// with the shared [`DataSourceExec`] extension-plan decoder. +/// +/// Registering the same type twice is a no-op; a different type under a taken +/// name, or an empty name, is an error. Source names are scoped to `dyn DataSource`. +/// Add these decoders to the same [`ProtoDecoderRegistry`] as the application's +/// plans and expressions with +/// [`SessionConfig::update_extension`](datafusion_execution::config::SessionConfig::update_extension), +/// so decoders added by other libraries are preserved. +/// +/// Built-in sources do not implement the extension trait: +/// +/// ```compile_fail +/// use datafusion_datasource::memory::MemorySourceConfig; +/// use datafusion_datasource::proto::{ProtoDecoderRegistry, register_data_source}; +/// let mut registry = ProtoDecoderRegistry::new(); +/// register_data_source::<MemorySourceConfig>(&mut registry).unwrap(); +/// ``` +/// +/// The same imports compile without attempting to register a built-in: +/// +/// ``` +/// use datafusion_datasource::memory::MemorySourceConfig; +/// #[allow(unused_imports)] +/// use datafusion_datasource::proto::{ProtoDecoderRegistry, register_data_source}; +/// let _ = (MemorySourceConfig::try_from_proto, ProtoDecoderRegistry::new()); +/// ``` +pub fn register_data_source<T: ExtensionDataSourceFromProto>( + registry: &mut ProtoDecoderRegistry, +) -> Result<()> { + // The shared wrapper may remain registered if source registration fails. + // This is safe: it still requires a registered source decoder and rejects + // unknown sources rather than decoding their payloads through a plan codec. + DataSourceExec::register_proto(registry)?; + registry.register_decoder::<dyn DataSource, T, DataSourceDecoder>( + T::NAME, + T::try_from_proto, + ) +} + +/// Every registered extension source name, in arbitrary order. +/// Names registered for other extension kinds are not included. +pub fn data_source_names(registry: &ProtoDecoderRegistry) -> impl Iterator<Item = &str> { + registry.names::<dyn DataSource>() +} + +/// Decode the wrapper through the same registry dispatch as other extension +/// plans. Only the source-specific lookup belongs to this crate. +impl ExtensionPlanFromProto for DataSourceExec { + const NAME: &'static str = "datafusion.DataSourceExec"; + + fn try_from_proto( + node: &PhysicalPlanNode, + ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + let extension = + expect_plan_variant!(node, PhysicalPlanType::Extension, "Extension"); + if !extension.inputs.is_empty() { + return plan_err!("DataSourceExec expects no input plans"); + } + let payload = extension.node.as_slice(); + let source = PhysicalExtensionDataSourceNode::decode(payload).map_err(|e| { + plan_datafusion_err!("Failed to decode DataSourceExec payload: {e}") + })?; + let Some(name) = source + .source_name + .as_deref() + .filter(|name| !name.is_empty()) + else { + return plan_err!( + "DataSourceExec payload is missing a non-empty source name" + ); + }; + let registry = ctx + .task_ctx() + .session_config() + .get_extension::<ProtoDecoderRegistry>(); + if let Some(decoder) = registry.as_ref().and_then(|registry| { + registry.decoder::<dyn DataSource, DataSourceDecoder>(name) + }) { + return decoder(&source, ctx).map(|source| Arc::new(Self::new(source)) as _); + } + let mut names: Vec<&str> = registry + .as_ref() + .map(|registry| data_source_names(registry).collect()) + .unwrap_or_default(); + names.sort_unstable(); Review Comment: Mutation testing: deleting this line survives. No test asserts the order of names in the message. Add an `assert_eq!` on the full message in the test that registers two names. ########## datafusion/datasource/src/mod.rs: ########## @@ -44,7 +44,7 @@ pub mod projection; /// Protobuf conversions for [`FileRange`], [`PartitionedFile`] and /// [`FileGroup`], gated on the `proto` feature. #[cfg(feature = "proto")] -mod proto; +pub mod proto; Review Comment: This makes `FileRange` / `PartitionedFile` / `FileGroup` conversions public as a side effect. They become semver surface. Keep `mod proto;` private and re-export only the trait (for example, from `crate::source`)? Also, `pub mod registry` in `proto.rs` gives a second path to the same items. ########## datafusion/proto-models/proto/datafusion.proto: ########## @@ -1058,6 +1058,15 @@ message PhysicalExtensionNode { optional string plan_name = 3; } +// Payload of a DataSourceExec encoded as a PhysicalExtensionNode. The outer +// plan name identifies DataSourceExec; source names use a separate registry kind. +message PhysicalExtensionDataSourceNode { + bytes node = 1; + // Namespaced source type name. The decoding session must register a source + // decoder for this name; source payloads are never passed to a plan codec. + optional string source_name = 2; Review Comment: - `optional`, but the decoder rejects a missing or empty value. Write `required in practice` in the comment, or use a plain `string`. - `bytes node` holds the source payload, not a node. Call it `payload`. The precursor's `ExtensionDataSourceNode { string name = 1; bytes payload = 2; }` makes both changes. ########## datafusion/datasource/src/proto/registry.rs: ########## @@ -0,0 +1,416 @@ +// 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 [`DataSource`]s. +//! +//! Custom sources serialize through the existing extension-plan wire format: +//! [`DataSourceExec`] is the registered plan type, and its payload identifies +//! the source decoder. Built-in sources keep their existing wire variants. + +use std::sync::Arc; + +use datafusion_common::{Result, plan_datafusion_err, plan_err}; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::expect_plan_variant; +use datafusion_physical_plan::proto::{ + ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx, ExtensionPlanFromProto, +}; +use datafusion_proto_models::ProtoDecoderRegistry; +use datafusion_proto_models::protobuf::physical_plan_node::PhysicalPlanType; +use datafusion_proto_models::protobuf::{ + PhysicalExtensionDataSourceNode, PhysicalPlanNode, +}; +use prost::Message; + +use crate::source::{DataSource, DataSourceExec}; + +/// The wire name of an extension [`DataSource`] and the constructor that +/// rebuilds it. Only custom sources implement this trait. +/// +/// Implement [`DataSource::try_to_proto`] using [`extension_data_source_node`]. +/// The helper writes a named `DataSourceExec` plan containing the source's +/// [`Self::NAME`] and payload. The source constructor receives that typed +/// payload and can read the decoding session through +/// [`ExecutionPlanDecodeCtx::task_ctx`]. Runtime dependencies should be rebuilt +/// from that session rather than serialized. +pub trait ExtensionDataSourceFromProto: DataSource + Sized { + /// Namespace this name with the owning crate, e.g. `"my_crate.MySource"`. + /// Names are scoped to `dyn DataSource`, independently of plan names. + const NAME: &'static str; + + /// Rebuild the source from its typed extension payload. + /// [`DataSourceExec`] supplies the execution-plan wrapper. + fn try_from_proto( + node: &PhysicalExtensionDataSourceNode, + ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn DataSource>>; +} + +/// Build an extension plan for source `T`, stamping the wrapper's plan name +/// and the source's registered name. +/// +/// This uses [`ExecutionPlanEncodeCtx::extension_node`] for [`DataSourceExec`], +/// whose payload contains `T::NAME` and the source bytes. The helper lives in +/// this crate because the plan encode context cannot depend on [`DataSource`]. +/// +/// A source implementing [`ExtensionDataSourceFromProto`] can use this in its +/// [`DataSource::try_to_proto`] hook: +/// +/// ```ignore +/// fn try_to_proto( +/// &self, +/// ctx: &ExecutionPlanEncodeCtx<'_>, +/// ) -> Result<Option<PhysicalPlanNode>> { +/// let payload = self.partition_id.to_le_bytes().to_vec(); +/// Ok(Some(extension_data_source_node::<Self>(ctx, payload)?)) +/// } +/// ``` +pub fn extension_data_source_node<T: ExtensionDataSourceFromProto>( + ctx: &ExecutionPlanEncodeCtx<'_>, + payload: Vec<u8>, +) -> Result<PhysicalPlanNode> { + let source = PhysicalExtensionDataSourceNode { + node: payload, + source_name: Some(T::NAME.to_string()), + }; + ctx.extension_node::<DataSourceExec>(source.encode_to_vec(), vec![]) +} + +// Private so the decoder representation can evolve without changing the API. +type DataSourceDecoder = fn( + &PhysicalExtensionDataSourceNode, + &ExecutionPlanDecodeCtx<'_>, +) -> Result<Arc<dyn DataSource>>; + +/// Register `T` under its [`ExtensionDataSourceFromProto::NAME`], together +/// with the shared [`DataSourceExec`] extension-plan decoder. +/// +/// Registering the same type twice is a no-op; a different type under a taken +/// name, or an empty name, is an error. Source names are scoped to `dyn DataSource`. +/// Add these decoders to the same [`ProtoDecoderRegistry`] as the application's +/// plans and expressions with +/// [`SessionConfig::update_extension`](datafusion_execution::config::SessionConfig::update_extension), +/// so decoders added by other libraries are preserved. +/// +/// Built-in sources do not implement the extension trait: +/// +/// ```compile_fail +/// use datafusion_datasource::memory::MemorySourceConfig; +/// use datafusion_datasource::proto::{ProtoDecoderRegistry, register_data_source}; +/// let mut registry = ProtoDecoderRegistry::new(); +/// register_data_source::<MemorySourceConfig>(&mut registry).unwrap(); +/// ``` +/// +/// The same imports compile without attempting to register a built-in: +/// +/// ``` +/// use datafusion_datasource::memory::MemorySourceConfig; +/// #[allow(unused_imports)] +/// use datafusion_datasource::proto::{ProtoDecoderRegistry, register_data_source}; +/// let _ = (MemorySourceConfig::try_from_proto, ProtoDecoderRegistry::new()); +/// ``` Review Comment: This doctest exists only to use the imports from the `compile_fail` doctest above. I suggest removing it. The `compile_fail` doctest alone shows the point. ########## datafusion/datasource/src/proto/registry.rs: ########## @@ -0,0 +1,416 @@ +// 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 [`DataSource`]s. +//! +//! Custom sources serialize through the existing extension-plan wire format: +//! [`DataSourceExec`] is the registered plan type, and its payload identifies +//! the source decoder. Built-in sources keep their existing wire variants. + +use std::sync::Arc; + +use datafusion_common::{Result, plan_datafusion_err, plan_err}; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::expect_plan_variant; +use datafusion_physical_plan::proto::{ + ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx, ExtensionPlanFromProto, +}; +use datafusion_proto_models::ProtoDecoderRegistry; +use datafusion_proto_models::protobuf::physical_plan_node::PhysicalPlanType; +use datafusion_proto_models::protobuf::{ + PhysicalExtensionDataSourceNode, PhysicalPlanNode, +}; +use prost::Message; + +use crate::source::{DataSource, DataSourceExec}; + +/// The wire name of an extension [`DataSource`] and the constructor that +/// rebuilds it. Only custom sources implement this trait. +/// +/// Implement [`DataSource::try_to_proto`] using [`extension_data_source_node`]. +/// The helper writes a named `DataSourceExec` plan containing the source's +/// [`Self::NAME`] and payload. The source constructor receives that typed +/// payload and can read the decoding session through +/// [`ExecutionPlanDecodeCtx::task_ctx`]. Runtime dependencies should be rebuilt +/// from that session rather than serialized. +pub trait ExtensionDataSourceFromProto: DataSource + Sized { + /// Namespace this name with the owning crate, e.g. `"my_crate.MySource"`. + /// Names are scoped to `dyn DataSource`, independently of plan names. + const NAME: &'static str; + + /// Rebuild the source from its typed extension payload. + /// [`DataSourceExec`] supplies the execution-plan wrapper. + fn try_from_proto( + node: &PhysicalExtensionDataSourceNode, + ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn DataSource>>; +} + +/// Build an extension plan for source `T`, stamping the wrapper's plan name +/// and the source's registered name. +/// +/// This uses [`ExecutionPlanEncodeCtx::extension_node`] for [`DataSourceExec`], +/// whose payload contains `T::NAME` and the source bytes. The helper lives in +/// this crate because the plan encode context cannot depend on [`DataSource`]. +/// +/// A source implementing [`ExtensionDataSourceFromProto`] can use this in its +/// [`DataSource::try_to_proto`] hook: +/// +/// ```ignore +/// fn try_to_proto( +/// &self, +/// ctx: &ExecutionPlanEncodeCtx<'_>, +/// ) -> Result<Option<PhysicalPlanNode>> { +/// let payload = self.partition_id.to_le_bytes().to_vec(); +/// Ok(Some(extension_data_source_node::<Self>(ctx, payload)?)) +/// } +/// ``` +pub fn extension_data_source_node<T: ExtensionDataSourceFromProto>( + ctx: &ExecutionPlanEncodeCtx<'_>, + payload: Vec<u8>, +) -> Result<PhysicalPlanNode> { + let source = PhysicalExtensionDataSourceNode { + node: payload, + source_name: Some(T::NAME.to_string()), + }; + ctx.extension_node::<DataSourceExec>(source.encode_to_vec(), vec![]) +} + +// Private so the decoder representation can evolve without changing the API. +type DataSourceDecoder = fn( + &PhysicalExtensionDataSourceNode, + &ExecutionPlanDecodeCtx<'_>, +) -> Result<Arc<dyn DataSource>>; + +/// Register `T` under its [`ExtensionDataSourceFromProto::NAME`], together +/// with the shared [`DataSourceExec`] extension-plan decoder. +/// +/// Registering the same type twice is a no-op; a different type under a taken +/// name, or an empty name, is an error. Source names are scoped to `dyn DataSource`. +/// Add these decoders to the same [`ProtoDecoderRegistry`] as the application's +/// plans and expressions with +/// [`SessionConfig::update_extension`](datafusion_execution::config::SessionConfig::update_extension), +/// so decoders added by other libraries are preserved. +/// +/// Built-in sources do not implement the extension trait: +/// +/// ```compile_fail +/// use datafusion_datasource::memory::MemorySourceConfig; +/// use datafusion_datasource::proto::{ProtoDecoderRegistry, register_data_source}; +/// let mut registry = ProtoDecoderRegistry::new(); +/// register_data_source::<MemorySourceConfig>(&mut registry).unwrap(); +/// ``` +/// +/// The same imports compile without attempting to register a built-in: +/// +/// ``` +/// use datafusion_datasource::memory::MemorySourceConfig; +/// #[allow(unused_imports)] +/// use datafusion_datasource::proto::{ProtoDecoderRegistry, register_data_source}; +/// let _ = (MemorySourceConfig::try_from_proto, ProtoDecoderRegistry::new()); +/// ``` +pub fn register_data_source<T: ExtensionDataSourceFromProto>( + registry: &mut ProtoDecoderRegistry, +) -> Result<()> { + // The shared wrapper may remain registered if source registration fails. + // This is safe: it still requires a registered source decoder and rejects + // unknown sources rather than decoding their payloads through a plan codec. + DataSourceExec::register_proto(registry)?; + registry.register_decoder::<dyn DataSource, T, DataSourceDecoder>( + T::NAME, + T::try_from_proto, + ) +} + +/// Every registered extension source name, in arbitrary order. +/// Names registered for other extension kinds are not included. +pub fn data_source_names(registry: &ProtoDecoderRegistry) -> impl Iterator<Item = &str> { Review Comment: Is this needed in the public API? Outside tests, it is used only in the error message at L183. `ProtoDecoderRegistry::names::<dyn DataSource>()` is already public. ```suggestion pub(crate) fn data_source_names(registry: &ProtoDecoderRegistry) -> impl Iterator<Item = &str> { ``` -- 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]
