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-23794-9dd4eeac32492370e8b182a9a4215e5b24db300e in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit e784e2cb5c43a58ec5900d4df93449c571a9aef5 Author: Burak Şen <[email protected]> AuthorDate: Thu Jul 23 20:13:51 2026 +0300 refactor(proto): migrate SortExec and SortPreservingMergeExec serde (#23794) ## Which issue does this PR close? - Closes #23505. ## Rationale for this change Part of epic #23494. Moves `SortExec` and `SortPreservingMergeExec` protobuf serialization from central dispatch. ## What changes are included in this PR? Add protobuf serialization and deserialization to the physical plan implementations and deprecate the corresponding central proto methods. The wire format remains unchanged. ## Are these changes tested? Yes, existing round-trip tests cover these plans. ## Are there any user-facing changes? Existing methods are deprecated with no immediate API change. --- datafusion/physical-plan/src/sorts/sort.rs | 114 ++++++++++ .../src/sorts/sort_preserving_merge.rs | 92 ++++++++ datafusion/proto/src/physical_plan/mod.rs | 246 +++++---------------- 3 files changed, 256 insertions(+), 196 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 685a5e4bbf..a9b754dee6 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -1553,6 +1553,120 @@ impl ExecutionPlan for SortExec { updated_node: Some(new_sort), }) } + #[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 expr = self + .expr() + .iter() + .map(|sort_expr| { + let sort_node = Box::new(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&sort_expr.expr)?)), + asc: !sort_expr.options.descending, + nulls_first: sort_expr.options.nulls_first, + }); + Ok(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Sort( + sort_node, + )), + }) + }) + .collect::<Result<Vec<_>>>()?; + let dynamic_filter = match self.dynamic_filter_expr() { + Some(df) => { + let df_expr: Arc<dyn PhysicalExpr> = df; + Some(ctx.encode_expr(&df_expr)?) + } + None => None, + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Sort(Box::new( + protobuf::SortExecNode { + input: Some(Box::new(input)), + expr, + fetch: match self.fetch() { + Some(n) => n as i64, + None => -1, + }, + preserve_partitioning: self.preserve_partitioning(), + dynamic_filter, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SortExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + use protobuf::physical_expr_node::ExprType; + let sort = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Sort, + "SortExec", + ); + let input = + ctx.decode_required_child(sort.input.as_deref(), "SortExec", "input")?; + let input_schema = input.schema(); + let exprs = sort + .expr + .iter() + .map(|expr| { + let Some(ExprType::Sort(sort_expr)) = expr.expr_type.as_ref() else { + return datafusion_common::internal_err!( + "SortExec expr must be a sort expression" + ); + }; + let expr_node = sort_expr.expr.as_deref().ok_or_else(|| { + internal_datafusion_err!( + "SortExec sort expression is missing its inner expr" + ) + })?; + Ok(PhysicalSortExpr { + expr: ctx.decode_expr(expr_node, input_schema.as_ref())?, + options: arrow::compute::SortOptions { + descending: !sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }, + }) + }) + .collect::<Result<Vec<_>>>()?; + let Some(ordering) = LexOrdering::new(exprs) else { + return datafusion_common::internal_err!("SortExec requires an ordering"); + }; + let fetch = (sort.fetch >= 0).then_some(sort.fetch as usize); + let new_sort = SortExec::new(ordering, input) + .with_fetch(fetch) + .with_preserve_partitioning(sort.preserve_partitioning); + + let new_sort = if let Some(df_proto) = &sort.dynamic_filter { + let df_expr = + ctx.decode_expr(df_proto, new_sort.input().schema().as_ref())?; + let df = (df_expr as Arc<dyn std::any::Any + Send + Sync>) + .downcast::<DynamicFilterPhysicalExpr>() + .map_err(|_| { + internal_datafusion_err!( + "SortExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" + ) + })?; + new_sort.with_dynamic_filter_expr(df)? + } else { + new_sort + }; + + Ok(Arc::new(new_sort)) + } } #[cfg(test)] diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index a25f3a2862..2add3e1eb8 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -437,6 +437,98 @@ impl ExecutionPlan for SortPreservingMergeExec { .with_fetch(self.fetch()), ))) } + #[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 expr = self + .expr() + .iter() + .map(|e| { + Ok(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Sort( + Box::new(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&e.expr)?)), + asc: !e.options.descending, + nulls_first: e.options.nulls_first, + }), + )), + }) + }) + .collect::<Result<Vec<_>>>()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::SortPreservingMerge( + Box::new(protobuf::SortPreservingMergeExecNode { + input: Some(Box::new(input)), + expr, + fetch: self.fetch().map(|f| f as i64).unwrap_or(-1), + }), + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SortPreservingMergeExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use arrow::compute::SortOptions; + use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; + use datafusion_proto_models::protobuf; + let spm = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::SortPreservingMerge, + "SortPreservingMergeExec", + ); + let input = ctx.decode_required_child( + spm.input.as_deref(), + "SortPreservingMergeExec", + "input", + )?; + let input_schema = input.schema(); + let exprs = spm + .expr + .iter() + .map(|e| { + let sort = match &e.expr_type { + Some(protobuf::physical_expr_node::ExprType::Sort(s)) => s, + _ => { + return internal_err!( + "SortPreservingMergeExec expression is not a sort expression" + ); + } + }; + let expr = ctx.decode_required_expr( + sort.expr.as_deref(), + input_schema.as_ref(), + "SortPreservingMergeExec", + "sort expression", + )?; + Ok(PhysicalSortExpr { + expr, + options: SortOptions { + descending: !sort.asc, + nulls_first: sort.nulls_first, + }, + }) + }) + .collect::<Result<Vec<_>>>()?; + let Some(ordering) = LexOrdering::new(exprs) else { + return internal_err!("SortPreservingMergeExec requires an ordering"); + }; + let fetch = (spm.fetch >= 0).then_some(spm.fetch as usize); + Ok(Arc::new( + SortPreservingMergeExec::new(ordering, input).with_fetch(fetch), + )) + } } #[cfg(test)] diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 89e5c3c49e..8acb979d78 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -21,7 +21,6 @@ use std::collections::HashMap; use std::fmt::Debug; use std::sync::Arc; -use arrow::compute::SortOptions; use arrow::datatypes::{IntervalMonthDayNanoType, Schema, SchemaRef}; use datafusion_catalog::memory::MemorySourceConfig; use datafusion_common::config::CsvOptions; @@ -829,11 +828,12 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::PlaceholderRow(placeholder) => { self.try_into_placeholder_row_physical_plan(placeholder, ctx) } - PhysicalPlanType::Sort(sort) => { - self.try_into_sort_physical_plan(sort, ctx, proto_converter) + PhysicalPlanType::Sort(_) => { + SortExec::try_from_proto(self.node(), &decode_ctx) + } + PhysicalPlanType::SortPreservingMerge(_) => { + SortPreservingMergeExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::SortPreservingMerge(sort) => self - .try_into_sort_preserving_merge_physical_plan(sort, ctx, proto_converter), PhysicalPlanType::Extension(extension) => { self.try_into_extension_physical_plan(extension, ctx, proto_converter) } @@ -977,14 +977,6 @@ pub trait PhysicalPlanNodeExt: Sized { ); } - if let Some(exec) = plan.downcast_ref::<SortExec>() { - return protobuf::PhysicalPlanNode::try_from_sort_exec( - exec, - codec, - proto_converter, - ); - } - if let Some(union) = plan.downcast_ref::<UnionExec>() { return protobuf::PhysicalPlanNode::try_from_union_exec( union, @@ -1001,14 +993,6 @@ pub trait PhysicalPlanNodeExt: Sized { ); } - if let Some(exec) = plan.downcast_ref::<SortPreservingMergeExec>() { - return protobuf::PhysicalPlanNode::try_from_sort_preserving_merge_exec( - exec, - codec, - proto_converter, - ); - } - if let Some(exec) = plan.downcast_ref::<NestedLoopJoinExec>() { return protobuf::PhysicalPlanNode::try_from_nested_loop_join_exec( exec, @@ -2115,130 +2099,48 @@ pub trait PhysicalPlanNodeExt: Sized { )) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `SortExec` deserializes itself via `SortExec::try_from_proto`" + )] fn try_into_sort_physical_plan( &self, sort: &protobuf::SortExecNode, ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<Arc<dyn ExecutionPlan>> { - let node = self.node(); - let input = into_physical_plan(&sort.input, ctx, proto_converter)?; - let exprs = sort - .expr - .iter() - .map(|expr| { - let expr = expr.expr_type.as_ref().ok_or_else(|| { - proto_error(format!( - "physical_plan::from_proto() Unexpected expr {node:?}" - )) - })?; - if let ExprType::Sort(sort_expr) = expr { - let expr = sort_expr - .expr - .as_ref() - .ok_or_else(|| { - proto_error(format!( - "physical_plan::from_proto() Unexpected sort expr {node:?}" - )) - })? - .as_ref(); - Ok(PhysicalSortExpr { - expr: proto_converter.proto_to_physical_expr( - expr, - input.schema().as_ref(), - ctx, - )?, - options: SortOptions { - descending: !sort_expr.asc, - nulls_first: sort_expr.nulls_first, - }, - }) - } else { - internal_err!( - "physical_plan::from_proto() {node:?}" - ) - } - }) - .collect::<Result<Vec<_>>>()?; - let Some(ordering) = LexOrdering::new(exprs) else { - return internal_err!("SortExec requires an ordering"); + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::Sort(Box::new(sort.clone()))), }; - let fetch = (sort.fetch >= 0).then_some(sort.fetch as _); - let new_sort = SortExec::new(ordering, input) - .with_fetch(fetch) - .with_preserve_partitioning(sort.preserve_partitioning); - - let new_sort = if let Some(dynamic_filter_proto) = &sort.dynamic_filter { - let dynamic_filter_expr = proto_converter.proto_to_physical_expr( - dynamic_filter_proto, - new_sort.input().schema().as_ref(), - ctx, - )?; - let df = (dynamic_filter_expr as Arc<dyn Any + Send + Sync>) - .downcast::<DynamicFilterPhysicalExpr>() - .map_err(|_| { - internal_datafusion_err!( - "SortExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" - ) - })?; - new_sort.with_dynamic_filter_expr(df)? - } else { - new_sort + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, }; - - Ok(Arc::new(new_sort)) + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + SortExec::try_from_proto(&node, &decode_ctx) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `SortPreservingMergeExec` deserializes itself via `SortPreservingMergeExec::try_from_proto`" + )] fn try_into_sort_preserving_merge_physical_plan( &self, sort: &protobuf::SortPreservingMergeExecNode, ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<Arc<dyn ExecutionPlan>> { - let node = self.node(); - let input = into_physical_plan(&sort.input, ctx, proto_converter)?; - let exprs = sort - .expr - .iter() - .map(|expr| { - let expr = expr.expr_type.as_ref().ok_or_else(|| { - proto_error(format!( - "physical_plan::from_proto() Unexpected expr {node:?}" - )) - })?; - if let ExprType::Sort(sort_expr) = expr { - let expr = sort_expr - .expr - .as_ref() - .ok_or_else(|| { - proto_error(format!( - "physical_plan::from_proto() Unexpected sort expr {node:?}" - )) - })? - .as_ref(); - Ok(PhysicalSortExpr { - expr: proto_converter.proto_to_physical_expr( - expr, - input.schema().as_ref(), - ctx, - )?, - options: SortOptions { - descending: !sort_expr.asc, - nulls_first: sort_expr.nulls_first, - }, - }) - } else { - internal_err!("physical_plan::from_proto() {node:?}") - } - }) - .collect::<Result<Vec<_>>>()?; - let Some(ordering) = LexOrdering::new(exprs) else { - return internal_err!("SortExec requires an ordering"); + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::SortPreservingMerge(Box::new( + sort.clone(), + ))), }; - let fetch = (sort.fetch >= 0).then_some(sort.fetch as _); - Ok(Arc::new( - SortPreservingMergeExec::new(ordering, input).with_fetch(fetch), - )) + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + SortPreservingMergeExec::try_from_proto(&node, &decode_ctx) } fn try_into_extension_physical_plan( @@ -3503,51 +3405,22 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `SortExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] fn try_from_sort_exec( exec: &SortExec, codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<protobuf::PhysicalPlanNode> { - let input = proto_converter.execution_plan_to_proto(exec.input(), codec)?; - let expr = exec - .expr() - .iter() - .map(|expr| { - let sort_expr = Box::new(protobuf::PhysicalSortExprNode { - expr: Some(Box::new( - proto_converter.physical_expr_to_proto(&expr.expr, codec)?, - )), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }); - Ok(protobuf::PhysicalExprNode { - expr_id: None, - expr_type: Some(ExprType::Sort(sort_expr)), - }) - }) - .collect::<Result<Vec<_>>>()?; - let dynamic_filter = exec - .dynamic_filter_expr() - .map(|df| { - let df_expr: Arc<dyn PhysicalExpr> = df as Arc<dyn PhysicalExpr>; - proto_converter.physical_expr_to_proto(&df_expr, codec) - }) - .transpose()?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Sort(Box::new( - protobuf::SortExecNode { - input: Some(Box::new(input)), - expr, - fetch: match exec.fetch() { - Some(n) => n as i64, - _ => -1, - }, - preserve_partitioning: exec.preserve_partitioning(), - dynamic_filter, - }, - ))), - }) + let encoder = ConverterPlanEncoder { + codec, + proto_converter, + }; + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + exec.try_to_proto(&encode_ctx)? + .ok_or_else(|| internal_datafusion_err!("SortExec is not serializable")) } fn try_from_union_exec( @@ -3594,41 +3467,22 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `SortPreservingMergeExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] fn try_from_sort_preserving_merge_exec( exec: &SortPreservingMergeExec, 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, - )?; - let expr = exec - .expr() - .iter() - .map(|expr| { - let sort_expr = Box::new(protobuf::PhysicalSortExprNode { - expr: Some(Box::new( - proto_converter.physical_expr_to_proto(&expr.expr, codec)?, - )), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }); - Ok(protobuf::PhysicalExprNode { - expr_id: None, - expr_type: Some(ExprType::Sort(sort_expr)), - }) - }) - .collect::<Result<Vec<_>>>()?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::SortPreservingMerge(Box::new( - protobuf::SortPreservingMergeExecNode { - input: Some(Box::new(input)), - expr, - fetch: exec.fetch().map(|f| f as i64).unwrap_or(-1), - }, - ))), + }; + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + exec.try_to_proto(&encode_ctx)?.ok_or_else(|| { + internal_datafusion_err!("SortPreservingMergeExec is not serializable") }) } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
