You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ga...@apache.org on 2026/08/31 16:49:21 UTC
(datafusion) 01/01: Add ComposedNamedPhysicalExtensionCodec
This is an automated email from the ASF dual-hosted git repository.
gabotechs pushed a commit to branch gabotechs/composed-named-physical-extension-codec
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 1da2d1dda2819be75cb08d90de4cc682b95c67c5
Author: Gabriel <ga...@datadoghq.com>
AuthorDate: Mon Aug 31 18:48:45 2026 +0200
Add ComposedNamedPhysicalExtensionCodec
---
.../examples/proto/composed_extension_codec.rs | 11 +-
datafusion/proto/src/physical_plan/mod.rs | 162 ++++++++++++++++++++-
2 files changed, 164 insertions(+), 9 deletions(-)
diff --git a/datafusion-examples/examples/proto/composed_extension_codec.rs b/datafusion-examples/examples/proto/composed_extension_codec.rs
index d3bcdc73df..12d498ef18 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 26d1a8ed83..030f5667f1 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<u8>,
}
+/// 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<u8>,
+}
+
pub struct DefaultPhysicalProtoConverter {}
impl PhysicalProtoConverterExtension for DefaultPhysicalProtoConverter {
@@ -1951,6 +1965,10 @@ 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<Arc<dyn PhysicalExtensionCodec>>,
@@ -2058,6 +2076,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<Cow<'static, str>, Arc<dyn PhysicalExtensionCodec>>,
+}
+
+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> {
+ 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<Cow<'static, str>>,
+ codec: Arc<dyn PhysicalExtensionCodec>,
+ ) -> Result<Self> {
+ 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<R>(
+ &self,
+ buf: &[u8],
+ decode: impl FnOnce(&dyn PhysicalExtensionCodec, &[u8]) -> Result<R>,
+ ) -> Result<R> {
+ 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::<Vec<_>>()
+ .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<u8>,
+ mut encode: impl FnMut(&dyn PhysicalExtensionCodec, &mut Vec<u8>) -> 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<dyn ExecutionPlan>],
+ ctx: &TaskContext,
+ proto_converter: &dyn PhysicalProtoConverterExtension,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.decode_protobuf(buf, |codec, data| {
+ codec.try_decode(data, inputs, ctx, proto_converter)
+ })
+ }
+
+ fn try_encode(
+ &self,
+ node: Arc<dyn ExecutionPlan>,
+ buf: &mut Vec<u8>,
+ 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<Arc<ScalarUDF>> {
+ self.decode_protobuf(buf, |codec, data| codec.try_decode_udf(name, data))
+ }
+
+ fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> {
+ self.encode_protobuf(buf, |codec, data| codec.try_encode_udf(node, data))
+ }
+
+ fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result<Arc<AggregateUDF>> {
+ self.decode_protobuf(buf, |codec, data| codec.try_decode_udaf(name, data))
+ }
+
+ fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec<u8>) -> 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.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org