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-23736-17634176a4f406a50b31a3711514831c9d3ecfa0 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 9caaf4b20a9b6a0cd63783f6e92e1d7f5761a9e0 Author: Phoenix <[email protected]> AuthorDate: Fri Jul 24 09:07:40 2026 +0800 refactor(proto): migrate symmetric hash join serde (#23736) ## 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 #23509. ## 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. --> SymmetricHashJoinExec serialization still depends on centralized plan dispatch, keeping its join, streaming-mode, and ordering wire logic outside the plan that owns the state. The retained compatibility helpers also duplicate that logic and can drift from normal dispatch. Move encoding and decoding behind the plan's ExecutionPlan hooks, route central dispatch through them, and reduce the deprecated helpers to thin adapters. Preserve explicit enum and optional-ordering representations, and expand round-trip coverage across the serialized join state. ## 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. --> ## Are these changes tested? <!-- 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)? --> Yes ## 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. --> Signed-off-by: Jiawei Zhao <[email protected]> --- .../physical-plan/src/joins/symmetric_hash_join.rs | 303 +++++++++++++++++++++ datafusion/proto/src/physical_plan/mod.rs | 263 ++---------------- .../proto/tests/cases/roundtrip_physical_plan.rs | 133 ++++++--- 3 files changed, 427 insertions(+), 272 deletions(-) diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index b95e5e1e3d..95f4c35871 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -627,6 +627,309 @@ impl ExecutionPlan for SymmetricHashJoinExec { Ok(None) } } + + #[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 partition_mode = match self.partition_mode() { + StreamJoinPartitionMode::SinglePartition => { + protobuf::StreamPartitionMode::SinglePartition + } + StreamJoinPartitionMode::Partitioned => { + protobuf::StreamPartitionMode::PartitionedExec + } + }; + let filter = self + .filter() + .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 encode_sort_exprs = + |exprs: Option<&LexOrdering>| -> Result<Vec<protobuf::PhysicalSortExprNode>> { + exprs + .map(|exprs| { + exprs + .iter() + .map(|expr| { + Ok(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&expr.expr)?)), + asc: !expr.options.descending, + nulls_first: expr.options.nulls_first, + }) + }) + .collect::<Result<Vec<_>>>() + }) + .transpose() + .map(Option::unwrap_or_default) + }; + let left_sort_exprs = encode_sort_exprs(self.left_sort_exprs())?; + let right_sort_exprs = encode_sort_exprs(self.right_sort_exprs())?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::SymmetricHashJoin( + Box::new(protobuf::SymmetricHashJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + join_type: join_type.into(), + partition_mode: partition_mode.into(), + null_equality: null_equality.into(), + filter, + left_sort_exprs, + right_sort_exprs, + }), + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SymmetricHashJoinExec { + /// Reconstruct a [`SymmetricHashJoinExec`] 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 datafusion_common::internal_datafusion_err; + use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; + use datafusion_proto_models::protobuf; + + let sym_join = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::SymmetricHashJoin, + "SymmetricHashJoinExec", + ); + let left = ctx.decode_required_child( + sym_join.left.as_deref(), + "SymmetricHashJoinExec", + "left", + )?; + let right = ctx.decode_required_child( + sym_join.right.as_deref(), + "SymmetricHashJoinExec", + "right", + )?; + let left_schema = left.schema(); + let right_schema = right.schema(); + let on = sym_join + .on + .iter() + .map(|columns| { + let left = ctx.decode_required_expr( + columns.left.as_ref(), + left_schema.as_ref(), + "SymmetricHashJoinExec", + "on.left", + )?; + let right = ctx.decode_required_expr( + columns.right.as_ref(), + right_schema.as_ref(), + "SymmetricHashJoinExec", + "on.right", + )?; + Ok((left, right)) + }) + .collect::<Result<JoinOn>>()?; + + let join_type = + match protobuf::JoinType::try_from(sym_join.join_type).map_err(|_| { + internal_datafusion_err!( + "SymmetricHashJoinExec: unknown JoinType {}", + sym_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(sym_join.null_equality) + .map_err(|_| { + internal_datafusion_err!( + "SymmetricHashJoinExec: unknown NullEquality {}", + sym_join.null_equality + ) + })? { + protobuf::NullEquality::NullEqualsNothing => NullEquality::NullEqualsNothing, + protobuf::NullEquality::NullEqualsNull => NullEquality::NullEqualsNull, + }; + let partition_mode = + match protobuf::StreamPartitionMode::try_from(sym_join.partition_mode) + .map_err(|_| { + internal_datafusion_err!( + "SymmetricHashJoinExec: unknown StreamPartitionMode {}", + sym_join.partition_mode + ) + })? { + protobuf::StreamPartitionMode::SinglePartition => { + StreamJoinPartitionMode::SinglePartition + } + protobuf::StreamPartitionMode::PartitionedExec => { + StreamJoinPartitionMode::Partitioned + } + }; + let filter = sym_join + .filter + .as_ref() + .map(|filter| -> Result<JoinFilter> { + let schema: Schema = filter + .schema + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "SymmetricHashJoinExec: JoinFilter missing schema" + ) + })? + .try_into()?; + let expression = ctx.decode_required_expr( + filter.expression.as_ref(), + &schema, + "SymmetricHashJoinExec", + "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!( + "SymmetricHashJoinExec: 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 decode_sort_exprs = |sort_exprs: &[protobuf::PhysicalSortExprNode], + schema: &Schema, + field: &str| + -> Result<Option<LexOrdering>> { + let sort_exprs = sort_exprs + .iter() + .map(|sort_expr| { + let expr = ctx.decode_required_expr( + sort_expr.expr.as_deref(), + schema, + "SymmetricHashJoinExec", + field, + )?; + Ok(PhysicalSortExpr { + expr, + options: arrow::compute::SortOptions { + descending: !sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }, + }) + }) + .collect::<Result<Vec<_>>>()?; + Ok(LexOrdering::new(sort_exprs)) + }; + let left_sort_exprs = decode_sort_exprs( + &sym_join.left_sort_exprs, + left_schema.as_ref(), + "left_sort_exprs", + )?; + let right_sort_exprs = decode_sort_exprs( + &sym_join.right_sort_exprs, + right_schema.as_ref(), + "right_sort_exprs", + )?; + + Self::try_new( + left, + right, + on, + filter, + &join_type, + null_equality, + left_sort_exprs, + right_sort_exprs, + partition_mode, + ) + .map(|exec| Arc::new(exec) as _) + } } /// A stream that issues [RecordBatch]es as they arrive from the right of the join. diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 1b122fcde6..d62bafa883 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -83,7 +83,7 @@ use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter}; use datafusion_physical_plan::joins::{ CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, - StreamJoinPartitionMode, SymmetricHashJoinExec, + SymmetricHashJoinExec, }; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::memory::LazyMemoryExec; @@ -802,12 +802,9 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::HashJoin(hashjoin) => { self.try_into_hash_join_physical_plan(hashjoin, ctx, proto_converter) } - PhysicalPlanType::SymmetricHashJoin(sym_join) => self - .try_into_symmetric_hash_join_physical_plan( - sym_join, - ctx, - proto_converter, - ), + PhysicalPlanType::SymmetricHashJoin(_) => { + SymmetricHashJoinExec::try_from_proto(self.node(), &decode_ctx) + } PhysicalPlanType::Union(_) => { UnionExec::try_from_proto(self.node(), &decode_ctx) } @@ -920,14 +917,6 @@ pub trait PhysicalPlanNodeExt: Sized { ); } - if let Some(exec) = plan.downcast_ref::<SymmetricHashJoinExec>() { - return protobuf::PhysicalPlanNode::try_from_symmetric_hash_join_exec( - exec, - codec, - proto_converter, - ); - } - if let Some(exec) = plan.downcast_ref::<AggregateExec>() { return protobuf::PhysicalPlanNode::try_from_aggregate_exec( exec, @@ -1857,129 +1846,27 @@ pub trait PhysicalPlanNodeExt: Sized { Ok(Arc::new(hash_join)) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `SymmetricHashJoinExec` deserializes itself via `SymmetricHashJoinExec::try_from_proto`" + )] fn try_into_symmetric_hash_join_physical_plan( &self, sym_join: &protobuf::SymmetricHashJoinExecNode, ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<Arc<dyn ExecutionPlan>> { - let left = into_physical_plan(&sym_join.left, ctx, proto_converter)?; - let right = into_physical_plan(&sym_join.right, ctx, proto_converter)?; - let left_schema = left.schema(); - let right_schema = right.schema(); - let on = sym_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<_>>()?; - let join_type = - protobuf::JoinType::try_from(sym_join.join_type).map_err(|_| { - proto_error(format!( - "Received a SymmetricHashJoin message with unknown JoinType {}", - sym_join.join_type - )) - })?; - let null_equality = protobuf::NullEquality::try_from(sym_join.null_equality) - .map_err(|_| { - proto_error(format!( - "Received a SymmetricHashJoin message with unknown NullEquality {}", - sym_join.null_equality - )) - })?; - let filter = sym_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 HashJoinNode message with JoinSide in Filter {}", - i.side)) - )?; - - Ok(ColumnIndex { - index: i.index as usize, - side: side.into(), - }) - }) - .collect::<Result<_>>()?; - - Ok(JoinFilter::new(expression, column_indices, Arc::new(schema))) - }) - .map_or(Ok(None), |v: Result<JoinFilter>| v.map(Some))?; - - let left_sort_exprs = parse_physical_sort_exprs( - &sym_join.left_sort_exprs, - ctx, - &left_schema, - proto_converter, - )?; - let left_sort_exprs = LexOrdering::new(left_sort_exprs); - - let right_sort_exprs = parse_physical_sort_exprs( - &sym_join.right_sort_exprs, + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::SymmetricHashJoin(Box::new( + sym_join.clone(), + ))), + }; + let decoder = ConverterPlanDecoder { ctx, - &right_schema, proto_converter, - )?; - let right_sort_exprs = LexOrdering::new(right_sort_exprs); - - let partition_mode = protobuf::StreamPartitionMode::try_from( - sym_join.partition_mode, - ) - .map_err(|_| { - proto_error(format!( - "Received a SymmetricHashJoin message with unknown PartitionMode {}", - sym_join.partition_mode - )) - })?; - let partition_mode = match partition_mode { - protobuf::StreamPartitionMode::SinglePartition => { - StreamJoinPartitionMode::SinglePartition - } - protobuf::StreamPartitionMode::PartitionedExec => { - StreamJoinPartitionMode::Partitioned - } }; - SymmetricHashJoinExec::try_new( - left, - right, - on, - filter, - &JoinType::from_proto(join_type), - NullEquality::from_proto(null_equality), - left_sort_exprs, - right_sort_exprs, - partition_mode, - ) - .map(|e| Arc::new(e) as _) + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + SymmetricHashJoinExec::try_from_proto(&node, &decode_ctx) } #[deprecated( @@ -2786,124 +2673,22 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `SymmetricHashJoinExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] fn try_from_symmetric_hash_join_exec( exec: &SymmetricHashJoinExec, 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 partition_mode = match exec.partition_mode() { - StreamJoinPartitionMode::SinglePartition => { - protobuf::StreamPartitionMode::SinglePartition - } - StreamJoinPartitionMode::Partitioned => { - protobuf::StreamPartitionMode::PartitionedExec - } }; - - let left_sort_exprs = exec - .left_sort_exprs() - .map(|exprs| { - exprs - .iter() - .map(|expr| { - Ok(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, - }) - }) - .collect::<Result<Vec<_>>>() - }) - .transpose()? - .unwrap_or(vec![]); - - let right_sort_exprs = exec - .right_sort_exprs() - .map(|exprs| { - exprs - .iter() - .map(|expr| { - Ok(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, - }) - }) - .collect::<Result<Vec<_>>>() - }) - .transpose()? - .unwrap_or(vec![]); - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::SymmetricHashJoin(Box::new( - protobuf::SymmetricHashJoinExecNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - on, - join_type: join_type.into(), - partition_mode: partition_mode.into(), - null_equality: null_equality.into(), - left_sort_exprs, - right_sort_exprs, - filter, - }, - ))), + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + exec.try_to_proto(&encode_ctx)?.ok_or_else(|| { + internal_datafusion_err!("SymmetricHashJoinExec is not serializable") }) } diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index e354105c76..889d42df40 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -71,6 +71,7 @@ use datafusion::physical_plan::expressions::{ cast, col, in_list, like, lit, }; use datafusion::physical_plan::filter::{FilterExec, FilterExecBuilder}; +use datafusion::physical_plan::joins::utils::{ColumnIndex, JoinFilter}; use datafusion::physical_plan::joins::{ HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, StreamJoinPartitionMode, SymmetricHashJoinExec, @@ -102,7 +103,7 @@ use datafusion_common::file_options::json_writer::JsonWriterOptions; use datafusion_common::parsers::CompressionTypeVariant; use datafusion_common::stats::Precision; use datafusion_common::{ - DataFusionError, NullEquality, Result, UnnestOptions, exec_datafusion_err, + DataFusionError, JoinSide, NullEquality, Result, UnnestOptions, exec_datafusion_err, internal_datafusion_err, internal_err, not_impl_err, }; use datafusion_datasource::file::FileSource; @@ -2098,17 +2099,66 @@ fn roundtrip_parquet_sink() -> Result<()> { #[test] fn roundtrip_sym_hash_join() -> Result<()> { - let field_a = Field::new("col", DataType::Int64, false); + 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()]); - let schema_right = Schema::new(vec![field_a]); + let schema_right = Schema::new(vec![field_b.clone()]); let on = vec![( - Arc::new(Column::new("col", schema_left.index_of("col")?)) as _, - Arc::new(Column::new("col", schema_right.index_of("col")?)) as _, + Arc::new(Column::new("col_a", schema_left.index_of("col_a")?)) as _, + Arc::new(Column::new("col_b", schema_right.index_of("col_b")?)) as _, )]; + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("col_a", 0)), + Operator::Gt, + Arc::new(Column::new("col_b", 1)), + )), + vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![field_a, field_b])), + ); let schema_left = Arc::new(schema_left); let schema_right = Arc::new(schema_right); - for join_type in &[ + let left_order: LexOrdering = [PhysicalSortExpr { + expr: Arc::new(Column::new("col_a", schema_left.index_of("col_a")?)), + options: SortOptions { + descending: true, + nulls_first: false, + }, + }] + .into(); + let right_order: LexOrdering = [PhysicalSortExpr { + expr: Arc::new(Column::new("col_b", schema_right.index_of("col_b")?)), + options: SortOptions { + descending: false, + nulls_first: true, + }, + }] + .into(); + let ordering_cases = [ + (None, None), + (Some(left_order.clone()), None), + (None, Some(right_order.clone())), + (Some(left_order), Some(right_order)), + ]; + let ordering_options = |ordering: Option<&LexOrdering>| { + ordering + .map(|ordering| ordering.iter().map(|expr| expr.options).collect::<Vec<_>>()) + }; + + for join_type in [ JoinType::Inner, JoinType::Left, JoinType::Right, @@ -2117,36 +2167,53 @@ fn roundtrip_sym_hash_join() -> Result<()> { JoinType::RightAnti, JoinType::LeftSemi, JoinType::RightSemi, + JoinType::LeftMark, + JoinType::RightMark, ] { - for partition_mode in &[ - StreamJoinPartitionMode::Partitioned, - StreamJoinPartitionMode::SinglePartition, + for null_equality in [ + NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull, ] { - for left_order in &[ - None, - LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::new(Column::new("col", schema_left.index_of("col")?)), - options: Default::default(), - }]), - ] { - for right_order in [ - None, - LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::new(Column::new("col", schema_right.index_of("col")?)), - options: Default::default(), - }]), + for filter in [None, Some(filter.clone())] { + for partition_mode in [ + StreamJoinPartitionMode::Partitioned, + StreamJoinPartitionMode::SinglePartition, ] { - roundtrip_test(Arc::new(SymmetricHashJoinExec::try_new( - Arc::new(EmptyExec::new(schema_left.clone())), - Arc::new(EmptyExec::new(schema_right.clone())), - on.clone(), - None, - join_type, - NullEquality::NullEqualsNothing, - left_order.clone(), - right_order, - *partition_mode, - )?))?; + for (left_order, right_order) in &ordering_cases { + let result = roundtrip_test_and_return( + Arc::new(SymmetricHashJoinExec::try_new( + Arc::new(EmptyExec::new(schema_left.clone())), + Arc::new(EmptyExec::new(schema_right.clone())), + on.clone(), + filter.clone(), + &join_type, + null_equality, + left_order.clone(), + right_order.clone(), + partition_mode, + )?), + &ctx, + &codec, + &proto_converter, + )?; + let result = + result.downcast_ref::<SymmetricHashJoinExec>().unwrap(); + assert_eq!(result.join_type(), &join_type); + assert_eq!(result.null_equality(), null_equality); + assert_eq!(result.partition_mode(), partition_mode); + assert_eq!( + ordering_options(result.left_sort_exprs()), + ordering_options(left_order.as_ref()) + ); + assert_eq!( + ordering_options(result.right_sort_exprs()), + ordering_options(right_order.as_ref()) + ); + assert_eq!( + result.filter().map(JoinFilter::column_indices), + filter.as_ref().map(JoinFilter::column_indices) + ); + } } } } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
