This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 9a50c67db0 refactor(proto): migrate sort merge join serde (#23712)
9a50c67db0 is described below
commit 9a50c67db0c8a1b23d4068419e9fefd0804c1b0d
Author: Phoenix <[email protected]>
AuthorDate: Tue Jul 21 14:24:11 2026 +0800
refactor(proto): migrate sort merge join serde (#23712)
## 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 #23508 .
## 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.
-->
`SortMergeJoinExec` serialization currently depends on centralized
encoding and decoding logic in `datafusion-proto`. This keeps protobuf
ownership separate from the execution plan and contributes to the
central type-dispatch chain tracked by #23494.
Moving serialization into `SortMergeJoinExec` keeps the implementation
with the plan that owns the relevant state and aligns it with the new
per-plan serialization model.
## 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.
-->
- Add `ExecutionPlan::try_to_proto` for `SortMergeJoinExec`.
- Add `SortMergeJoinExec::try_from_proto`.
- Preserve serialization of:
- Both child plans.
- Join-key expressions.
- Join type and null-equality behavior.
- Optional join filters and column origins.
- Per-key sort options.
- Convert `JoinType`, `JoinSide`, and `NullEquality` through exhaustive
by-name matches.
- Route decoding through the plan-local implementation and remove the
centralized encoding downcast.
- Retain the old public conversion helpers as deprecated compatibility
APIs.
- Preserve the existing protobuf wire representation.
- Expand roundtrip coverage for every join type, both null-equality
modes, filters, column origins, and non-default sort options.
## 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.
-->
There are no runtime behavior or protobuf wire-format changes.
The old `PhysicalPlanNodeExt` conversion helpers remain available but
are deprecated in favor of the plan-local serialization hooks.
---------
Signed-off-by: Jiawei Zhao <[email protected]>
---
.../src/joins/sort_merge_join/exec.rs | 242 +++++++++++++++++++++
datafusion/proto/src/physical_plan/mod.rs | 217 +++---------------
.../proto/tests/cases/roundtrip_physical_plan.rs | 67 ++++--
3 files changed, 317 insertions(+), 209 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 5d3621c492..304f547bdd 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<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+ 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::<Result<Vec<_>>>()?;
+
+ 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<protobuf::JoinFilter> {
+ 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<Arc<dyn ExecutionPlan>> {
+ 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::<Result<JoinOn>>()?;
+
+ 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<JoinFilter> {
+ 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::<Result<Vec<_>>>()?;
+ 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 fa18b1ffc6..cea334e42a 100644
--- a/datafusion/proto/src/physical_plan/mod.rs
+++ b/datafusion/proto/src/physical_plan/mod.rs
@@ -125,8 +125,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;
@@ -862,8 +862,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)
@@ -949,14 +949,6 @@ pub trait PhysicalPlanNodeExt: Sized {
);
}
- if let Some(exec) = plan.downcast_ref::<SortMergeJoinExec>() {
- return protobuf::PhysicalPlanNode::try_from_sort_merge_join_exec(
- exec,
- codec,
- proto_converter,
- );
- }
-
if let Some(exec) = plan.downcast_ref::<CrossJoinExec>() {
return protobuf::PhysicalPlanNode::try_from_cross_join_exec(
exec,
@@ -2546,112 +2538,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<Arc<dyn ExecutionPlan>> {
- 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::<Result<Vec<_>>>()?;
-
- Ok(JoinFilter::new(
- expression,
- column_indices,
- Arc::new(schema),
- ))
- })
- .map_or(Ok(None), |v: Result<JoinFilter>| 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::<Result<_>>()?;
-
- 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(
@@ -3212,90 +3119,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<protobuf::PhysicalPlanNode> {
- 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::<Result<_>>()?;
- 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<protobuf::JoinFilter>| 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 8027c0fa78..2671f4d015 100644
--- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
+++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
@@ -2788,6 +2788,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()]);
@@ -2818,26 +2821,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::<SortMergeJoinExec>().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(())
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
