adriangb commented on code in PR #26807: URL: https://github.com/apache/datafusion/pull/26807#discussion_r4393405463
########## 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<Arc<dyn DataSink>>; +} + +/// 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<Option<PhysicalPlanNode>> { +/// let payload = self.partition_id.to_le_bytes().to_vec(); +/// Ok(Some(extension_data_sink_node::<Self>(ctx, exec, payload)?)) +/// } +/// ``` +pub fn extension_data_sink_node<T: ExtensionDataSinkFromProto>( + ctx: &ExecutionPlanEncodeCtx<'_>, + exec: &DataSinkExec, + payload: Vec<u8>, +) -> Result<PhysicalPlanNode> { + let sink = PhysicalExtensionDataSinkNode { + node: payload, + sink_name: Some(T::NAME.to_string()), + sort_order: exec.encode_sort_order(ctx)?, + }; + ctx.extension_node::<DataSinkExec>(sink.encode_to_vec(), vec![exec.input()]) +} Review Comment: The sink receives `&DataSinkExec` only to encode the wrapper's input and sort order. That gives the sink access to its parent. With `DataSink::try_to_proto_sink -> DataSinkNode`, `DataSinkExec` encodes `input` and `sort_order`, and this function can be removed. ########## 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<Arc<dyn DataSink>>; +} Review Comment: Same point as in #26555: pass `payload: &[u8]`. `sort_order` belongs to `DataSinkExec`, not to the sink. ```suggestion fn try_from_proto( payload: &[u8], ctx: &ExecutionPlanDecodeCtx<'_>, ) -> Result<Arc<dyn DataSink>>; ``` ########## 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<Arc<dyn DataSink>>; +} + +/// 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<Option<PhysicalPlanNode>> { +/// let payload = self.partition_id.to_le_bytes().to_vec(); +/// Ok(Some(extension_data_sink_node::<Self>(ctx, exec, payload)?)) +/// } +/// ``` +pub fn extension_data_sink_node<T: ExtensionDataSinkFromProto>( + ctx: &ExecutionPlanEncodeCtx<'_>, + exec: &DataSinkExec, + payload: Vec<u8>, +) -> Result<PhysicalPlanNode> { + let sink = PhysicalExtensionDataSinkNode { + node: payload, + sink_name: Some(T::NAME.to_string()), + sort_order: exec.encode_sort_order(ctx)?, + }; + ctx.extension_node::<DataSinkExec>(sink.encode_to_vec(), vec![exec.input()]) +} + +// Private so the decoder representation can evolve without changing the API. +type DataSinkDecoder = fn( + &PhysicalExtensionDataSinkNode, + &ExecutionPlanDecodeCtx<'_>, +) -> Result<Arc<dyn DataSink>>; + +/// 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<T: ExtensionDataSinkFromProto>( + 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::<dyn DataSink, T, DataSinkDecoder>(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<Item = &str> { Review Comment: Same as `data_source_names`: used outside tests only by the error message at L169. ```suggestion pub(crate) fn data_sink_names(registry: &ProtoDecoderRegistry) -> impl Iterator<Item = &str> { ``` ########## 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<Arc<dyn DataSink>>; +} + +/// 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<Option<PhysicalPlanNode>> { +/// let payload = self.partition_id.to_le_bytes().to_vec(); +/// Ok(Some(extension_data_sink_node::<Self>(ctx, exec, payload)?)) +/// } +/// ``` +pub fn extension_data_sink_node<T: ExtensionDataSinkFromProto>( + ctx: &ExecutionPlanEncodeCtx<'_>, + exec: &DataSinkExec, + payload: Vec<u8>, +) -> Result<PhysicalPlanNode> { + let sink = PhysicalExtensionDataSinkNode { + node: payload, + sink_name: Some(T::NAME.to_string()), + sort_order: exec.encode_sort_order(ctx)?, + }; + ctx.extension_node::<DataSinkExec>(sink.encode_to_vec(), vec![exec.input()]) +} + +// Private so the decoder representation can evolve without changing the API. +type DataSinkDecoder = fn( + &PhysicalExtensionDataSinkNode, + &ExecutionPlanDecodeCtx<'_>, +) -> Result<Arc<dyn DataSink>>; + +/// 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<T: ExtensionDataSinkFromProto>( + 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::<dyn DataSink, T, DataSinkDecoder>(T::NAME, T::try_from_proto) +} Review Comment: Same as #26555: validate first, then register the wrapper. ```suggestion if T::NAME.is_empty() { return plan_err!("extension DataSink name must not be empty"); } registry .register_decoder::<dyn DataSink, T, DataSinkDecoder>(T::NAME, T::try_from_proto)?; DataSinkExec::register_proto(registry) } ``` ########## 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<Arc<dyn DataSink>>; +} + +/// 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<Option<PhysicalPlanNode>> { +/// let payload = self.partition_id.to_le_bytes().to_vec(); +/// Ok(Some(extension_data_sink_node::<Self>(ctx, exec, payload)?)) +/// } +/// ``` +pub fn extension_data_sink_node<T: ExtensionDataSinkFromProto>( + ctx: &ExecutionPlanEncodeCtx<'_>, + exec: &DataSinkExec, + payload: Vec<u8>, +) -> Result<PhysicalPlanNode> { + let sink = PhysicalExtensionDataSinkNode { + node: payload, + sink_name: Some(T::NAME.to_string()), + sort_order: exec.encode_sort_order(ctx)?, + }; + ctx.extension_node::<DataSinkExec>(sink.encode_to_vec(), vec![exec.input()]) +} + +// Private so the decoder representation can evolve without changing the API. +type DataSinkDecoder = fn( + &PhysicalExtensionDataSinkNode, + &ExecutionPlanDecodeCtx<'_>, +) -> Result<Arc<dyn DataSink>>; + +/// 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<T: ExtensionDataSinkFromProto>( + 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::<dyn DataSink, T, DataSinkDecoder>(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<Item = &str> { + registry.names::<dyn DataSink>() +} + +/// 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<Arc<dyn ExecutionPlan>> { + 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::<ProtoDecoderRegistry>(); + if let Some(decoder) = registry + .as_ref() + .and_then(|registry| registry.decoder::<dyn DataSink, DataSinkDecoder>(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}" + ) + } Review Comment: This duplicates `registry.rs` L180–197 except for the noun. Put it in one shared helper, for example `fn unknown_extension_err(kind: &str, name: &str, names: impl Iterator<Item = &str>) -> DataFusionError`. ########## datafusion/datasource/src/proto.rs: ########## @@ -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; Review Comment: Same as #26555: keep these modules private. All items are also re-exported below, so this is a second public path to the same items. ########## 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; Review Comment: `sort_order` belongs to `DataSinkExec`, not to the sink payload. In the precursor (#26824) it moves to `DataSinkExecNode`, and the sink message is `ExtensionDataSinkNode { string name = 1; bytes payload = 2; }`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
