From a535900fff70fd23ecd9ba73703494b5891acf91 Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Mon, 20 Jul 2026 14:49:42 +0800 Subject: [PATCH 1/2] refactor(proto): migrate sort merge join serde SortMergeJoin serialization still depends on centralized plan downcasts, keeping protobuf ownership separate from the execution plan. Move encoding and decoding behind its ExecutionPlan hooks while preserving existing wire fields. Keep the legacy helpers as deprecated compatibility APIs and verify join enums and sort options roundtrip. Closes #23508 Signed-off-by: Jiawei Zhao --- .../src/joins/sort_merge_join/exec.rs | 242 ++++++++++++++++++ datafusion/proto/src/physical_plan/mod.rs | 20 +- .../tests/cases/roundtrip_physical_plan.rs | 67 +++-- 3 files changed, 299 insertions(+), 30 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 5d3621c49219b..304f547bdd570 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -652,4 +652,246 @@ impl ExecutionPlan for SortMergeJoinExec { self.null_equality, )?))) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + let on = self + .on() + .iter() + .map(|(left, right)| { + Ok(protobuf::JoinOn { + left: Some(ctx.encode_expr(left)?), + right: Some(ctx.encode_expr(right)?), + }) + }) + .collect::>>()?; + + let join_type = match self.join_type() { + JoinType::Inner => protobuf::JoinType::Inner, + JoinType::Left => protobuf::JoinType::Left, + JoinType::Right => protobuf::JoinType::Right, + JoinType::Full => protobuf::JoinType::Full, + JoinType::LeftSemi => protobuf::JoinType::Leftsemi, + JoinType::RightSemi => protobuf::JoinType::Rightsemi, + JoinType::LeftAnti => protobuf::JoinType::Leftanti, + JoinType::RightAnti => protobuf::JoinType::Rightanti, + JoinType::LeftMark => protobuf::JoinType::Leftmark, + JoinType::RightMark => protobuf::JoinType::Rightmark, + }; + let null_equality = match self.null_equality() { + NullEquality::NullEqualsNothing => protobuf::NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull => protobuf::NullEquality::NullEqualsNull, + }; + let filter = self + .filter() + .as_ref() + .map(|filter| -> Result { + let expression = ctx.encode_expr(filter.expression())?; + let column_indices = filter + .column_indices() + .iter() + .map(|column_index| { + let side = match column_index.side { + JoinSide::Left => protobuf::JoinSide::LeftSide, + JoinSide::Right => protobuf::JoinSide::RightSide, + JoinSide::None => protobuf::JoinSide::None, + }; + protobuf::ColumnIndex { + index: column_index.index as u32, + side: side.into(), + } + }) + .collect(); + Ok(protobuf::JoinFilter { + expression: Some(expression), + column_indices, + schema: Some(filter.schema().as_ref().try_into()?), + }) + }) + .transpose()?; + let sort_options = self + .sort_options() + .iter() + .map(|options| protobuf::SortExprNode { + expr: None, + asc: !options.descending, + nulls_first: options.nulls_first, + }) + .collect(); + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::SortMergeJoin(Box::new( + protobuf::SortMergeJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + join_type: join_type.into(), + filter, + sort_options, + null_equality: null_equality.into(), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SortMergeJoinExec { + /// Reconstruct a [`SortMergeJoinExec`] 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> { + use crate::joins::utils::ColumnIndex; + use arrow::datatypes::Schema; + use datafusion_common::internal_datafusion_err; + use datafusion_proto_models::protobuf; + + let sort_join = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::SortMergeJoin, + "SortMergeJoinExec", + ); + let left = ctx.decode_required_child( + sort_join.left.as_deref(), + "SortMergeJoinExec", + "left", + )?; + let right = ctx.decode_required_child( + sort_join.right.as_deref(), + "SortMergeJoinExec", + "right", + )?; + let left_schema = left.schema(); + let right_schema = right.schema(); + let on = sort_join + .on + .iter() + .map(|columns| { + let left = ctx.decode_required_expr( + columns.left.as_ref(), + left_schema.as_ref(), + "SortMergeJoinExec", + "on.left", + )?; + let right = ctx.decode_required_expr( + columns.right.as_ref(), + right_schema.as_ref(), + "SortMergeJoinExec", + "on.right", + )?; + Ok((left, right)) + }) + .collect::>()?; + + let join_type = + match protobuf::JoinType::try_from(sort_join.join_type).map_err(|_| { + internal_datafusion_err!( + "SortMergeJoinExec: unknown JoinType {}", + sort_join.join_type + ) + })? { + protobuf::JoinType::Inner => JoinType::Inner, + protobuf::JoinType::Left => JoinType::Left, + protobuf::JoinType::Right => JoinType::Right, + protobuf::JoinType::Full => JoinType::Full, + protobuf::JoinType::Leftsemi => JoinType::LeftSemi, + protobuf::JoinType::Rightsemi => JoinType::RightSemi, + protobuf::JoinType::Leftanti => JoinType::LeftAnti, + protobuf::JoinType::Rightanti => JoinType::RightAnti, + protobuf::JoinType::Leftmark => JoinType::LeftMark, + protobuf::JoinType::Rightmark => JoinType::RightMark, + }; + let null_equality = match protobuf::NullEquality::try_from( + sort_join.null_equality, + ) + .map_err(|_| { + internal_datafusion_err!( + "SortMergeJoinExec: unknown NullEquality {}", + sort_join.null_equality + ) + })? { + protobuf::NullEquality::NullEqualsNothing => NullEquality::NullEqualsNothing, + protobuf::NullEquality::NullEqualsNull => NullEquality::NullEqualsNull, + }; + let filter = sort_join + .filter + .as_ref() + .map(|filter| -> Result { + let schema: Schema = filter + .schema + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "SortMergeJoinExec: JoinFilter missing schema" + ) + })? + .try_into()?; + let expression = ctx.decode_required_expr( + filter.expression.as_ref(), + &schema, + "SortMergeJoinExec", + "filter.expression", + )?; + let column_indices = filter + .column_indices + .iter() + .map(|column_index| { + let side = protobuf::JoinSide::try_from(column_index.side) + .map_err(|_| { + internal_datafusion_err!( + "SortMergeJoinExec: unknown JoinSide {}", + column_index.side + ) + })?; + let side = match side { + protobuf::JoinSide::LeftSide => JoinSide::Left, + protobuf::JoinSide::RightSide => JoinSide::Right, + protobuf::JoinSide::None => JoinSide::None, + }; + Ok(ColumnIndex { + index: column_index.index as usize, + side, + }) + }) + .collect::>>()?; + Ok(JoinFilter::new( + expression, + column_indices, + Arc::new(schema), + )) + }) + .transpose()?; + let sort_options = sort_join + .sort_options + .iter() + .map(|options| SortOptions { + descending: !options.asc, + nulls_first: options.nulls_first, + }) + .collect(); + + Ok(Arc::new(Self::try_new( + left, + right, + on, + filter, + join_type, + sort_options, + null_equality, + )?)) + } } diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index e5f8aa072db71..7ad9c585dc1bd 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -856,8 +856,8 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::GenerateSeries(generate_series) => { self.try_into_generate_series_physical_plan(generate_series) } - PhysicalPlanType::SortMergeJoin(sort_join) => { - self.try_into_sort_join(sort_join, ctx, proto_converter) + PhysicalPlanType::SortMergeJoin(_) => { + SortMergeJoinExec::try_from_proto(self.node(), &decode_ctx) } PhysicalPlanType::AsyncFunc(async_func) => { self.try_into_async_func_physical_plan(async_func, ctx, proto_converter) @@ -951,14 +951,6 @@ pub trait PhysicalPlanNodeExt: Sized { ); } - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_sort_merge_join_exec( - exec, - codec, - proto_converter, - ); - } - if let Some(exec) = plan.downcast_ref::() { return protobuf::PhysicalPlanNode::try_from_cross_join_exec( exec, @@ -2596,6 +2588,10 @@ pub trait PhysicalPlanNodeExt: Sized { protobuf::GenerateSeriesName::GsRange => "range", } } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `SortMergeJoinExec` deserializes itself via `SortMergeJoinExec::try_from_proto`" + )] fn try_into_sort_join( &self, sort_join: &SortMergeJoinExecNode, @@ -3254,6 +3250,10 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `SortMergeJoinExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] fn try_from_sort_merge_join_exec( exec: &SortMergeJoinExec, codec: &dyn PhysicalExtensionCodec, diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index cbd6fd912abef..d8c920412ece7 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -2724,6 +2724,9 @@ async fn analyze_roundtrip_unoptimized() -> Result<()> { #[test] fn roundtrip_sort_merge_join() -> Result<()> { + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; let field_a = Field::new("col_a", DataType::Int64, false); let field_b = Field::new("col_b", DataType::Int64, false); let schema_left = Schema::new(vec![field_a.clone()]); @@ -2754,26 +2757,50 @@ fn roundtrip_sort_merge_join() -> Result<()> { let schema_left = Arc::new(schema_left); let schema_right = Arc::new(schema_right); - for filter in [None, Some(filter)] { - for join_type in [ - JoinType::Inner, - JoinType::Left, - JoinType::Right, - JoinType::Full, - JoinType::LeftAnti, - JoinType::RightAnti, - JoinType::LeftSemi, - JoinType::RightSemi, - ] { - roundtrip_test(Arc::new(SortMergeJoinExec::try_new( - Arc::new(EmptyExec::new(schema_left.clone())), - Arc::new(EmptyExec::new(schema_right.clone())), - on.clone(), - filter.clone(), - join_type, - vec![Default::default()], - NullEquality::NullEqualsNothing, - )?))?; + let sort_options = vec![SortOptions { + descending: true, + nulls_first: false, + }]; + for null_equality in [ + NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull, + ] { + for filter in [None, Some(filter.clone())] { + for join_type in [ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftAnti, + JoinType::RightAnti, + JoinType::LeftSemi, + JoinType::RightSemi, + JoinType::LeftMark, + JoinType::RightMark, + ] { + let result = roundtrip_test_and_return( + Arc::new(SortMergeJoinExec::try_new( + Arc::new(EmptyExec::new(schema_left.clone())), + Arc::new(EmptyExec::new(schema_right.clone())), + on.clone(), + filter.clone(), + join_type, + sort_options.clone(), + null_equality, + )?), + &ctx, + &codec, + &proto_converter, + )?; + let result = result.downcast_ref::().unwrap(); + assert_eq!(result.join_type(), join_type); + assert_eq!(result.null_equality(), null_equality); + assert_eq!(result.sort_options(), sort_options); + assert_eq!( + result.filter().as_ref().map(|f| f.column_indices()), + filter.as_ref().map(|f| f.column_indices()) + ); + } } } Ok(()) From 54ad24aae7ea47ae93d9f35cfb733666f93e9624 Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Tue, 21 Jul 2026 13:33:39 +0800 Subject: [PATCH 2/2] refactor(proto): delegate deprecated sort merge join serde SortMergeJoinExec serde now lives in per-plan hooks, but the deprecated PhysicalPlanNodeExt helpers kept independent implementations. Duplicating join, filter, and sort-option conversions allowed the compatibility path to drift from normal dispatch. Build the existing encode and decode adapter contexts in those helpers and forward to SortMergeJoinExec's canonical hooks. Remove the protobuf sort-expression alias that the duplicate encoder no longer needs. Signed-off-by: Jiawei Zhao --- datafusion/proto/src/physical_plan/mod.rs | 197 ++-------------------- 1 file changed, 18 insertions(+), 179 deletions(-) diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 7ad9c585dc1bd..89e91cd545a50 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -120,8 +120,8 @@ 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, SortExprNode, SortMergeJoinExecNode, - proto_error, window_agg_exec_node, + self, ListUnnest as ProtoListUnnest, SortMergeJoinExecNode, proto_error, + window_agg_exec_node, }; pub mod from_proto; @@ -2598,106 +2598,17 @@ pub trait PhysicalPlanNodeExt: Sized { ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let left = into_physical_plan(&sort_join.left, ctx, proto_converter)?; - let left_schema = left.schema(); - let right = into_physical_plan(&sort_join.right, ctx, proto_converter)?; - let right_schema = right.schema(); - - let filter = sort_join - .filter - .as_ref() - .map(|f| { - let schema = f - .schema - .as_ref() - .ok_or_else(|| proto_error("Missing JoinFilter schema"))? - .try_into()?; - - let expression = proto_converter.proto_to_physical_expr( - f.expression.as_ref().ok_or_else(|| { - proto_error("Unexpected empty filter expression") - })?, - &schema, - ctx, - )?; - let column_indices = f - .column_indices - .iter() - .map(|i| { - let side = - protobuf::JoinSide::try_from(i.side).map_err(|_| { - proto_error(format!( - "Received a SortMergeJoinExecNode message with JoinSide in Filter {}", - i.side - )) - })?; - - Ok(ColumnIndex { - index: i.index as usize, - side: side.into(), - }) - }) - .collect::>>()?; - - Ok(JoinFilter::new( - expression, - column_indices, - Arc::new(schema), - )) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - let join_type = - protobuf::JoinType::try_from(sort_join.join_type).map_err(|_| { - proto_error(format!( - "Received a SortMergeJoinExecNode message with unknown JoinType {}", - sort_join.join_type - )) - })?; - - let null_equality = protobuf::NullEquality::try_from(sort_join.null_equality) - .map_err(|_| { - proto_error(format!( - "Received a SortMergeJoinExecNode message with unknown NullEquality {}", - sort_join.null_equality - )) - })?; - - let sort_options = sort_join - .sort_options - .iter() - .map(|e| SortOptions { - descending: !e.asc, - nulls_first: e.nulls_first, - }) - .collect(); - let on = sort_join - .on - .iter() - .map(|col| { - let left = proto_converter.proto_to_physical_expr( - &col.left.clone().unwrap(), - left_schema.as_ref(), - ctx, - )?; - let right = proto_converter.proto_to_physical_expr( - &col.right.clone().unwrap(), - right_schema.as_ref(), - ctx, - )?; - Ok((left, right)) - }) - .collect::>()?; - - Ok(Arc::new(SortMergeJoinExec::try_new( - left, - right, - on, - filter, - JoinType::from_proto(join_type), - sort_options, - NullEquality::from_proto(null_equality), - )?)) + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::SortMergeJoin(Box::new( + sort_join.clone(), + ))), + }; + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + SortMergeJoinExec::try_from_proto(&node, &decode_ctx) } fn try_into_generate_series_physical_plan( @@ -3259,85 +3170,13 @@ pub trait PhysicalPlanNodeExt: Sized { codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result { - let left = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.left().to_owned(), - codec, - proto_converter, - )?; - let right = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.right().to_owned(), + let encoder = ConverterPlanEncoder { codec, proto_converter, - )?; - let on = exec - .on() - .iter() - .map(|tuple| { - let l = proto_converter.physical_expr_to_proto(&tuple.0, codec)?; - let r = proto_converter.physical_expr_to_proto(&tuple.1, codec)?; - Ok::<_, DataFusionError>(protobuf::JoinOn { - left: Some(l), - right: Some(r), - }) - }) - .collect::>()?; - let join_type = protobuf::JoinType::from_proto(exec.join_type().to_owned()); - let null_equality = protobuf::NullEquality::from_proto(exec.null_equality()); - let filter = exec - .filter() - .as_ref() - .map(|f| { - let expression = - proto_converter.physical_expr_to_proto(f.expression(), codec)?; - let column_indices = f - .column_indices() - .iter() - .map(|i| { - let side: protobuf::JoinSide = i.side.to_owned().into(); - protobuf::ColumnIndex { - index: i.index as u32, - side: side.into(), - } - }) - .collect(); - let schema = f.schema().as_ref().try_into()?; - Ok(protobuf::JoinFilter { - expression: Some(expression), - column_indices, - schema: Some(schema), - }) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - let sort_options = exec - .sort_options() - .iter() - .map( - |SortOptions { - descending, - nulls_first, - }| { - SortExprNode { - expr: None, - asc: !*descending, - nulls_first: *nulls_first, - } - }, - ) - .collect(); - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::SortMergeJoin(Box::new( - SortMergeJoinExecNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - on, - join_type: join_type.into(), - null_equality: null_equality.into(), - filter, - sort_options, - }, - ))), + }; + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + exec.try_to_proto(&encode_ctx)?.ok_or_else(|| { + internal_datafusion_err!("SortMergeJoinExec is not serializable") }) }