diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 3511609e2e9b..4722329ea55a 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -473,6 +473,42 @@ 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 { + 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 +719,42 @@ 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 { + 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 cea334e42aac..09c524b605d3 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -814,11 +814,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) @@ -1001,22 +1001,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, @@ -2057,30 +2041,46 @@ 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, 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( + 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, 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( @@ -3583,48 +3583,42 @@ 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, 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( + 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, 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(