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..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; @@ -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,112 +2588,27 @@ 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, 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( @@ -3254,90 +3161,22 @@ 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, 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") }) } 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(())