diff --git a/datafusion/physical-plan/src/empty.rs b/datafusion/physical-plan/src/empty.rs index 44a6f444dc4b5..809be6a778aa1 100644 --- a/datafusion/physical-plan/src/empty.rs +++ b/datafusion/physical-plan/src/empty.rs @@ -185,6 +185,53 @@ impl ExecutionPlan for EmptyExec { Ok(Arc::new(stats)) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let schema = self.schema().as_ref().try_into()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Empty( + protobuf::EmptyExecNode { + schema: Some(schema), + partitions: self + .properties() + .output_partitioning() + .partition_count() as u32, + }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl EmptyExec { + /// Reconstruct an [`EmptyExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + _ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let empty = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Empty, + "EmptyExec", + ); + let schema = empty.schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "EmptyExec is missing required field 'schema'" + ) + })?; + let schema = Arc::new(arrow::datatypes::Schema::try_from(schema)?); + // Zero means the field was absent in the previous wire format. + let partitions = empty.partitions.max(1) as usize; + Ok(Arc::new(EmptyExec::new(schema).with_partitions(partitions))) + } } #[cfg(test)] diff --git a/datafusion/physical-plan/src/placeholder_row.rs b/datafusion/physical-plan/src/placeholder_row.rs index 20d267331b2aa..ef71166f97416 100644 --- a/datafusion/physical-plan/src/placeholder_row.rs +++ b/datafusion/physical-plan/src/placeholder_row.rs @@ -186,6 +186,55 @@ impl ExecutionPlan for PlaceholderRowExec { None, ))) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let schema = self.schema().as_ref().try_into()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow( + protobuf::PlaceholderRowExecNode { + schema: Some(schema), + partitions: self + .properties() + .output_partitioning() + .partition_count() as u32, + }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl PlaceholderRowExec { + /// Reconstruct a [`PlaceholderRowExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + _ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let placeholder = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow, + "PlaceholderRowExec", + ); + let schema = placeholder.schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "PlaceholderRowExec is missing required field 'schema'" + ) + })?; + let schema = Arc::new(Schema::try_from(schema)?); + // Zero means the field was absent in the previous wire format. + let partitions = placeholder.partitions.max(1) as usize; + Ok(Arc::new( + PlaceholderRowExec::new(schema).with_partitions(partitions), + )) + } } #[cfg(test)] diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index cea334e42aace..9c604e0e659c5 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -823,11 +823,11 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::CrossJoin(crossjoin) => { self.try_into_cross_join_physical_plan(crossjoin, ctx, proto_converter) } - PhysicalPlanType::Empty(empty) => { - self.try_into_empty_physical_plan(empty, ctx, proto_converter) + PhysicalPlanType::Empty(_) => { + EmptyExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::PlaceholderRow(placeholder) => { - self.try_into_placeholder_row_physical_plan(placeholder, ctx) + PhysicalPlanType::PlaceholderRow(_) => { + PlaceholderRowExec::try_from_proto(self.node(), &decode_ctx) } PhysicalPlanType::Sort(sort) => { self.try_into_sort_physical_plan(sort, ctx, proto_converter) @@ -965,16 +965,6 @@ pub trait PhysicalPlanNodeExt: Sized { ); } - if let Some(empty) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_empty_exec(empty, codec); - } - - if let Some(empty) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_placeholder_row_exec( - empty, codec, - ); - } - if let Some(data_source_exec) = plan.downcast_ref::() && let Some(node) = protobuf::PhysicalPlanNode::try_from_data_source_exec( data_source_exec, @@ -2096,31 +2086,48 @@ pub trait PhysicalPlanNodeExt: Sized { Ok(Arc::new(CrossJoinExec::new(left, right))) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `EmptyExec` deserializes itself via `EmptyExec::try_from_proto`" + )] fn try_into_empty_physical_plan( &self, empty: &protobuf::EmptyExecNode, - _ctx: &PhysicalPlanDecodeContext<'_>, - _proto_converter: &dyn PhysicalProtoConverterExtension, + ctx: &PhysicalPlanDecodeContext<'_>, + proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let schema = Arc::new(convert_required!(empty.schema)?); - // A zero (absent) partition count comes from a plan encoded before the - // field existed, which always meant a single partition. - let partitions = empty.partitions.max(1) as usize; - Ok(Arc::new(EmptyExec::new(schema).with_partitions(partitions))) + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::Empty(empty.clone())), + }; + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + EmptyExec::try_from_proto(&node, &decode_ctx) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `PlaceholderRowExec` deserializes itself via `PlaceholderRowExec::try_from_proto`" + )] fn try_into_placeholder_row_physical_plan( &self, placeholder: &protobuf::PlaceholderRowExecNode, - _ctx: &PhysicalPlanDecodeContext<'_>, + ctx: &PhysicalPlanDecodeContext<'_>, ) -> Result> { - let schema = Arc::new(convert_required!(placeholder.schema)?); - // A zero (absent) partition count comes from a plan encoded before the - // field existed, which always meant a single partition. - let partitions = placeholder.partitions.max(1) as usize; - Ok(Arc::new( - PlaceholderRowExec::new(schema).with_partitions(partitions), - )) + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::PlaceholderRow( + placeholder.clone(), + )), + }; + let proto_converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx, + proto_converter: &proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + PlaceholderRowExec::try_from_proto(&node, &decode_ctx) } fn try_into_sort_physical_plan( @@ -3267,33 +3274,41 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `EmptyExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] fn try_from_empty_exec( empty: &EmptyExec, - _codec: &dyn PhysicalExtensionCodec, + codec: &dyn PhysicalExtensionCodec, ) -> Result { - let schema = empty.schema().as_ref().try_into()?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Empty(protobuf::EmptyExecNode { - schema: Some(schema), - partitions: empty.properties().output_partitioning().partition_count() - as u32, - })), + let proto_converter = DefaultPhysicalProtoConverter {}; + let encoder = ConverterPlanEncoder { + codec, + proto_converter: &proto_converter, + }; + let ctx = ExecutionPlanEncodeCtx::new(&encoder); + empty.try_to_proto(&ctx)?.ok_or_else(|| { + internal_datafusion_err!("EmptyExec::try_to_proto returned None") }) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `PlaceholderRowExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] fn try_from_placeholder_row_exec( - empty: &PlaceholderRowExec, - _codec: &dyn PhysicalExtensionCodec, + placeholder: &PlaceholderRowExec, + codec: &dyn PhysicalExtensionCodec, ) -> Result { - let schema = empty.schema().as_ref().try_into()?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::PlaceholderRow( - protobuf::PlaceholderRowExecNode { - schema: Some(schema), - partitions: empty.properties().output_partitioning().partition_count() - as u32, - }, - )), + let proto_converter = DefaultPhysicalProtoConverter {}; + let encoder = ConverterPlanEncoder { + codec, + proto_converter: &proto_converter, + }; + let ctx = ExecutionPlanEncodeCtx::new(&encoder); + placeholder.try_to_proto(&ctx)?.ok_or_else(|| { + internal_datafusion_err!("PlaceholderRowExec::try_to_proto returned None") }) }