gabotechs commented on code in PR #26555: URL: https://github.com/apache/datafusion/pull/26555#discussion_r4336903193
########## datafusion/datasource/src/proto/registry.rs: ########## @@ -0,0 +1,389 @@ +// 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, + register_execution_plan, +}; +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. Libraries should add to the application's registry +/// rather than replace it with their own. +/// +/// 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<()> { + register_execution_plan::<DataSourceExec>(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"); + } Review Comment: The fact that `DataSource`s are entities nested under `DataSourceExec` (an `ExecutionPlan` impl) complicates things a bit. Having that kind of nesting makes things a bit weird, as `ExecutionPlan`s are also serializable entities themselves. One weird thing that happens because of this: `DataSourceExec` nodes baked by Parquet or CSV have their own proto representation, which are handled outside of the registry. It's just when the `DataSource` is baked by something other than Parquet or CSV when this `impl ExtensionPlanFromProto for DataSourceExec` mechanisms kicks in. This makes the approach of [de]serializing custom DataSources and built-in DataSource significantly different. Something clean that comes to mind is: rather than separating [de]serialization logic for built-in entities VS custom entities, use this registry path for everything, and dogfood it with every built-in implementation that needs to go over the wire. This is a very large refactor though... -- 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]
