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..d3e59835d1363 100644 --- a/datafusion/datasource/src/proto.rs +++ b/datafusion/datasource/src/proto.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +//! Serialization support for data sources and sinks. +//! +//! Custom sources use the name-keyed decoder API in [`registry`](crate::proto::registry). +//! Custom sinks use the equivalent API in [`sink_registry`](crate::proto::sink_registry). +//! //! Protobuf conversions for the file-scan leaf types owned by this crate: //! [`FileRange`], [`PartitionedFile`] and [`FileGroup`]. //! @@ -30,6 +35,19 @@ //! 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 mod sink_registry; + +pub use datafusion_proto_models::ProtoDecoderRegistry; +pub use registry::{ + ExtensionDataSourceFromProto, data_source_names, extension_data_source_node, + register_data_source, +}; +pub use sink_registry::{ + ExtensionDataSinkFromProto, data_sink_names, extension_data_sink_node, + register_data_sink, +}; + 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/proto/sink_registry.rs b/datafusion/datasource/src/proto/sink_registry.rs new file mode 100644 index 0000000000000..5f3fd5d3ae4c7 --- /dev/null +++ b/datafusion/datasource/src/proto/sink_registry.rs @@ -0,0 +1,355 @@ +// 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 [`DataSink`]s. +//! +//! Custom sinks serialize through the existing extension-plan wire format: +//! [`DataSinkExec`] is the registered plan type, and its payload identifies +//! the sink decoder. Built-in sinks 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::{ + PhysicalExtensionDataSinkNode, PhysicalPlanNode, +}; +use prost::Message; + +use crate::sink::{DataSink, DataSinkExec}; + +/// The wire name of an extension [`DataSink`] and the constructor that +/// rebuilds it. Only custom sinks implement this trait. +/// +/// Implement [`DataSink::try_to_proto`] using [`extension_data_sink_node`]. +/// The helper writes a named `DataSinkExec` plan containing the sink's +/// [`Self::NAME`] and payload. The sink 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 ExtensionDataSinkFromProto: DataSink + Sized { + /// Namespace this name with the owning crate, e.g. `"my_crate.MySink"`. + /// Names are scoped to `dyn DataSink`, independently of plan names. + const NAME: &'static str; + + /// Rebuild the sink from its typed extension payload. + /// [`DataSinkExec`] supplies the execution-plan wrapper. + fn try_from_proto( + node: &PhysicalExtensionDataSinkNode, + ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result>; +} + +/// Build an extension plan for sink `T`, stamping the wrapper's plan name +/// and the sink's registered name. +/// +/// This uses [`ExecutionPlanEncodeCtx::extension_node`] for [`DataSinkExec`], +/// whose payload contains `T::NAME`, the sink bytes, and the optional sort order. +/// The input plan is encoded as the wrapper's only child. The helper lives in +/// this crate because the plan encode context cannot depend on [`DataSink`]. +/// +/// A sink implementing [`ExtensionDataSinkFromProto`] can use this in its +/// [`DataSink::try_to_proto`] hook: +/// +/// ```ignore +/// fn try_to_proto( +/// &self, +/// exec: &DataSinkExec, +/// ctx: &ExecutionPlanEncodeCtx<'_>, +/// ) -> Result> { +/// let payload = self.partition_id.to_le_bytes().to_vec(); +/// Ok(Some(extension_data_sink_node::(ctx, exec, payload)?)) +/// } +/// ``` +pub fn extension_data_sink_node( + ctx: &ExecutionPlanEncodeCtx<'_>, + exec: &DataSinkExec, + payload: Vec, +) -> Result { + let sink = PhysicalExtensionDataSinkNode { + node: payload, + sink_name: Some(T::NAME.to_string()), + sort_order: exec.encode_sort_order(ctx)?, + }; + ctx.extension_node::(sink.encode_to_vec(), vec![exec.input()]) +} + +// Private so the decoder representation can evolve without changing the API. +type DataSinkDecoder = fn( + &PhysicalExtensionDataSinkNode, + &ExecutionPlanDecodeCtx<'_>, +) -> Result>; + +/// Register `T` under its [`ExtensionDataSinkFromProto::NAME`], together +/// with the shared [`DataSinkExec`] 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. Sink names are scoped to `dyn DataSink`. +/// 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. +pub fn register_data_sink( + registry: &mut ProtoDecoderRegistry, +) -> Result<()> { + // The shared wrapper may remain registered if sink registration fails. + // This is safe: it still requires a registered sink decoder and rejects + // unknown sinks rather than decoding their payloads through a plan codec. + DataSinkExec::register_proto(registry)?; + registry + .register_decoder::(T::NAME, T::try_from_proto) +} + +/// Every registered extension sink name, in arbitrary order. +/// Names registered for other extension kinds are not included. +pub fn data_sink_names(registry: &ProtoDecoderRegistry) -> impl Iterator { + registry.names::() +} + +/// Decode the wrapper through the same registry dispatch as other extension +/// plans. Only the sink-specific lookup belongs to this crate. +impl ExtensionPlanFromProto for DataSinkExec { + const NAME: &'static str = "datafusion.DataSinkExec"; + + fn try_from_proto( + node: &PhysicalPlanNode, + ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + let extension = + expect_plan_variant!(node, PhysicalPlanType::Extension, "Extension"); + let [input] = extension.inputs.as_slice() else { + return plan_err!("DataSinkExec expects exactly one input plan"); + }; + let payload = extension.node.as_slice(); + let sink = PhysicalExtensionDataSinkNode::decode(payload).map_err(|e| { + plan_datafusion_err!("Failed to decode DataSinkExec payload: {e}") + })?; + let Some(name) = sink.sink_name.as_deref().filter(|name| !name.is_empty()) else { + return plan_err!("DataSinkExec payload is missing a non-empty sink name"); + }; + let registry = ctx + .task_ctx() + .session_config() + .get_extension::(); + if let Some(decoder) = registry + .as_ref() + .and_then(|registry| registry.decoder::(name)) + { + let input = ctx.decode_child(input)?; + let sort_order = Self::decode_sort_order( + sink.sort_order.as_ref(), + ctx, + input.schema().as_ref(), + )?; + let sink = decoder(&sink, ctx)?; + return Ok(Arc::new(Self::new(input, sink, sort_order))); + } + let mut names: Vec<&str> = registry + .as_ref() + .map(|registry| data_sink_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 DataSink '{name}'. Register \ + the sink in the ProtoDecoderRegistry attached to the decoding session's \ + SessionConfig. Registered extension data sinks: {names}" + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::SchemaRef; + use async_trait::async_trait; + use datafusion_common::internal_err; + use datafusion_execution::TaskContext; + use datafusion_physical_plan::proto::execution_plan_names; + use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, SendableRecordBatchStream, + }; + use std::fmt; + + // These fixtures only exercise registration; their decoders are never invoked. + macro_rules! registration_sink { + ($name:ident) => { + #[derive(Debug)] + struct $name; + + impl DisplayAs for $name { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut fmt::Formatter, + ) -> fmt::Result { + write!(f, stringify!($name)) + } + } + + #[async_trait] + impl DataSink for $name { + fn schema(&self) -> &SchemaRef { + unreachable!("registration-only fixture") + } + async fn write_all( + &self, + _data: SendableRecordBatchStream, + _context: &Arc, + ) -> Result { + unreachable!("registration-only fixture") + } + } + }; + } + + registration_sink!(RegisteredSink); + impl ExtensionDataSinkFromProto for RegisteredSink { + const NAME: &'static str = "test.RegisteredSink"; + + fn try_from_proto( + _node: &PhysicalExtensionDataSinkNode, + _ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + internal_err!("registration-only fixture") + } + } + + registration_sink!(CollidingSink); + impl ExtensionDataSinkFromProto for CollidingSink { + const NAME: &'static str = "test.RegisteredSink"; + + fn try_from_proto( + _node: &PhysicalExtensionDataSinkNode, + _ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + internal_err!("registration-only fixture") + } + } + + registration_sink!(OtherSink); + impl ExtensionDataSinkFromProto for OtherSink { + const NAME: &'static str = "other.OtherSink"; + + fn try_from_proto( + _node: &PhysicalExtensionDataSinkNode, + _ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + internal_err!("registration-only fixture") + } + } + + registration_sink!(UnnamedSink); + impl ExtensionDataSinkFromProto for UnnamedSink { + const NAME: &'static str = ""; + + fn try_from_proto( + _node: &PhysicalExtensionDataSinkNode, + _ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + internal_err!("registration-only fixture") + } + } + + #[test] + fn register_is_idempotent_and_registers_one_shared_wrapper() -> Result<()> { + let mut registry = ProtoDecoderRegistry::new(); + register_data_sink::(&mut registry)?; + register_data_sink::(&mut registry)?; + assert_eq!( + data_sink_names(®istry).collect::>(), + [RegisteredSink::NAME] + ); + register_data_sink::(&mut registry)?; + let mut names: Vec<_> = data_sink_names(®istry).collect(); + names.sort_unstable(); + assert_eq!(names, [OtherSink::NAME, RegisteredSink::NAME]); + assert_eq!( + execution_plan_names(®istry).collect::>(), + [DataSinkExec::NAME] + ); + Ok(()) + } + + #[test] + fn register_rejects_a_name_collision_between_types() -> Result<()> { + let mut registry = ProtoDecoderRegistry::new(); + register_data_sink::(&mut registry)?; + let err = register_data_sink::(&mut registry).unwrap_err(); + assert!(err.to_string().contains("already registered"), "{err}"); + assert!(err.to_string().contains(RegisteredSink::NAME), "{err}"); + assert_eq!( + data_sink_names(®istry).collect::>(), + [RegisteredSink::NAME] + ); + Ok(()) + } + + #[test] + fn register_rejects_an_empty_name_and_may_retain_the_wrapper() { + let mut registry = ProtoDecoderRegistry::new(); + let err = register_data_sink::(&mut registry).unwrap_err(); + assert!(err.to_string().contains("empty name"), "{err}"); + assert!(data_sink_names(®istry).next().is_none()); + assert_eq!( + execution_plan_names(®istry).collect::>(), + [DataSinkExec::NAME] + ); + } + + #[test] + fn sink_and_source_names_are_independent() -> Result<()> { + use crate::source::DataSource; + let mut registry = ProtoDecoderRegistry::new(); + registry.register_decoder::(RegisteredSink::NAME, 7)?; + assert!(data_sink_names(®istry).next().is_none()); + register_data_sink::(&mut registry)?; + assert_eq!( + data_sink_names(®istry).collect::>(), + [RegisteredSink::NAME] + ); + assert_eq!( + registry.decoder::(RegisteredSink::NAME), + Some(&7) + ); + Ok(()) + } + + #[test] + fn wrapper_collision_does_not_register_the_sink() -> Result<()> { + let mut registry = ProtoDecoderRegistry::new(); + registry.register_decoder::(DataSinkExec::NAME, 7)?; + let err = register_data_sink::(&mut registry).unwrap_err(); + assert!(err.to_string().contains("already registered"), "{err}"); + assert!(data_sink_names(®istry).next().is_none()); + assert_eq!( + registry.decoder::(DataSinkExec::NAME), + Some(&7) + ); + Ok(()) + } +} diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 6851c9d7f48cb..3d5d42dcc051a 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -77,6 +77,10 @@ pub trait DataSink: Any + DisplayAs + Debug + Send + Sync { /// /// Implementations can use `ctx` to encode the input plan, sink-specific /// expressions, and [`DataSinkExec::encode_sort_order`]. + /// Custom sinks can use + /// [`extension_data_sink_node`](crate::proto::extension_data_sink_node) to + /// encode a named wrapper and register their decoder with + /// [`register_data_sink`](crate::proto::register_data_sink). /// /// Returning `Ok(None)` lets the caller try its extension codec instead. #[cfg(feature = "proto")] 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..00974b89918bb 100644 --- a/datafusion/physical-plan/src/proto/registry.rs +++ b/datafusion/physical-plan/src/proto/registry.rs @@ -26,8 +26,11 @@ //! 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` and `DataSinkExec` are exceptions: +//! their registered decoders wrap custom sources and sinks, while built-in scans +//! and sinks 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 +50,10 @@ 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 do the shared `DataSourceExec` and +/// `DataSinkExec` wrappers for custom sources and sinks. 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 +116,10 @@ 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` and `DataSinkExec` implement it to wrap custom + /// sources and sinks; 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..625af19c324de 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1058,6 +1058,25 @@ 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; +} + +// Payload of a DataSinkExec encoded as a PhysicalExtensionNode. The outer +// plan carries the input; sink names use a separate registry kind. +message PhysicalExtensionDataSinkNode { + bytes node = 1; + // Namespaced sink type name. Once the DataSinkExec wrapper decoder is + // registered, sink lookup and decoding errors do not fall back to a plan codec. + optional string sink_name = 2; + PhysicalSortExprNodeCollection sort_order = 3; +} + // 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..394b897c6bfa2 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -20282,6 +20282,250 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { deserializer.deserialize_struct("datafusion.PhysicalExprNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for PhysicalExtensionDataSinkNode { + #[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.sink_name.is_some() { + len += 1; + } + if self.sort_order.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalExtensionDataSinkNode", 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.sink_name.as_ref() { + struct_ser.serialize_field("sinkName", v)?; + } + if let Some(v) = self.sort_order.as_ref() { + struct_ser.serialize_field("sortOrder", v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalExtensionDataSinkNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "node", + "sink_name", + "sinkName", + "sort_order", + "sortOrder", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Node, + SinkName, + SortOrder, + } + 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), + "sinkName" | "sink_name" => Ok(GeneratedField::SinkName), + "sortOrder" | "sort_order" => Ok(GeneratedField::SortOrder), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalExtensionDataSinkNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalExtensionDataSinkNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut node__ = None; + let mut sink_name__ = None; + let mut sort_order__ = 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::SinkName => { + if sink_name__.is_some() { + return Err(serde::de::Error::duplicate_field("sinkName")); + } + sink_name__ = map_.next_value()?; + } + GeneratedField::SortOrder => { + if sort_order__.is_some() { + return Err(serde::de::Error::duplicate_field("sortOrder")); + } + sort_order__ = map_.next_value()?; + } + } + } + Ok(PhysicalExtensionDataSinkNode { + node: node__.unwrap_or_default(), + sink_name: sink_name__, + sort_order: sort_order__, + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalExtensionDataSinkNode", 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..f0029fcb6767f 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1634,6 +1634,30 @@ 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>, +} +/// Payload of a DataSinkExec encoded as a PhysicalExtensionNode. The outer +/// plan carries the input; sink names use a separate registry kind. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalExtensionDataSinkNode { + #[prost(bytes = "vec", tag = "1")] + pub node: ::prost::alloc::vec::Vec, + /// Namespaced sink type name. Once the DataSinkExec wrapper decoder is + /// registered, sink lookup and decoding errors do not fall back to a plan codec. + #[prost(string, optional, tag = "2")] + pub sink_name: ::core::option::Option<::prost::alloc::string::String>, + #[prost(message, optional, tag = "3")] + pub sort_order: ::core::option::Option, +} /// physical expressions #[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalExprNode { diff --git a/datafusion/proto/tests/cases/plans/data_sink_registry.rs b/datafusion/proto/tests/cases/plans/data_sink_registry.rs new file mode 100644 index 0000000000000..c1ec5f8a493a9 --- /dev/null +++ b/datafusion/proto/tests/cases/plans/data_sink_registry.rs @@ -0,0 +1,531 @@ +// 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 sinks, including execution after decoding. + +use std::fmt::{self, Formatter}; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use datafusion::arrow::array::{ArrayRef, Int64Array, UInt64Array}; +use datafusion::arrow::compute::SortOptions; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{LexRequirement, PhysicalSortRequirement}; +use datafusion::physical_plan::proto::{ + ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx, ExtensionPlanFromProto, +}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, SendableRecordBatchStream, collect, +}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::{DataFusionError, Result, internal_err}; +use datafusion_datasource::memory::MemorySourceConfig; +use datafusion_datasource::proto::{ + ExtensionDataSinkFromProto, ProtoDecoderRegistry, extension_data_sink_node, + register_data_sink, +}; +use datafusion_datasource::sink::{DataSink, DataSinkExec}; +use datafusion_datasource::source::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::{ + PhysicalExtensionDataSinkNode, PhysicalExtensionNode, PhysicalPlanNode, +}; +use prost::Message; + +#[derive(Debug, Default)] +struct WrittenValues(Mutex>); + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])) +} + +// Only the DataSink boilerplate is shared. Each decoder remains explicit. +macro_rules! test_sink { + ($name:ident, $($hook:item)*) => { + #[derive(Debug)] + struct $name { + value: i64, + schema: SchemaRef, + written: Arc, + } + + impl $name { + fn new(value: i64, written: Arc) -> Self { + Self { value, schema: schema(), written } + } + } + + impl DisplayAs for $name { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "{}({})", stringify!($name), self.value) + } + } + + #[async_trait] + impl DataSink for $name { + fn schema(&self) -> &SchemaRef { &self.schema } + + async fn write_all( + &self, + data: SendableRecordBatchStream, + _context: &Arc, + ) -> Result { + let mut rows = 0; + for batch in datafusion::physical_plan::common::collect(data).await? { + let values = batch.column(0).as_any().downcast_ref::().unwrap(); + self.written.0.lock().unwrap().extend(values.values().iter().map(|value| value + self.value)); + rows += batch.num_rows() as u64; + } + Ok(rows) + } + $($hook)* + } + }; +} + +test_sink!( + RegisteredSink, + fn try_to_proto( + &self, + exec: &DataSinkExec, + ctx: &ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(Some(extension_data_sink_node::( + ctx, + exec, + self.value.to_le_bytes().to_vec(), + )?)) + } +); + +impl ExtensionDataSinkFromProto for RegisteredSink { + const NAME: &'static str = "test.RegisteredSink"; + + fn try_from_proto( + node: &PhysicalExtensionDataSinkNode, + ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + let written = ctx + .task_ctx() + .session_config() + .get_extension::() + .unwrap_or_default(); + Ok(Arc::new(Self::new(read_value(&node.node)?, written))) + } +} + +test_sink!( + OtherSink, + fn try_to_proto( + &self, + exec: &DataSinkExec, + ctx: &ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(Some(extension_data_sink_node::( + ctx, + exec, + self.value.to_le_bytes().to_vec(), + )?)) + } +); + +impl ExtensionDataSinkFromProto for OtherSink { + // Deliberately share the wrapper's plan name: kinds must remain independent. + const NAME: &'static str = "datafusion.DataSinkExec"; + + fn try_from_proto( + node: &PhysicalExtensionDataSinkNode, + _ctx: &ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + Ok(Arc::new(Self::new(read_value(&node.node)?, Arc::default()))) + } +} + +test_sink!(LegacySink,); + +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 input() -> Result> { + let values: ArrayRef = Arc::new(Int64Array::from(vec![3, 5])); + let batch = RecordBatch::try_new(schema(), vec![values])?; + Ok(DataSourceExec::from_data_source( + MemorySourceConfig::try_new(&[vec![batch]], schema(), None)?, + )) +} + +fn plan( + value: i64, + sort_order: Option, +) -> Result> { + Ok(Arc::new(DataSinkExec::new( + input()?, + Arc::new(RegisteredSink::new(value, Arc::default())), + sort_order, + ))) +} + +fn session( + register: impl FnOnce(&mut ProtoDecoderRegistry) -> Result<()>, +) -> Result { + let mut config = SessionConfig::new(); + config.update_extension(register)?; + Ok(SessionContext::new_with_config(config)) +} + +fn encode(plan: Arc) -> Result { + PhysicalPlanNode::try_from_physical_plan(plan, &DefaultPhysicalExtensionCodec {}) +} + +fn decode( + node: &PhysicalPlanNode, + ctx: &SessionContext, + codec: &dyn PhysicalExtensionCodec, +) -> Result> { + let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice()).unwrap(); + node.try_into_physical_plan(ctx.task_ctx().as_ref(), codec) +} + +fn extension(node: &PhysicalPlanNode) -> &PhysicalExtensionNode { + let Some(PhysicalPlanType::Extension(extension)) = &node.physical_plan_type else { + panic!("expected extension plan"); + }; + extension +} + +fn extension_mut(node: &mut PhysicalPlanNode) -> &mut PhysicalExtensionNode { + let Some(PhysicalPlanType::Extension(extension)) = &mut node.physical_plan_type + else { + panic!("expected extension plan"); + }; + extension +} + +fn update_payload( + node: &mut PhysicalPlanNode, + update: impl FnOnce(&mut PhysicalExtensionDataSinkNode), +) { + let mut sink = + PhysicalExtensionDataSinkNode::decode(extension(node).node.as_slice()).unwrap(); + update(&mut sink); + extension_mut(node).node = sink.encode_to_vec(); +} + +fn decoded_sink(plan: &Arc) -> &T { + plan.downcast_ref::() + .unwrap() + .sink() + .downcast_ref::() + .unwrap() +} + +#[tokio::test] +async fn sink_registry_roundtrip_executes_with_receiving_session_state() -> Result<()> { + let written = Arc::new(WrittenValues::default()); + let mut config = SessionConfig::new().with_extension(Arc::clone(&written)); + config.update_extension(register_data_sink::)?; + let reader = SessionContext::new_with_config(config); + let codec = DefaultPhysicalExtensionCodec {}; + let original = plan(7, None)?; + let node = encode(original.clone())?; + assert_eq!( + extension(&node).plan_name.as_deref(), + Some(DataSinkExec::NAME) + ); + assert_eq!(extension(&node).inputs.len(), 1); + let payload = + PhysicalExtensionDataSinkNode::decode(extension(&node).node.as_slice()).unwrap(); + assert_eq!(payload.sink_name.as_deref(), Some(RegisteredSink::NAME)); + assert_eq!(payload.node, 7_i64.to_le_bytes()); + assert!(payload.sort_order.is_none()); + let bytes = physical_plan_to_bytes_with_extension_codec(original, &codec)?; + let decoded = physical_plan_from_bytes_with_extension_codec( + &bytes, + &reader.task_ctx(), + &codec, + )?; + assert_eq!(decoded_sink::(&decoded).value, 7); + let batches = collect(decoded, reader.task_ctx()).await?; + assert_eq!( + batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 2 + ); + assert_eq!(*written.0.lock().unwrap(), [10, 12]); + Ok(()) +} + +#[test] +fn sink_input_and_sort_requirements_roundtrip() -> Result<()> { + let order: LexRequirement = [PhysicalSortRequirement::new( + Arc::new(Column::new("a", 0)), + Some(SortOptions { + descending: true, + nulls_first: false, + }), + )] + .into(); + let original = plan(7, Some(order.clone()))?; + let reader = session(register_data_sink::)?; + let decoded = decode( + &encode(original.clone())?, + &reader, + &DefaultPhysicalExtensionCodec {}, + )?; + let decoded = decoded.downcast_ref::().unwrap(); + let original = original.downcast_ref::().unwrap(); + assert_eq!(decoded.sort_order(), &Some(order)); + assert_eq!(decoded.input().schema(), original.input().schema()); + assert_eq!(decoded.schema(), original.schema()); + Ok(()) +} + +#[derive(Debug)] +struct GreedyCodec; + +impl PhysicalExtensionCodec for GreedyCodec { + fn try_decode( + &self, + _buf: &[u8], + inputs: &[Arc], + _ctx: &TaskContext, + _converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + Ok(Arc::new(DataSinkExec::new( + Arc::clone(&inputs[0]), + Arc::new(LegacySink::new(0, Arc::default())), + None, + ))) + } + fn try_encode( + &self, + _node: Arc, + _buf: &mut Vec, + _converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + internal_err!("unexpected codec encode") + } +} + +#[test] +fn registered_sink_does_not_call_plan_codec() -> Result<()> { + let reader = session(register_data_sink::)?; + let decoded = decode(&encode(plan(9, None)?)?, &reader, &GreedyCodec)?; + assert_eq!(decoded_sink::(&decoded).value, 9); + Ok(()) +} + +#[test] +fn malformed_sink_and_wrapper_payloads_do_not_fall_back_to_codec() -> Result<()> { + let reader = session(register_data_sink::)?; + let mut node = encode(plan(9, None)?)?; + update_payload(&mut node, |sink| sink.node = vec![0]); + let err = decode(&node, &reader, &GreedyCodec).unwrap_err(); + assert!(matches!(err.find_root(), DataFusionError::Internal(_))); + assert!(err.to_string().contains("expected eight bytes"), "{err}"); + extension_mut(&mut node).node = vec![0]; + let err = decode(&node, &reader, &GreedyCodec).unwrap_err(); + assert!( + err.to_string() + .contains("Failed to decode DataSinkExec payload"), + "{err}" + ); + Ok(()) +} + +#[test] +fn sink_wrapper_requires_exactly_one_input() -> Result<()> { + let reader = session(register_data_sink::)?; + for count in [0, 2] { + let mut node = encode(plan(9, None)?)?; + extension_mut(&mut node).inputs = vec![encode(input()?)?; count]; + let err = decode(&node, &reader, &GreedyCodec).unwrap_err(); + assert!( + err.to_string().contains("expects exactly one input plan"), + "{err}" + ); + } + Ok(()) +} + +#[test] +fn invalid_input_and_ordering_do_not_fall_back_to_codec() -> Result<()> { + let reader = session(register_data_sink::)?; + let mut node = encode(plan(9, None)?)?; + extension_mut(&mut node).inputs[0] = PhysicalPlanNode::default(); + assert!(decode(&node, &reader, &GreedyCodec).is_err()); + + let mut node = encode(plan(9, None)?)?; + update_payload(&mut node, |sink| { + sink.sort_order = + Some(datafusion_proto::protobuf::PhysicalSortExprNodeCollection { + physical_sort_expr_nodes: vec![ + datafusion_proto::protobuf::PhysicalSortExprNode::default(), + ], + }); + }); + let err = decode(&node, &reader, &GreedyCodec).unwrap_err(); + assert!( + err.to_string() + .contains("Unexpected empty physical expression"), + "{err}" + ); + Ok(()) +} + +#[test] +fn unnamed_sinks_do_not_fall_back_to_plan_codec() -> Result<()> { + let reader = session(register_data_sink::)?; + for name in [None, Some(String::new())] { + let mut node = encode(plan(9, None)?)?; + update_payload(&mut node, |sink| sink.sink_name = name); + let err = decode(&node, &reader, &GreedyCodec).unwrap_err(); + assert!( + err.to_string().contains("missing a non-empty sink name"), + "{err}" + ); + } + Ok(()) +} + +#[test] +fn unregistered_sinks_report_registered_names_without_codec_fallback() -> Result<()> { + for (reader, expected) in [ + (session(DataSinkExec::register_proto)?, "none"), + (session(register_data_sink::)?, OtherSink::NAME), + ] { + let err = decode(&encode(plan(9, None)?)?, &reader, &GreedyCodec).unwrap_err(); + assert!(matches!(err.find_root(), DataFusionError::Plan(_))); + let message = err.to_string(); + assert!(message.contains(RegisteredSink::NAME), "{message}"); + assert!( + message.contains(&format!("Registered extension data sinks: {expected}")), + "{message}" + ); + } + Ok(()) +} + +#[test] +fn unregistered_wrapper_uses_existing_plan_codec_fallback() -> Result<()> { + let reader = SessionContext::new(); + let node = encode(plan(9, None)?)?; + let decoded = decode(&node, &reader, &GreedyCodec)?; + assert_eq!(decoded_sink::(&decoded).value, 0); + let err = decode(&node, &reader, &DefaultPhysicalExtensionCodec {}).unwrap_err(); + assert!(err.to_string().contains(DataSinkExec::NAME), "{err}"); + Ok(()) +} + +#[test] +fn independent_sink_types_and_matching_plan_names_coexist() -> Result<()> { + let reader = session(|registry| { + register_data_sink::(registry)?; + register_data_sink::(registry) + })?; + let decoded = decode( + &encode(plan(11, None)?)?, + &reader, + &DefaultPhysicalExtensionCodec {}, + )?; + assert_eq!(decoded_sink::(&decoded).value, 11); + let other = Arc::new(DataSinkExec::new( + input()?, + Arc::new(OtherSink::new(22, Arc::default())), + None, + )); + let decoded = decode(&encode(other)?, &reader, &DefaultPhysicalExtensionCodec {})?; + assert_eq!(decoded_sink::(&decoded).value, 22); + Ok(()) +} + +#[derive(Debug)] +struct LegacyCodec; + +impl PhysicalExtensionCodec for LegacyCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + _ctx: &TaskContext, + _converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + Ok(Arc::new(DataSinkExec::new( + Arc::clone(&inputs[0]), + Arc::new(LegacySink::new(read_value(buf)?, Arc::default())), + None, + ))) + } + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + _converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + let sink = decoded_sink::(&node); + buf.extend_from_slice(&sink.value.to_le_bytes()); + Ok(()) + } +} + +#[test] +fn codec_owned_sinks_keep_the_legacy_wire_format() -> Result<()> { + let reader = session(register_data_sink::)?; + let original = Arc::new(DataSinkExec::new( + input()?, + Arc::new(LegacySink::new(33, Arc::default())), + None, + )); + let node = PhysicalPlanNode::try_from_physical_plan(original, &LegacyCodec)?; + assert!(extension(&node).plan_name.is_none()); + assert_eq!(extension(&node).node, 33_i64.to_le_bytes()); + let decoded = decode(&node, &reader, &LegacyCodec)?; + assert_eq!(decoded_sink::(&decoded).value, 33); + Ok(()) +} + +#[cfg(feature = "json")] +#[test] +fn sink_registry_json_roundtrip() -> Result<()> { + let reader = session(register_data_sink::)?; + let node = encode(plan(44, None)?)?; + let json = serde_json::to_string(&node).unwrap(); + assert!(json.contains(DataSinkExec::NAME)); + let node = serde_json::from_str(&json).unwrap(); + let decoded = decode(&node, &reader, &DefaultPhysicalExtensionCodec {})?; + assert_eq!(decoded_sink::(&decoded).value, 44); + Ok(()) +} 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..73eacad5586d4 100644 --- a/datafusion/proto/tests/cases/plans/mod.rs +++ b/datafusion/proto/tests/cases/plans/mod.rs @@ -34,6 +34,8 @@ use datafusion_proto::physical_plan::{ use std::sync::Arc; mod aggregates; +mod data_sink_registry; +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..3801e3cc19dcd 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,62 @@ 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. + +### Custom `DataSink` decoder registration + +Custom sinks can implement `datafusion_datasource::proto::ExtensionDataSinkFromProto` +and use `extension_data_sink_node::(ctx, exec, payload)?` in +`DataSink::try_to_proto`. As with sources, register each sink on every decoding +session: + +```rust,ignore +let mut config = SessionConfig::new(); +config.update_extension(|registry: &mut ProtoDecoderRegistry| { + register_data_sink::(registry) +})?; +``` + +Registration adds both the sink decoder and the shared `DataSinkExec` plan decoder. +The existing extension-plan wire format carries one input plan and a payload +containing the sink name, sink bytes, and optional sort requirements. Constructors +receive `PhysicalExtensionDataSinkNode` and `ExecutionPlanDecodeCtx`; DataFusion +reconstructs the input and ordering and supplies the `DataSinkExec` wrapper. +Sink names are independent of source, plan, and expression names. + +Once the wrapper decoder is registered, missing sink names, unregistered sinks, +and decoding errors fail directly without falling back to a plan codec. If the +wrapper itself is unregistered, the existing execution-plan codec fallback applies +to its payload. Upgrade readers and register sinks before writers emit the new +wrapper payload. Built-in and existing codec-encoded sinks 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 +1251,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` and `DataSinkExec` are exceptions: they implement the trait to decode custom sources and sinks through shared extension-plan wrappers, while built-in scans and sinks 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.