From c118cbcc7c6911f617921d36239197d529bff9c2 Mon Sep 17 00:00:00 2001 From: buraksenn Date: Wed, 22 Jul 2026 12:36:10 +0300 Subject: [PATCH 1/2] Migrate SortExec and SortPreservingMergeExec protobuf serialization --- datafusion/physical-plan/src/sorts/sort.rs | 114 ++++++++ .../src/sorts/sort_preserving_merge.rs | 92 +++++++ datafusion/proto/src/physical_plan/mod.rs | 245 ++++-------------- 3 files changed, 256 insertions(+), 195 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 685a5e4bbf55a..a9b754dee68bd 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> { + 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::>>()?; + let dynamic_filter = match self.dynamic_filter_expr() { + Some(df) => { + let df_expr: Arc = 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> { + 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::>>()?; + 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) + .downcast::() + .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 a25f3a2862b4d..2add3e1eb82f0 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> { + 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::>>()?; + 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> { + 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::>>()?; + 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 cea334e42aace..a0bc36e69a040 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -829,11 +829,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) } @@ -993,14 +994,6 @@ pub trait PhysicalPlanNodeExt: Sized { ); } - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_sort_exec( - exec, - codec, - proto_converter, - ); - } - if let Some(union) = plan.downcast_ref::() { return protobuf::PhysicalPlanNode::try_from_union_exec( union, @@ -1017,14 +1010,6 @@ pub trait PhysicalPlanNodeExt: Sized { ); } - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_sort_preserving_merge_exec( - exec, - codec, - proto_converter, - ); - } - if let Some(exec) = plan.downcast_ref::() { return protobuf::PhysicalPlanNode::try_from_nested_loop_join_exec( exec, @@ -2123,130 +2108,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> { - 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::>>()?; - 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) - .downcast::() - .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> { - 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::>>()?; - 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( @@ -3536,51 +3439,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 { - 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::>>()?; - let dynamic_filter = exec - .dynamic_filter_expr() - .map(|df| { - let df_expr: Arc = df as Arc; - 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( @@ -3627,41 +3501,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 { - 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::>>()?; - 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") }) } From 895d9441ee6c9039a0357568df36c1d9111e60a6 Mon Sep 17 00:00:00 2001 From: buraksenn Date: Wed, 22 Jul 2026 14:24:06 +0300 Subject: [PATCH 2/2] fix cargo clippy --- datafusion/proto/src/physical_plan/mod.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index a0bc36e69a040..8f5f82887b9e9 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;