This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23739-e784e2cb5c43a58ec5900d4df93449c571a9aef5 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit f4a94277784a2df0027c413d1274efd3d79497ab Author: Phoenix <[email protected]> AuthorDate: Fri Jul 24 01:21:17 2026 +0800 refactor(proto): migrate UnnestExec serde (#23739) ## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123. --> - Closes #23510. ## Rationale for this change <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> The central protobuf dispatcher must know every built-in execution plan, which separates UnnestExec serialization from the type that owns it. Move encoding and decoding to plan-local hooks while retaining the deprecated helpers as compatibility wrappers. Extend roundtrip coverage to include non-default unnest options. ## What changes are included in this PR? <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> ## Are these changes tested? Yes <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? --> ## Are there any user-facing changes? <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> <!-- If there are any breaking changes to public APIs, please add the `api change` label. --> Signed-off-by: Jiawei Zhao <[email protected]> --- datafusion/physical-plan/src/unnest.rs | 121 +++++++++++++++++++++ datafusion/proto/src/physical_plan/mod.rs | 87 +++++---------- .../proto/tests/cases/roundtrip_physical_plan.rs | 21 +++- 3 files changed, 165 insertions(+), 64 deletions(-) diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 01c2f3ae27..fbe8492299 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/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)] diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 8acb979d78..f3f1d88e08 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -99,7 +99,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; @@ -123,10 +123,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) @@ -1027,14 +1024,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)? @@ -2398,32 +2387,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 { @@ -3713,39 +3695,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( diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 89647d5da3..e354105c76 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
