From 774dfd5d6285dae7454501b5e9c6bf46863f4931 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Wed, 21 Oct 2026 11:19:27 +0200 Subject: [PATCH] feat: register extension DataSource decoders through DataSourceExec --- Cargo.lock | 1 + datafusion/datasource/Cargo.toml | 2 + datafusion/datasource/src/mod.rs | 2 +- datafusion/datasource/src/proto.rs | 12 + datafusion/datasource/src/proto/registry.rs | 416 ++++++++++++ datafusion/datasource/src/source.rs | 4 + .../physical-plan/src/proto/registry.rs | 16 +- .../proto-models/proto/datafusion.proto | 9 + .../proto-models/src/generated/pbjson.rs | 113 ++++ .../proto-models/src/generated/prost.rs | 11 + .../tests/cases/plans/data_source_registry.rs | 608 ++++++++++++++++++ datafusion/proto/tests/cases/plans/mod.rs | 1 + .../library-user-guide/upgrading/56.0.0.md | 32 +- 13 files changed, 1219 insertions(+), 8 deletions(-) create mode 100644 datafusion/datasource/src/proto/registry.rs create mode 100644 datafusion/proto/tests/cases/plans/data_source_registry.rs diff --git a/Cargo.lock b/Cargo.lock index cbf3c06465bef..d3a0db08987a4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1976,6 +1976,7 @@ dependencies = [ "object_store", "parking_lot 0.13.0", "pin-project-lite", + "prost", "rand 0.9.4", "tempfile", "tokio", diff --git a/datafusion/datasource/Cargo.toml b/datafusion/datasource/Cargo.toml index dd4111cf958e8..4fa81875f857a 100644 --- a/datafusion/datasource/Cargo.toml +++ b/datafusion/datasource/Cargo.toml @@ -39,6 +39,7 @@ default = ["compression"] # consumers that never serialize plans pay nothing. Mirrors the `proto` feature # on `datafusion-physical-plan`. proto = [ + "dep:prost", "dep:datafusion-proto-models", "datafusion-physical-plan/proto", ] @@ -75,6 +76,7 @@ log = { workspace = true } object_store = { workspace = true } parking_lot = { workspace = true } pin-project-lite = { workspace = true } +prost = { workspace = true, optional = true } rand = { workspace = true } tempfile = { workspace = true, optional = true } tokio = { workspace = true } diff --git a/datafusion/datasource/src/mod.rs b/datafusion/datasource/src/mod.rs index bbe14280a6b12..5d80fb9e2fc77 100644 --- a/datafusion/datasource/src/mod.rs +++ b/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; pub mod schema_adapter; pub mod sink; pub mod source; diff --git a/datafusion/datasource/src/proto.rs b/datafusion/datasource/src/proto.rs index 6dc6e2ee45ffd..2d226622c2985 100644 --- a/datafusion/datasource/src/proto.rs +++ b/datafusion/datasource/src/proto.rs @@ -15,6 +15,10 @@ // specific language governing permissions and limitations // under the License. +//! Serialization support for data sources. +//! +//! Custom sources use the name-keyed decoder API in [`registry`](crate::proto::registry). +//! //! Protobuf conversions for the file-scan leaf types owned by this crate: //! [`FileRange`], [`PartitionedFile`] and [`FileGroup`]. //! @@ -30,6 +34,14 @@ //! the value, and the orphan rule allows it here because one side of each //! conversion is a type this crate owns. +pub mod registry; + +pub use datafusion_proto_models::ProtoDecoderRegistry; +pub use registry::{ + ExtensionDataSourceFromProto, data_source_names, extension_data_source_node, + register_data_source, +}; + use std::sync::Arc; use chrono::{TimeZone, Utc}; diff --git a/datafusion/datasource/src/proto/registry.rs b/datafusion/datasource/src/proto/registry.rs new file mode 100644 index 0000000000000..7d9176867d31c --- /dev/null +++ b/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>; +} + +/// 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> { +/// let payload = self.partition_id.to_le_bytes().to_vec(); +/// Ok(Some(extension_data_source_node::(ctx, payload)?)) +/// } +/// ``` +pub fn extension_data_source_node( + ctx: &ExecutionPlanEncodeCtx<'_>, + payload: Vec, +) -> Result { + let source = PhysicalExtensionDataSourceNode { + node: payload, + source_name: Some(T::NAME.to_string()), + }; + ctx.extension_node::(source.encode_to_vec(), vec![]) +} + +// Private so the decoder representation can evolve without changing the API. +type DataSourceDecoder = fn( + &PhysicalExtensionDataSourceNode, + &ExecutionPlanDecodeCtx<'_>, +) -> Result>; + +/// 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::(&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( + 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::( + 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 { + registry.names::() +} + +/// 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> { + 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::(); + if let Some(decoder) = registry.as_ref().and_then(|registry| { + registry.decoder::(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(); + let names = if names.is_empty() { + "none".to_string() + } else { + names.join(", ") + }; + plan_err!( + "No decoder is registered for the extension DataSource '{name}'. Register \ + the source in the ProtoDecoderRegistry attached to the decoding session's \ + SessionConfig. Registered extension data sources: {names}" + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fmt::{self, Formatter}; + + use datafusion_common::{Statistics, internal_err, tree_node::TreeNodeRecursion}; + use datafusion_execution::TaskContext; + use datafusion_physical_expr::projection::ProjectionExprs; + use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; + use datafusion_physical_plan::proto::execution_plan_names; + use datafusion_physical_plan::{DisplayFormatType, SendableRecordBatchStream}; + + use crate::memory::MemorySourceConfig; + + // Registration-only fixtures. The memory source supplies DataSource's + // unrelated methods; these tests never invoke a decoder. + macro_rules! registration_source { + ($name:ident) => { + #[derive(Debug)] + struct $name(MemorySourceConfig); + + impl DataSource for $name { + fn open( + &self, + partition: usize, + ctx: Arc, + ) -> Result { + self.0.open(partition, ctx) + } + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + self.0.fmt_as(t, f) + } + fn output_partitioning(&self) -> Partitioning { + self.0.output_partitioning() + } + fn eq_properties(&self) -> EquivalenceProperties { + self.0.eq_properties() + } + fn partition_statistics( + &self, + partition: Option, + ) -> Result> { + self.0.partition_statistics(partition) + } + fn with_fetch( + &self, + limit: Option, + ) -> Option> { + self.0.with_fetch(limit) + } + fn fetch(&self) -> Option { + self.0.fetch() + } + fn try_swapping_with_projection( + &self, + projection: &ProjectionExprs, + ) -> Result>> { + self.0.try_swapping_with_projection(projection) + } + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + self.0.apply_expressions(f) + } + } + }; + } + + registration_source!(RegisteredSource); + + impl ExtensionDataSourceFromProto for RegisteredSource { + const NAME: &'static str = "test.RegisteredSource"; + + fn try_from_proto( + _node: &PhysicalExtensionDataSourceNode, + _ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + internal_err!("registration-only test source") + } + } + + registration_source!(CollidingSource); + + impl ExtensionDataSourceFromProto for CollidingSource { + const NAME: &'static str = "test.RegisteredSource"; + + fn try_from_proto( + _node: &PhysicalExtensionDataSourceNode, + _ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + internal_err!("registration-only test source") + } + } + + registration_source!(OtherSource); + + impl ExtensionDataSourceFromProto for OtherSource { + const NAME: &'static str = "other.OtherSource"; + + fn try_from_proto( + _node: &PhysicalExtensionDataSourceNode, + _ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + internal_err!("registration-only test source") + } + } + + registration_source!(UnnamedSource); + + impl ExtensionDataSourceFromProto for UnnamedSource { + const NAME: &'static str = ""; + + fn try_from_proto( + _node: &PhysicalExtensionDataSourceNode, + _ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + internal_err!("registration-only test source") + } + } + + #[test] + fn register_is_idempotent_for_the_same_type() -> Result<()> { + let mut registry = ProtoDecoderRegistry::new(); + register_data_source::(&mut registry)?; + register_data_source::(&mut registry)?; + assert_eq!( + data_source_names(®istry).collect::>(), + [RegisteredSource::NAME] + ); + Ok(()) + } + + #[test] + fn register_rejects_a_name_collision_between_types() -> Result<()> { + let mut registry = ProtoDecoderRegistry::new(); + register_data_source::(&mut registry)?; + let err = register_data_source::(&mut registry) + .expect_err("colliding registration must fail"); + let message = err.to_string(); + assert!(message.contains(RegisteredSource::NAME), "{message}"); + assert!(message.contains("already registered"), "{message}"); + assert_eq!( + data_source_names(®istry).collect::>(), + [RegisteredSource::NAME] + ); + Ok(()) + } + + #[test] + fn register_keeps_distinct_names_apart() -> Result<()> { + let mut registry = ProtoDecoderRegistry::new(); + register_data_source::(&mut registry)?; + register_data_source::(&mut registry)?; + let mut names: Vec<_> = data_source_names(®istry).collect(); + names.sort_unstable(); + assert_eq!(names, [OtherSource::NAME, RegisteredSource::NAME]); + Ok(()) + } + + #[test] + fn register_rejects_an_empty_source_name() { + let mut registry = ProtoDecoderRegistry::new(); + let err = register_data_source::(&mut registry) + .expect_err("an empty NAME must fail"); + assert!(err.to_string().contains("empty name"), "{err}"); + assert!(data_source_names(®istry).next().is_none()); + assert_eq!( + execution_plan_names(®istry).collect::>(), + [DataSourceExec::NAME] + ); + } + + #[test] + fn register_rejects_a_wrapper_collision_without_registering_the_source() -> Result<()> + { + let mut registry = ProtoDecoderRegistry::new(); + registry + .register_decoder::(DataSourceExec::NAME, 7)?; + let err = register_data_source::(&mut registry) + .expect_err("a colliding wrapper registration must fail"); + assert!(err.to_string().contains("already registered"), "{err}"); + assert!(data_source_names(®istry).next().is_none()); + assert_eq!( + registry.decoder::(DataSourceExec::NAME), + Some(&7) + ); + Ok(()) + } + + #[test] + fn a_name_registered_for_another_kind_is_not_a_data_source() -> Result<()> { + trait OtherKind {} + let mut registry = ProtoDecoderRegistry::new(); + registry.register_decoder::(RegisteredSource::NAME, 7)?; + assert!(data_source_names(®istry).next().is_none()); + register_data_source::(&mut registry)?; + assert_eq!( + data_source_names(®istry).collect::>(), + [RegisteredSource::NAME] + ); + Ok(()) + } + + #[test] + fn register_source_also_registers_the_shared_plan_decoder() -> Result<()> { + let mut registry = ProtoDecoderRegistry::new(); + register_data_source::(&mut registry)?; + register_data_source::(&mut registry)?; + assert_eq!( + execution_plan_names(®istry).collect::>(), + [DataSourceExec::NAME] + ); + Ok(()) + } +} diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index 32a300ff31b74..a071992e158fc 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -290,6 +290,10 @@ pub trait DataSource: Any + Send + Sync + Debug { /// Serialize this data source to a full [`PhysicalPlanNode`] (a /// `DataSourceExec` wrapping this source), if it knows how. /// + /// Custom sources can implement [`ExtensionDataSourceFromProto`](crate::proto::ExtensionDataSourceFromProto) + /// and use [`extension_data_source_node`](crate::proto::extension_data_source_node) + /// to encode the `DataSourceExec` wrapper and decode through the session registry. + /// /// This is the `DataSource` analog of /// [`ExecutionPlan::try_to_proto`]. /// [`DataSourceExec::try_to_proto`](crate::source::DataSourceExec) delegates diff --git a/datafusion/physical-plan/src/proto/registry.rs b/datafusion/physical-plan/src/proto/registry.rs index 0be6615231d6b..a30b8f3cac41f 100644 --- a/datafusion/physical-plan/src/proto/registry.rs +++ b/datafusion/physical-plan/src/proto/registry.rs @@ -26,8 +26,10 @@ //! An extension plan implements [`ExtensionPlanFromProto`] instead, and //! [`ExtensionPlanFromProto::register_proto`] 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. +//! so registration order does not matter. Built-ins with dedicated wire variants +//! generally cannot be registered. `DataSourceExec` is an exception: its registered +//! decoder wraps custom data sources, while built-in scans keep their wire variants. +//! 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. @@ -47,8 +49,9 @@ 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 +/// Extension plans implement this, as does the shared `DataSourceExec` wrapper +/// for custom data sources. Other built-ins keep their dedicated wire variants. +/// One impl block carries both halves, and /// [`ExecutionPlan::try_to_proto`] writes the matching node: /// /// ```ignore @@ -111,8 +114,9 @@ pub trait ExtensionPlanFromProto: ExecutionPlan + Sized { /// `config.update_extension(MyExec::register_proto)?`, which keeps the decoders /// that other libraries added. See [`ProtoDecoderRegistry`]. /// - /// A built-in plan does not implement this trait, so it cannot be - /// registered: + /// Most built-in plans do not implement this trait, so they cannot be + /// registered. `DataSourceExec` implements it to wrap custom data sources; + /// built-ins with dedicated wire variants, such as `FilterExec`, do not: /// /// ```compile_fail /// use datafusion_physical_plan::filter::FilterExec; diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 6c4ef8f0a4b2f..de1eb317e50db 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/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; +} + // physical expressions message PhysicalExprNode { // Was date_time_interval_expr diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 725cceeb75970..3c0ba8ad28ffa 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -20282,6 +20282,119 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { deserializer.deserialize_struct("datafusion.PhysicalExprNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for PhysicalExtensionDataSourceNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.node.is_empty() { + len += 1; + } + if self.source_name.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalExtensionDataSourceNode", len)?; + if !self.node.is_empty() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("node", pbjson::private::base64::encode(&self.node).as_str())?; + } + if let Some(v) = self.source_name.as_ref() { + struct_ser.serialize_field("sourceName", v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalExtensionDataSourceNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "node", + "source_name", + "sourceName", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Node, + SourceName, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "node" => Ok(GeneratedField::Node), + "sourceName" | "source_name" => Ok(GeneratedField::SourceName), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalExtensionDataSourceNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalExtensionDataSourceNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut node__ = None; + let mut source_name__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Node => { + if node__.is_some() { + return Err(serde::de::Error::duplicate_field("node")); + } + node__ = + Some(map_.next_value::<::pbjson::private::BytesDeserialize<_>>()?.0) + ; + } + GeneratedField::SourceName => { + if source_name__.is_some() { + return Err(serde::de::Error::duplicate_field("sourceName")); + } + source_name__ = map_.next_value()?; + } + } + } + Ok(PhysicalExtensionDataSourceNode { + node: node__.unwrap_or_default(), + source_name: source_name__, + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalExtensionDataSourceNode", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for PhysicalExtensionExprNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 1034d0443dbce..3979f1291d6dd 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1634,6 +1634,17 @@ pub struct PhysicalExtensionNode { #[prost(string, optional, tag = "3")] pub plan_name: ::core::option::Option<::prost::alloc::string::String>, } +/// Payload of a DataSourceExec encoded as a PhysicalExtensionNode. The outer +/// plan name identifies DataSourceExec; source names use a separate registry kind. +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct PhysicalExtensionDataSourceNode { + #[prost(bytes = "vec", tag = "1")] + pub node: ::prost::alloc::vec::Vec, + /// Namespaced source type name. The decoding session must register a source + /// decoder for this name; source payloads are never passed to a plan codec. + #[prost(string, optional, tag = "2")] + pub source_name: ::core::option::Option<::prost::alloc::string::String>, +} /// physical expressions #[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalExprNode { diff --git a/datafusion/proto/tests/cases/plans/data_source_registry.rs b/datafusion/proto/tests/cases/plans/data_source_registry.rs new file mode 100644 index 0000000000000..d6f3bbbb16202 --- /dev/null +++ b/datafusion/proto/tests/cases/plans/data_source_registry.rs @@ -0,0 +1,608 @@ +// 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. + +//! Public serialization APIs for custom sources, including execution after decoding. + +use std::fmt::{self, Formatter}; +use std::sync::Arc; + +use datafusion::arrow::array::{ArrayRef, Int64Array}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::projection::ProjectionExprs; +use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion::physical_plan::proto::{ + ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx, ExtensionPlanFromProto, +}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, PlanProperties, + SendableRecordBatchStream, collect, +}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{DataFusionError, Result, Statistics, internal_err}; +use datafusion_datasource::memory::MemorySourceConfig; +use datafusion_datasource::proto::{ + ExtensionDataSourceFromProto, ProtoDecoderRegistry, extension_data_source_node, + register_data_source, +}; +use datafusion_datasource::source::{DataSource, DataSourceExec}; +use datafusion_proto::bytes::{ + physical_plan_from_bytes_with_extension_codec, + physical_plan_to_bytes_with_extension_codec, +}; +use datafusion_proto::physical_plan::{ + AsExecutionPlan, DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, + PhysicalProtoConverterExtension, +}; +use datafusion_proto::protobuf::physical_plan_node::PhysicalPlanType; +use datafusion_proto::protobuf::{ + PhysicalExtensionDataSourceNode, PhysicalExtensionNode, PhysicalPlanNode, +}; +use prost::Message; + +#[derive(Debug)] +struct ReaderState(i64); + +// Share the DataSource boilerplate, as the extension-plan tests do for their +// execution plans. Each source's serialization behavior remains explicit. +macro_rules! memory_source { + ($name:ident, $($hook:item)*) => { + #[derive(Debug, Clone)] + struct $name { + value: i64, + memory: MemorySourceConfig, + } + + impl $name { + fn new(value: i64, reader_state: Arc) -> Result { + let array: ArrayRef = Arc::new(Int64Array::from(vec![value + reader_state.0])); + let batch = RecordBatch::try_from_iter([("a", array)])?; + let memory = MemorySourceConfig::try_new(&[vec![batch.clone()]], batch.schema(), None)?; + Ok(Self { value, memory }) + } + } + + impl DataSource for $name { + fn open(&self, partition: usize, ctx: Arc) -> Result { + self.memory.open(partition, ctx) + } + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "{}({})", stringify!($name), self.value) + } + fn output_partitioning(&self) -> Partitioning { self.memory.output_partitioning() } + fn eq_properties(&self) -> EquivalenceProperties { self.memory.eq_properties() } + fn partition_statistics(&self, partition: Option) -> Result> { + self.memory.partition_statistics(partition) + } + fn with_fetch(&self, _limit: Option) -> Option> { None } + fn fetch(&self) -> Option { None } + fn try_swapping_with_projection(&self, _projection: &ProjectionExprs) -> Result>> { Ok(None) } + fn apply_expressions(&self, f: &mut dyn FnMut(&Arc) -> Result) -> Result { + self.memory.apply_expressions(f) + } + $($hook)* + } + }; +} + +memory_source!( + RegisteredSource, + fn try_to_proto( + &self, + ctx: &ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(Some(extension_data_source_node::( + ctx, + self.value.to_le_bytes().to_vec(), + )?)) + } +); + +impl ExtensionDataSourceFromProto for RegisteredSource { + const NAME: &'static str = "test.RegisteredSource"; + + fn try_from_proto( + node: &PhysicalExtensionDataSourceNode, + ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + let value = read_value(&node.node)?; + let state = ctx + .task_ctx() + .session_config() + .get_extension::() + .unwrap_or_else(|| Arc::new(ReaderState(0))); + Ok(Arc::new(Self::new(value, state)?)) + } +} + +memory_source!( + OtherSource, + fn try_to_proto( + &self, + ctx: &ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(Some(extension_data_source_node::( + ctx, + self.value.to_le_bytes().to_vec(), + )?)) + } +); + +impl ExtensionDataSourceFromProto for OtherSource { + const NAME: &'static str = "other.OtherSource"; + + fn try_from_proto( + node: &PhysicalExtensionDataSourceNode, + ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + let value = read_value(&node.node)?; + let state = ctx + .task_ctx() + .session_config() + .get_extension::() + .unwrap_or_else(|| Arc::new(ReaderState(0))); + Ok(Arc::new(Self::new(value, state)?)) + } +} + +// No serialization hook: this source is owned by the legacy codec. +memory_source!(LegacySource,); + +fn read_value(buf: &[u8]) -> Result { + let Ok(bytes) = buf.try_into() else { + return internal_err!("expected eight bytes"); + }; + Ok(i64::from_le_bytes(bytes)) +} + +fn source(value: i64) -> Result> { + Ok(DataSourceExec::from_data_source(RegisteredSource::new( + value, + Arc::new(ReaderState(0)), + )?)) +} + +fn session( + register: impl FnOnce(&mut ProtoDecoderRegistry) -> Result<()>, +) -> Result { + let mut registry = ProtoDecoderRegistry::new(); + register(&mut registry)?; + let config = SessionConfig::new().with_extension(Arc::new(registry)); + Ok(SessionContext::new_with_config(config)) +} + +fn decoded_source(plan: &Arc) -> &T { + plan.downcast_ref::() + .unwrap() + .data_source() + .downcast_ref::() + .unwrap() +} + +fn encode(plan: Arc) -> Result { + PhysicalPlanNode::try_from_physical_plan(plan, &DefaultPhysicalExtensionCodec {}) +} + +fn decode( + node: &PhysicalPlanNode, + ctx: &SessionContext, + codec: &dyn PhysicalExtensionCodec, +) -> Result> { + // Exercise the actual extension-plan wire format. + let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice()).unwrap(); + node.try_into_physical_plan(ctx.task_ctx().as_ref(), codec) +} + +fn extension_of(node: &PhysicalPlanNode) -> &PhysicalExtensionNode { + let Some(PhysicalPlanType::Extension(extension)) = &node.physical_plan_type else { + panic!("expected an extension plan"); + }; + extension +} + +fn extension_of_mut(node: &mut PhysicalPlanNode) -> &mut PhysicalExtensionNode { + let Some(PhysicalPlanType::Extension(extension)) = &mut node.physical_plan_type + else { + panic!("expected an extension plan"); + }; + extension +} + +fn source_payload(node: &PhysicalPlanNode) -> PhysicalExtensionDataSourceNode { + PhysicalExtensionDataSourceNode::decode(extension_of(node).node.as_slice()) + .expect("valid source payload") +} + +fn update_source_payload( + node: &mut PhysicalPlanNode, + update: impl FnOnce(&mut PhysicalExtensionDataSourceNode), +) { + let mut source = source_payload(node); + update(&mut source); + extension_of_mut(node).node = source.encode_to_vec(); +} + +#[tokio::test] +async fn source_registry_roundtrip_executes_without_codec() -> Result<()> { + let mut registry = ProtoDecoderRegistry::new(); + register_data_source::(&mut registry)?; + let config = SessionConfig::new().with_extension(Arc::new(registry)); + let reader = SessionContext::new_with_config(config); + let codec = DefaultPhysicalExtensionCodec {}; + let original = source(7)?; + let node = encode(original.clone())?; + assert_eq!( + extension_of(&node).plan_name.as_deref(), + Some(DataSourceExec::NAME) + ); + assert!(extension_of(&node).inputs.is_empty()); + let payload = source_payload(&node); + assert_eq!(payload.source_name.as_deref(), Some(RegisteredSource::NAME)); + assert_eq!(payload.node, 7_i64.to_le_bytes()); + let bytes = physical_plan_to_bytes_with_extension_codec(original.clone(), &codec)?; + let decoded = physical_plan_from_bytes_with_extension_codec( + &bytes, + reader.task_ctx().as_ref(), + &codec, + )?; + let source = decoded_source::(&decoded); + assert_eq!(source.value, 7); + assert_eq!(original.schema(), decoded.schema()); + let batches = collect(decoded, reader.task_ctx()).await?; + assert_eq!( + batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 7 + ); + Ok(()) +} + +#[tokio::test] +async fn source_decode_reads_decoding_session() -> Result<()> { + let mut registry = ProtoDecoderRegistry::new(); + register_data_source::(&mut registry)?; + let config = SessionConfig::new() + .with_extension(Arc::new(ReaderState(100))) + .with_extension(Arc::new(registry)); + let reader = SessionContext::new_with_config(config); + let decoded = decode( + &encode(source(7)?)?, + &reader, + &DefaultPhysicalExtensionCodec {}, + )?; + let batches = collect(decoded, reader.task_ctx()).await?; + assert_eq!( + batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 107 + ); + Ok(()) +} + +#[test] +fn independent_source_types_share_one_registry() -> Result<()> { + let ctx = session(|r| { + register_data_source::(r)?; + register_data_source::(r) + })?; + let codec = DefaultPhysicalExtensionCodec {}; + let decoded = decode(&encode(source(11)?)?, &ctx, &codec)?; + assert_eq!(decoded_source::(&decoded).value, 11); + + let other = + DataSourceExec::from_data_source(OtherSource::new(22, Arc::new(ReaderState(0)))?); + let decoded = decode(&encode(other)?, &ctx, &codec)?; + assert_eq!(decoded_source::(&decoded).value, 22); + Ok(()) +} + +/// Codec for sources that leave serialization to PhysicalExtensionCodec. +#[derive(Debug)] +struct LegacyCodec {} + +impl PhysicalExtensionCodec for LegacyCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + _ctx: &TaskContext, + _converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + assert!(inputs.is_empty()); + Ok(DataSourceExec::from_data_source(LegacySource::new( + read_value(buf)?, + Arc::new(ReaderState(0)), + )?)) + } + fn try_encode( + &self, + plan: Arc, + buf: &mut Vec, + _converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + buf.extend_from_slice(&decoded_source::(&plan).value.to_le_bytes()); + Ok(()) + } +} + +// This codec would accept even a malformed source payload. It makes accidental +// fallback after a registered decoder's error observable. +#[derive(Debug)] +struct GreedyCodec {} + +impl PhysicalExtensionCodec for GreedyCodec { + fn try_decode( + &self, + _buf: &[u8], + inputs: &[Arc], + _ctx: &TaskContext, + _converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + assert!(inputs.is_empty()); + Ok(DataSourceExec::from_data_source(LegacySource::new( + 0, + Arc::new(ReaderState(0)), + )?)) + } + fn try_encode( + &self, + _plan: Arc, + _buf: &mut Vec, + _converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + internal_err!("GreedyCodec does not encode") + } +} + +#[test] +fn registered_source_does_not_call_plan_codec() -> Result<()> { + let ctx = session(register_data_source::)?; + let decoded = decode(&encode(source(9)?)?, &ctx, &GreedyCodec {})?; + assert_eq!(decoded_source::(&decoded).value, 9); + Ok(()) +} + +#[test] +fn registered_decoder_failure_does_not_fall_back_to_codec() -> Result<()> { + let ctx = session(register_data_source::)?; + let mut node = encode(source(9)?)?; + // The source's real decoder requires eight bytes for its i64 payload. + update_source_payload(&mut node, |source| source.node = vec![0]); + let err = decode(&node, &ctx, &GreedyCodec {}) + .expect_err("a registered decoder's failure must be fatal"); + assert!(err.to_string().contains("expected eight bytes"), "{err}"); + Ok(()) +} + +#[test] +fn malformed_wrapper_payload_does_not_fall_back_to_codec() -> Result<()> { + let ctx = session(register_data_source::)?; + let mut node = encode(source(9)?)?; + // Zero is not a valid protobuf tag. + extension_of_mut(&mut node).node = vec![0]; + let err = decode(&node, &ctx, &GreedyCodec {}).unwrap_err(); + assert!( + err.to_string() + .contains("Failed to decode DataSourceExec payload"), + "{err}" + ); + Ok(()) +} + +#[test] +fn source_wrapper_rejects_input_plans() -> Result<()> { + let ctx = session(register_data_source::)?; + let mut node = encode(source(9)?)?; + extension_of_mut(&mut node).inputs.push(encode(source(1)?)?); + let err = decode(&node, &ctx, &GreedyCodec {}).unwrap_err(); + assert!( + err.to_string() + .contains("DataSourceExec expects no input plans"), + "{err}" + ); + Ok(()) +} + +#[test] +fn unregistered_sources_do_not_fall_back_to_plan_codec() -> Result<()> { + let node = encode(source(9)?)?; + for ctx in [ + session(DataSourceExec::register_proto)?, + session(register_data_source::)?, + ] { + // The greedy codec accepts these bytes, so falling back would silently + // return the wrong plan instead of reporting the missing registration. + let err = decode(&node, &ctx, &GreedyCodec {}) + .expect_err("an unregistered source must not reach the plan codec"); + assert!(matches!(err.find_root(), DataFusionError::Plan(_))); + let message = err.to_string(); + assert!(message.contains("No decoder is registered"), "{message}"); + assert!(message.contains(RegisteredSource::NAME), "{message}"); + } + Ok(()) +} + +#[test] +fn unnamed_sources_do_not_fall_back_to_plan_codec() -> Result<()> { + let ctx = session(register_data_source::)?; + for source_name in [None, Some(String::new())] { + let mut node = encode(source(9)?)?; + update_source_payload(&mut node, |source| source.source_name = source_name); + let err = decode(&node, &ctx, &GreedyCodec {}) + .expect_err("an unnamed source must not reach the plan codec"); + assert!(matches!(err.find_root(), DataFusionError::Plan(_))); + assert!( + err.to_string().contains("missing a non-empty source name"), + "{err}" + ); + } + Ok(()) +} + +#[test] +fn missing_source_registration_reports_registered_names() -> Result<()> { + let ctx = session(register_data_source::)?; + let node = encode(source(9)?)?; + let err = decode(&node, &ctx, &DefaultPhysicalExtensionCodec {}).unwrap_err(); + let message = err.to_string(); + assert!(message.contains(RegisteredSource::NAME), "{message}"); + assert!(message.contains(OtherSource::NAME), "{message}"); + assert!(message.contains("No decoder is registered"), "{message}"); + let err = decode( + &node, + &session(DataSourceExec::register_proto)?, + &DefaultPhysicalExtensionCodec {}, + ) + .unwrap_err(); + assert!(err.to_string().contains("extension data sources: none")); + Ok(()) +} + +#[test] +fn unregistered_wrapper_uses_existing_plan_codec_fallback() -> Result<()> { + let node = encode(source(9)?)?; + let ctx = SessionContext::new(); + // With no DataSourceExec decoder registered, the outer node is an unknown + // extension plan. The existing plan-codec policy still applies to it. + let decoded = decode(&node, &ctx, &GreedyCodec {})?; + assert_eq!(decoded_source::(&decoded).value, 0); + let err = decode(&node, &ctx, &DefaultPhysicalExtensionCodec {}).unwrap_err(); + assert!(err.to_string().contains(DataSourceExec::NAME), "{err}"); + Ok(()) +} + +#[test] +fn codec_owned_sources_keep_the_legacy_plan_wire_format() -> Result<()> { + let ctx = session(register_data_source::)?; + let codec = LegacyCodec {}; + let node = PhysicalPlanNode::try_from_physical_plan( + DataSourceExec::from_data_source(LegacySource::new( + 33, + Arc::new(ReaderState(0)), + )?), + &codec, + )?; + let Some(PhysicalPlanType::Extension(extension)) = &node.physical_plan_type else { + panic!("expected legacy extension node") + }; + assert_eq!(extension.plan_name, None); + assert_eq!( + decoded_source::(&decode(&node, &ctx, &codec)?).value, + 33 + ); + Ok(()) +} + +/// A custom plan with the same wire name as RegisteredSource. +#[derive(Debug)] +struct NamedPlan(Arc); + +impl DisplayAs for NamedPlan { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "NamedPlan") + } +} + +impl ExecutionPlan for NamedPlan { + fn name(&self) -> &str { + "NamedPlan" + } + fn properties(&self) -> &Arc { + self.0.properties() + } + fn children(&self) -> Vec<&Arc> { + vec![&self.0] + } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(Self(Arc::clone(&children[0])))) + } + fn execute( + &self, + partition: usize, + ctx: Arc, + ) -> Result { + self.0.execute(partition, ctx) + } + fn try_to_proto( + &self, + ctx: &ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(Some(ctx.extension_node::(vec![], self.children())?)) + } +} + +impl ExtensionPlanFromProto for NamedPlan { + const NAME: &'static str = RegisteredSource::NAME; + fn try_from_proto( + node: &PhysicalPlanNode, + ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + let Some(PhysicalPlanType::Extension(extension)) = &node.physical_plan_type + else { + return internal_err!("expected plan extension"); + }; + if extension.inputs.len() != 1 { + return internal_err!("expected one child"); + } + Ok(Arc::new(Self(ctx.decode_child(&extension.inputs[0])?))) + } +} + +#[test] +fn plan_and_source_with_same_name_roundtrip_together() -> Result<()> { + let ctx = session(|r| { + NamedPlan::register_proto(r)?; + register_data_source::(r) + })?; + let plan: Arc = Arc::new(NamedPlan(source(42)?)); + let decoded = decode(&encode(plan)?, &ctx, &DefaultPhysicalExtensionCodec {})?; + let decoded = decoded.downcast_ref::().unwrap(); + assert_eq!(decoded_source::(&decoded.0).value, 42); + Ok(()) +} + +#[cfg(feature = "json")] +#[test] +fn source_registry_json_roundtrip() -> Result<()> { + let ctx = session(register_data_source::)?; + let node = encode(source(55)?)?; + let json = serde_json::to_string(&node).unwrap(); + assert!(json.contains("planName")); + assert!(json.contains(DataSourceExec::NAME)); + let node = serde_json::from_str(&json).unwrap(); + let decoded = decode(&node, &ctx, &DefaultPhysicalExtensionCodec {})?; + assert_eq!(decoded_source::(&decoded).value, 55); + Ok(()) +} diff --git a/datafusion/proto/tests/cases/plans/mod.rs b/datafusion/proto/tests/cases/plans/mod.rs index 69a6256af9263..621a36cb3f218 100644 --- a/datafusion/proto/tests/cases/plans/mod.rs +++ b/datafusion/proto/tests/cases/plans/mod.rs @@ -34,6 +34,7 @@ use datafusion_proto::physical_plan::{ use std::sync::Arc; mod aggregates; +mod data_source_registry; mod dispatch; mod dynamic_filters; mod expr_registry; diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index ba87076b3de10..45b218adb42ed 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -25,6 +25,35 @@ in this section pertains to features and changes that have already been merged to the main branch and are awaiting release in this version. +### Custom `DataSource` decoder registration + +Custom sources can now implement +`datafusion_datasource::proto::ExtensionDataSourceFromProto` and write their +`DataSource::try_to_proto` node with +`extension_data_source_node::(ctx, payload)?`. +`register_data_source::(&mut registry)` registers both the source and +the shared `DataSourceExec` plan decoder in the session's `ProtoDecoderRegistry`. + +```rust,ignore +let mut config = SessionConfig::new(); +config.update_extension(|registry: &mut ProtoDecoderRegistry| { + register_data_source::(registry) +})?; +``` + +The existing plan-extension wire format identifies `datafusion.DataSourceExec`. +Its payload carries the source name and source bytes. The wrapper's decoder +looks up the source and reconstructs `DataSourceExec` from `Arc`. +Source constructors receive `PhysicalExtensionDataSourceNode` and can read the +receiving session through `ExecutionPlanDecodeCtx::task_ctx()`. + +Once the wrapper decoder is registered, missing source names, unregistered +sources, and source-decoder errors fail directly without passing source bytes +to a plan codec. If the wrapper decoder itself is absent, the existing +execution-plan codec fallback applies to the wrapper payload. Register sources +on every reader before writers start emitting the new wrapper payload. +Built-in scans and existing codec-encoded sources keep their wire formats. + ### Session time zone now applies to mixed timestamp operations When `datafusion.execution.time_zone` is set, DataFusion now uses it to @@ -1195,7 +1224,8 @@ let mut config = SessionConfig::new(); config.update_extension(MyExec::register_proto)?; ``` -Only extension plans implement `ExtensionPlanFromProto`, and registration is the trait's provided `register_proto` method. A built-in plan has a wire variant of its own and keeps its inherent `try_from_proto`, so `FilterExec::register_proto(&mut registry)` does not compile. +Registration uses the `ExtensionPlanFromProto::register_proto` method, which built-in plans with dedicated wire variants generally do not implement. For example, `FilterExec::register_proto(&mut registry)` does not compile. +`DataSourceExec` is an exception: it implements the trait to decode custom sources through the shared extension-plan wrapper, while built-in scans keep their dedicated wire variants. `ProtoDecoderRegistry` is shared by every kind of extension DataFusion can decode, not only plans. Its key is the trait plus the name, so a plan and an expression may use the same name, and one session carries one registry rather than one per kind.