diff --git a/datafusion-examples/examples/proto/composed_extension_codec.rs b/datafusion-examples/examples/proto/composed_extension_codec.rs index d3bcdc73df7f8..12d498ef18db6 100644 --- a/datafusion-examples/examples/proto/composed_extension_codec.rs +++ b/datafusion-examples/examples/proto/composed_extension_codec.rs @@ -43,7 +43,7 @@ use datafusion::physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use datafusion::physical_plan::{DisplayAs, ExecutionPlan}; use datafusion::prelude::SessionContext; use datafusion_proto::physical_plan::{ - AsExecutionPlan, ComposedPhysicalExtensionCodec, PhysicalExtensionCodec, + AsExecutionPlan, ComposedNamedPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension, }; use datafusion_proto::protobuf; @@ -58,12 +58,9 @@ pub fn composed_extension_codec() -> Result<()> { }); let ctx = SessionContext::new(); - // Position in this list is important as it will be used for decoding. - // If new codec is added it should go to last position. - let composed_codec = ComposedPhysicalExtensionCodec::new(vec![ - Arc::new(ParentPhysicalExtensionCodec {}), - Arc::new(ChildPhysicalExtensionCodec {}), - ]); + let composed_codec = ComposedNamedPhysicalExtensionCodec::default() + .with_type_named_codec(ParentPhysicalExtensionCodec {})? + .with_type_named_codec(ChildPhysicalExtensionCodec {})?; // Serialize execution plan to proto let proto: protobuf::PhysicalPlanNode = diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 26d1a8ed83d36..5e95e260227fc 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -16,6 +16,7 @@ // under the License. use std::any::Any; +use std::borrow::Cow; use std::cell::RefCell; use std::collections::HashMap; use std::fmt::Debug; @@ -25,8 +26,8 @@ use arrow::datatypes::{IntervalMonthDayNanoType, Schema, SchemaRef}; use datafusion_catalog::memory::MemorySourceConfig; use datafusion_common::utils::{usize_from_wire, usize_to_wire}; use datafusion_common::{ - DataFusionError, Result, internal_datafusion_err, internal_err, not_impl_err, - plan_err, + DataFusionError, Result, config_err, internal_datafusion_err, internal_err, + not_impl_err, plan_err, }; use datafusion_datasource_arrow::source::ArrowSource; #[cfg(feature = "avro")] @@ -1771,6 +1772,19 @@ struct DataEncoderTuple { pub blob: Vec, } +/// NamedDataEncoderTuple captures the name of the encoder +/// in the codec map that was used to encode the data and actual encoded data +#[derive(Clone, PartialEq, prost::Message)] +struct NamedDataEncoderTuple { + /// The name of the encoder used to encode data + /// (to be used for decoding) + #[prost(string, tag = 1)] + pub name: String, + + #[prost(bytes, tag = 2)] + pub blob: Vec, +} + pub struct DefaultPhysicalProtoConverter {} impl PhysicalProtoConverterExtension for DefaultPhysicalProtoConverter { @@ -1951,11 +1965,16 @@ impl PhysicalProtoConverterExtension for DeduplicatingProtoConverter { /// A PhysicalExtensionCodec that tries one of multiple inner codecs /// until one works +#[deprecated( + since = "56.0.0", + note = "Please use `ComposedNamedPhysicalExtensionCodec`" +)] #[derive(Debug)] pub struct ComposedPhysicalExtensionCodec { codecs: Vec>, } +#[expect(deprecated)] impl ComposedPhysicalExtensionCodec { // Position in this codecs list is important as it will be used for decoding. // If new codec is added it should go to last position. @@ -2017,6 +2036,7 @@ impl ComposedPhysicalExtensionCodec { } } +#[expect(deprecated)] impl PhysicalExtensionCodec for ComposedPhysicalExtensionCodec { fn try_decode( &self, @@ -2058,6 +2078,146 @@ impl PhysicalExtensionCodec for ComposedPhysicalExtensionCodec { } } +/// A PhysicalExtensionCodec that tries one of multiple inner codecs until one works. +/// The name of the codec that successfully encoded an [`ExecutionPlan`] is stored in the +/// encoded payload, and the codec with that exact name will be used for decoding. +#[derive(Default, Debug)] +pub struct ComposedNamedPhysicalExtensionCodec { + codecs: HashMap, Arc>, +} + +impl ComposedNamedPhysicalExtensionCodec { + /// Registers a new [`PhysicalExtensionCodec`] with a name extracted from its Rust type. + /// For example, if the implementation is called `MyPhysicalExtensionCodec`, then the codec + /// will be registered with the "MyPhysicalExtensionCodec" name. + pub fn with_type_named_codec( + self, + codec: impl PhysicalExtensionCodec, + ) -> Result { + self.with_codec(std::any::type_name_of_val(&codec), Arc::new(codec)) + } + + /// Registers a new [`PhysicalExtensionCodec`] with the provided name. + pub fn with_codec( + mut self, + name: impl Into>, + codec: Arc, + ) -> Result { + let name = name.into(); + if self.codecs.contains_key(&name) { + return config_err!( + "Two PhysicalExtensionCodecs with the same name ('{name}') were registered" + ); + } + self.codecs.insert(name, codec); + Ok(self) + } + + fn decode_protobuf( + &self, + buf: &[u8], + decode: impl FnOnce(&dyn PhysicalExtensionCodec, &[u8]) -> Result, + ) -> Result { + let proto = NamedDataEncoderTuple::decode(buf) + .map_err(|e| internal_datafusion_err!("{e}"))?; + + let name = proto.name.as_str(); + + let Some(codec) = self.codecs.get(&Cow::Borrowed(name)) else { + let available_names = self + .codecs + .keys() + .map(|v| v.as_ref()) + .collect::>() + .join(", "); + return internal_err!( + "The message was encoded by a codec with name '{name}', but this codec is not available in the current ComposedNamedPhysicalExtensionCodec. Available codecs are: {available_names}" + ); + }; + + decode(codec.as_ref(), &proto.blob) + } + + fn encode_protobuf( + &self, + buf: &mut Vec, + mut encode: impl FnMut(&dyn PhysicalExtensionCodec, &mut Vec) -> Result<()>, + ) -> Result<()> { + let mut data = vec![]; + let mut last_err = None; + let mut encoder_name = None; + + // find the encoder + for (name, codec) in &self.codecs { + match encode(codec.as_ref(), &mut data) { + Ok(_) => { + encoder_name = Some(name); + break; + } + Err(err) => last_err = Some(err), + } + } + + let encoder_name = encoder_name.ok_or_else(|| { + last_err.unwrap_or_else(|| { + DataFusionError::NotImplemented( + "Empty list of composed named codecs".to_owned(), + ) + }) + })?; + + // encode with encoder position + let proto = NamedDataEncoderTuple { + name: encoder_name.to_string(), + blob: data, + }; + proto + .encode(buf) + .map_err(|e| internal_datafusion_err!("{e}")) + } +} + +impl PhysicalExtensionCodec for ComposedNamedPhysicalExtensionCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &TaskContext, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + self.decode_protobuf(buf, |codec, data| { + codec.try_decode(data, inputs, ctx, proto_converter) + }) + } + + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + self.encode_protobuf(buf, |codec, data| { + codec.try_encode(Arc::clone(&node), data, proto_converter) + }) + } + + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { + self.decode_protobuf(buf, |codec, data| codec.try_decode_udf(name, data)) + } + + fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { + self.encode_protobuf(buf, |codec, data| codec.try_encode_udf(node, data)) + } + + fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { + self.decode_protobuf(buf, |codec, data| codec.try_decode_udaf(name, data)) + } + + fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { + self.encode_protobuf(buf, |codec, data| codec.try_encode_udaf(node, data)) + } +} + /// Adapter backing [`ExecutionPlanEncodeCtx`] for plans migrated to the /// `try_to_proto` hook (#22419). Routes child-plan and child-expr encoding back /// through the central converter so nested plans honor their own hooks.