refactor(proto): migrate UnnestExec serde

datafusion/physical-plan/src/unnest.rs

@@ -283,6 +283,127 @@ impl ExecutionPlan for UnnestExec {
     fn metrics(&self) -> Option<MetricsSet> {
         Some(self.metrics.clone_inner())
     }
+
+    #[cfg(feature = "proto")]
+    fn try_to_proto(
+        &self,
+        ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
+    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+        use datafusion_proto_models::protobuf;
+
+        let input = ctx.encode_child(self.input())?;
+        let schema = self.schema().as_ref().try_into()?;
+        let list_type_columns = self
+            .list_column_indices()
+            .iter()
+            .map(|column| protobuf::ListUnnest {
+                index_in_input_schema: column.index_in_input_schema as _,
+                depth: column.depth as _,
+            })
+            .collect();
+        let struct_type_columns = self
+            .struct_column_indices()
+            .iter()
+            .map(|index| *index as _)
+            .collect();
+        let options = protobuf::UnnestOptions {
+            preserve_nulls: self.options().preserve_nulls,
+            recursions: self
+                .options()
+                .recursions
+                .iter()
+                .map(|recursion| protobuf::RecursionUnnestOption {
+                    input_column: Some((&recursion.input_column).into()),
+                    output_column: Some((&recursion.output_column).into()),
+                    depth: recursion.depth as _,
+                })
+                .collect(),
+        };
+
+        Ok(Some(protobuf::PhysicalPlanNode {
+            physical_plan_type: Some(
+                protobuf::physical_plan_node::PhysicalPlanType::Unnest(Box::new(
+                    protobuf::UnnestExecNode {
+                        input: Some(Box::new(input)),
+                        schema: Some(schema),
+                        list_type_columns,
+                        struct_type_columns,
+                        options: Some(options),
+                    },
+                )),
+            ),
+        }))
+    }
+}
+
+#[cfg(feature = "proto")]
+impl UnnestExec {
+    /// Reconstruct an [`UnnestExec`] from its protobuf representation.
+    ///
+    /// The exact inverse of [`ExecutionPlan::try_to_proto`].
+    ///
+    /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto
+    pub fn try_from_proto(
+        node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
+        ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        use datafusion_proto_models::protobuf;
+
+        let unnest = crate::expect_plan_variant!(
+            node,
+            protobuf::physical_plan_node::PhysicalPlanType::Unnest,
+            "UnnestExec",
+        );
+        let input =
+            ctx.decode_required_child(unnest.input.as_deref(), "UnnestExec", "input")?;
+        let schema: Schema = unnest
+            .schema
+            .as_ref()
+            .ok_or_else(|| {
+                datafusion_common::internal_datafusion_err!(
+                    "UnnestExec is missing required field 'schema'"
+                )
+            })?
+            .try_into()?;
+        let list_column_indices = unnest
+            .list_type_columns
+            .iter()
+            .map(|column| ListUnnest {
+                index_in_input_schema: column.index_in_input_schema as _,
+                depth: column.depth as _,
+            })
+            .collect();
+        let struct_column_indices = unnest
+            .struct_type_columns
+            .iter()
+            .map(|index| *index as _)
+            .collect();
+        let options = unnest.options.as_ref().ok_or_else(|| {
+            datafusion_common::internal_datafusion_err!(
+                "UnnestExec is missing required field 'options'"
+            )
+        })?;
+        let options = UnnestOptions {
+            preserve_nulls: options.preserve_nulls,
+            recursions: options
+                .recursions
+                .iter()
+                .map(|recursion| datafusion_common::RecursionUnnestOption {
+                    input_column: recursion.input_column.as_ref().unwrap().into(),
+                    output_column: recursion.output_column.as_ref().unwrap().into(),
+                    depth: recursion.depth as _,
+                })
+                .collect(),
+        };
+
+        Ok(Arc::new(UnnestExec::new(
+            input,
+            list_column_indices,
+            struct_column_indices,
+            Arc::new(schema),
+            options,
+        )?))
+    }
 }
 
 #[derive(Clone, Debug)]

datafusion/proto/src/physical_plan/mod.rs

@@ -100,7 +100,7 @@ use datafusion_physical_plan::scalar_subquery::{ScalarSubqueryExec, ScalarSubque
 use datafusion_physical_plan::sorts::sort::SortExec;
 use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion_physical_plan::union::{InterleaveExec, UnionExec};
-use datafusion_physical_plan::unnest::{ListUnnest, UnnestExec};
+use datafusion_physical_plan::unnest::UnnestExec;
 use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowAggExec};
 use datafusion_physical_plan::{ExecutionPlan, InputOrderMode, PhysicalExpr, WindowExpr};
 use prost::Message;
@@ -124,10 +124,7 @@ use crate::physical_plan::to_proto::{
 use crate::protobuf::physical_aggregate_expr_node::AggregateFunction;
 use crate::protobuf::physical_expr_node::ExprType;
 use crate::protobuf::physical_plan_node::PhysicalPlanType;
-use crate::protobuf::{
-    self, ListUnnest as ProtoListUnnest, SortMergeJoinExecNode, proto_error,
-    window_agg_exec_node,
-};
+use crate::protobuf::{self, SortMergeJoinExecNode, proto_error, window_agg_exec_node};
 
 pub mod from_proto;
 pub mod to_proto;
@@ -853,8 +850,8 @@ pub trait PhysicalPlanNodeExt: Sized {
             PhysicalPlanType::ParquetSink(sink) => {
                 self.try_into_parquet_sink_physical_plan(sink, ctx, proto_converter)
             }
-            PhysicalPlanType::Unnest(unnest) => {
-                self.try_into_unnest_physical_plan(unnest, ctx, proto_converter)
+            PhysicalPlanType::Unnest(_) => {
+                UnnestExec::try_from_proto(self.node(), &decode_ctx)
             }
             PhysicalPlanType::Cooperative(_) => {
                 CooperativeExec::try_from_proto(self.node(), &decode_ctx)
@@ -1059,14 +1056,6 @@ pub trait PhysicalPlanNodeExt: Sized {
             return Ok(node);
         }
 
-        if let Some(exec) = plan.downcast_ref::<UnnestExec>() {
-            return protobuf::PhysicalPlanNode::try_from_unnest_exec(
-                exec,
-                codec,
-                proto_converter,
-            );
-        }
-
         if let Some(exec) = plan.downcast_ref::<LazyMemoryExec>()
             && let Some(node) =
                 protobuf::PhysicalPlanNode::try_from_lazy_memory_exec(exec)?
@@ -2504,32 +2493,25 @@ pub trait PhysicalPlanNodeExt: Sized {
         panic!("Trying to use ParquetSink without `parquet` feature enabled");
     }
 
+    #[deprecated(
+        since = "55.0.0",
+        note = "unused by DataFusion; `UnnestExec` deserializes itself via `UnnestExec::try_from_proto`"
+    )]
     fn try_into_unnest_physical_plan(
         &self,
         unnest: &protobuf::UnnestExecNode,
         ctx: &PhysicalPlanDecodeContext<'_>,
         proto_converter: &dyn PhysicalProtoConverterExtension,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        let input = into_physical_plan(&unnest.input, ctx, proto_converter)?;
-
-        Ok(Arc::new(UnnestExec::new(
-            input,
-            unnest
-                .list_type_columns
-                .iter()
-                .map(|c| ListUnnest {
-                    index_in_input_schema: c.index_in_input_schema as _,
-                    depth: c.depth as _,
-                })
-                .collect(),
-            unnest.struct_type_columns.iter().map(|c| *c as _).collect(),
-            Arc::new(convert_required!(unnest.schema)?),
-            unnest
-                .options
-                .as_ref()
-                .map(datafusion_common::UnnestOptions::from_proto)
-                .ok_or_else(|| proto_error("Missing required field in protobuf"))?,
-        )?))
+        let node = protobuf::PhysicalPlanNode {
+            physical_plan_type: Some(PhysicalPlanType::Unnest(Box::new(unnest.clone()))),
+        };
+        let decoder = ConverterPlanDecoder {
+            ctx,
+            proto_converter,
+        };
+        let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder);
+        UnnestExec::try_from_proto(&node, &decode_ctx)
     }
 
     fn generate_series_name_to_str(name: protobuf::GenerateSeriesName) -> &'static str {
@@ -3892,39 +3874,22 @@ pub trait PhysicalPlanNodeExt: Sized {
         Ok(None)
     }
 
+    #[deprecated(
+        since = "55.0.0",
+        note = "unused by DataFusion; `UnnestExec` serializes itself via `ExecutionPlan::try_to_proto`"
+    )]
     fn try_from_unnest_exec(
         exec: &UnnestExec,
         codec: &dyn PhysicalExtensionCodec,
         proto_converter: &dyn PhysicalProtoConverterExtension,
     ) -> Result<protobuf::PhysicalPlanNode> {
-        let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter(
-            exec.input().to_owned(),
+        let encoder = ConverterPlanEncoder {
             codec,
             proto_converter,
-        )?;
-
-        Ok(protobuf::PhysicalPlanNode {
-            physical_plan_type: Some(PhysicalPlanType::Unnest(Box::new(
-                protobuf::UnnestExecNode {
-                    input: Some(Box::new(input)),
-                    schema: Some(exec.schema().try_into()?),
-                    list_type_columns: exec
-                        .list_column_indices()
-                        .iter()
-                        .map(|c| ProtoListUnnest {
-                            index_in_input_schema: c.index_in_input_schema as _,
-                            depth: c.depth as _,
-                        })
-                        .collect(),
-                    struct_type_columns: exec
-                        .struct_column_indices()
-                        .iter()
-                        .map(|c| *c as _)
-                        .collect(),
-                    options: Some(protobuf::UnnestOptions::from_proto(exec.options())),
-                },
-            ))),
-        })
+        };
+        let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder);
+        exec.try_to_proto(&encode_ctx)?
+            .ok_or_else(|| internal_datafusion_err!("UnnestExec is not serializable"))
     }
 
     #[deprecated(

datafusion/proto/tests/cases/roundtrip_physical_plan.rs

@@ -2251,7 +2251,14 @@ fn roundtrip_unnest() -> Result<()> {
     let output_schema =
         Arc::new(Schema::new(vec![fa, fb0, fc1, fc2, fd0, fe1, fe2, fe3]));
     let input = Arc::new(EmptyExec::new(input_schema));
-    let options = UnnestOptions::default();
+    let options = UnnestOptions {
+        preserve_nulls: false,
+        recursions: vec![datafusion_common::RecursionUnnestOption {
+            input_column: datafusion_common::Column::new_unqualified("b"),
+            output_column: datafusion_common::Column::new_unqualified("b"),
+            depth: 2,
+        }],
+    };
     let unnest = UnnestExec::new(
         input,
         vec![
@@ -2270,9 +2277,17 @@ fn roundtrip_unnest() -> Result<()> {
         ],
         vec![2, 4],
         output_schema,
-        options,
+        options.clone(),
     )?;
-    roundtrip_test(Arc::new(unnest))
+    let ctx = SessionContext::new();
+    let codec = DefaultPhysicalExtensionCodec {};
+    let proto_converter = DefaultPhysicalProtoConverter {};
+    let result =
+        roundtrip_test_and_return(Arc::new(unnest), &ctx, &codec, &proto_converter)?;
+    let result = result.downcast_ref::<UnnestExec>().unwrap();
+    assert_eq!(result.options(), &options);
+
+    Ok(())
 }
 
 #[tokio::test]