From 84aee77b0fee8ce460f7853345c841e708830af7 Mon Sep 17 00:00:00 2001 From: buraksenn Date: Fri, 17 Jul 2026 18:02:18 +0000 Subject: [PATCH 1/2] Migrate union exec protobuf serialization --- datafusion/physical-plan/src/union.rs | 86 +++++++++++++++++++++++ datafusion/proto/src/physical_plan/mod.rs | 40 +++++------ 2 files changed, 106 insertions(+), 20 deletions(-) diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 3511609e2e9b1..5963ac8d15dc6 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -473,6 +473,49 @@ impl ExecutionPlan for UnionExec { // on all children (either pushed down or via FilterExec) Ok(propagation) } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let inputs = ctx.encode_children(self.inputs())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Union( + protobuf::UnionExecNode { inputs }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl UnionExec { + /// Reconstruct a [`UnionExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: each child in + /// `node.inputs` is decoded recursively via the [`ExecutionPlanDecodeCtx`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let union = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Union, + "UnionExec", + ); + let inputs = union + .inputs + .iter() + .map(|input| ctx.decode_child(input)) + .collect::>>()?; + UnionExec::try_new(inputs) + } } /// Combines multiple input streams by interleaving them. @@ -683,6 +726,49 @@ impl ExecutionPlan for InterleaveExec { fn benefits_from_input_partitioning(&self) -> Vec { vec![false; self.children().len()] } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let inputs = ctx.encode_children(self.inputs())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Interleave( + protobuf::InterleaveExecNode { inputs }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl InterleaveExec { + /// Reconstruct an [`InterleaveExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: each child in + /// `node.inputs` is decoded recursively via the [`ExecutionPlanDecodeCtx`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let interleave = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Interleave, + "InterleaveExec", + ); + let inputs = interleave + .inputs + .iter() + .map(|input| ctx.decode_child(input)) + .collect::>>()?; + Ok(Arc::new(InterleaveExec::try_new(inputs)?)) + } } /// Returns true if all inputs have the same [`Partitioning::Hash`] or [`Partitioning::Range`] diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index e5f8aa072db71..c0aa0772cb46b 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -808,11 +808,11 @@ pub trait PhysicalPlanNodeExt: Sized { ctx, proto_converter, ), - PhysicalPlanType::Union(union) => { - self.try_into_union_physical_plan(union, ctx, proto_converter) + PhysicalPlanType::Union(_) => { + UnionExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Interleave(interleave) => { - self.try_into_interleave_physical_plan(interleave, ctx, proto_converter) + PhysicalPlanType::Interleave(_) => { + InterleaveExec::try_from_proto(self.node(), &decode_ctx) } PhysicalPlanType::CrossJoin(crossjoin) => { self.try_into_cross_join_physical_plan(crossjoin, ctx, proto_converter) @@ -1028,22 +1028,6 @@ pub trait PhysicalPlanNodeExt: Sized { ); } - if let Some(union) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_union_exec( - union, - codec, - proto_converter, - ); - } - - if let Some(interleave) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_interleave_exec( - interleave, - codec, - proto_converter, - ); - } - if let Some(exec) = plan.downcast_ref::() { return protobuf::PhysicalPlanNode::try_from_sort_preserving_merge_exec( exec, @@ -2115,6 +2099,10 @@ pub trait PhysicalPlanNodeExt: Sized { .map(|e| Arc::new(e) as _) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `UnionExec` deserializes itself via `UnionExec::try_from_proto`" + )] fn try_into_union_physical_plan( &self, union: &protobuf::UnionExecNode, @@ -2128,6 +2116,10 @@ pub trait PhysicalPlanNodeExt: Sized { UnionExec::try_new(inputs) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `InterleaveExec` deserializes itself via `InterleaveExec::try_from_proto`" + )] fn try_into_interleave_physical_plan( &self, interleave: &protobuf::InterleaveExecNode, @@ -3786,6 +3778,10 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `UnionExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] fn try_from_union_exec( union: &UnionExec, codec: &dyn PhysicalExtensionCodec, @@ -3808,6 +3804,10 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `InterleaveExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] fn try_from_interleave_exec( interleave: &InterleaveExec, codec: &dyn PhysicalExtensionCodec, From 7c11901dc773e20ba1a706dbe58a589b6bd72d4b Mon Sep 17 00:00:00 2001 From: buraksenn Date: Tue, 21 Jul 2026 05:49:01 +0000 Subject: [PATCH 2/2] Delegate deprecated union serde shims --- datafusion/proto/src/physical_plan/mod.rs | 74 +++++++++++------------ 1 file changed, 34 insertions(+), 40 deletions(-) diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index c0aa0772cb46b..5e00ce9401856 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -2109,11 +2109,15 @@ pub trait PhysicalPlanNodeExt: Sized { ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let mut inputs: Vec> = vec![]; - for input in &union.inputs { - inputs.push(proto_converter.proto_to_execution_plan(input, ctx)?); - } - UnionExec::try_new(inputs) + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::Union(union.clone())), + }; + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + UnionExec::try_from_proto(&node, &decode_ctx) } #[deprecated( @@ -2126,11 +2130,15 @@ pub trait PhysicalPlanNodeExt: Sized { ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let mut inputs: Vec> = vec![]; - for input in &interleave.inputs { - inputs.push(proto_converter.proto_to_execution_plan(input, ctx)?); - } - Ok(Arc::new(InterleaveExec::try_new(inputs)?)) + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::Interleave(interleave.clone())), + }; + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + InterleaveExec::try_from_proto(&node, &decode_ctx) } fn try_into_cross_join_physical_plan( @@ -3787,21 +3795,14 @@ pub trait PhysicalPlanNodeExt: Sized { codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result { - let mut inputs: Vec = vec![]; - for input in union.inputs() { - inputs.push( - protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - input.to_owned(), - codec, - proto_converter, - )?, - ); - } - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Union(protobuf::UnionExecNode { - inputs, - })), - }) + let encoder = ConverterPlanEncoder { + codec, + proto_converter, + }; + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + union + .try_to_proto(&encode_ctx)? + .ok_or_else(|| internal_datafusion_err!("UnionExec is not serializable")) } #[deprecated( @@ -3813,21 +3814,14 @@ pub trait PhysicalPlanNodeExt: Sized { codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result { - let mut inputs: Vec = vec![]; - for input in interleave.inputs() { - inputs.push( - protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - input.to_owned(), - codec, - proto_converter, - )?, - ); - } - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Interleave( - protobuf::InterleaveExecNode { inputs }, - )), - }) + let encoder = ConverterPlanEncoder { + codec, + proto_converter, + }; + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + interleave + .try_to_proto(&encode_ctx)? + .ok_or_else(|| internal_datafusion_err!("InterleaveExec is not serializable")) } fn try_from_sort_preserving_merge_exec(