diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index a8162f137ed0a..76af9b0c29218 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -243,6 +243,40 @@ async fn test_remove_unnecessary_sort5() -> Result<()> { Ok(()) } +#[tokio::test] +async fn test_propagate_in_right_mark_join() -> Result<()> { + let left_schema = create_test_schema2()?; + let right_schema = create_test_schema3()?; + let left_input = memory_exec(&left_schema); + let parquet_ordering = [sort_expr("a", &right_schema)].into(); + let right_input = + parquet_exec_with_sort(right_schema.clone(), vec![parquet_ordering]); + let on = vec![( + Arc::new(Column::new_with_schema("col_a", &left_schema)?) as _, + Arc::new(Column::new_with_schema("c", &right_schema)?) as _, + )]; + let join = hash_join_exec(left_input, right_input, on, None, &JoinType::RightMark)?; + let physical_plan = sort_exec([sort_expr("a", &join.schema())].into(), join); + + let test = EnforceSortingTest::new(physical_plan).with_repartition_sorts(true); + assert_snapshot!(test.run(), @r" + Input Plan: + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + HashJoinExec: mode=Partitioned, join_type=RightMark, on=[(col_a@0, c@2)] + DataSourceExec: partitions=1, partition_sizes=[0] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [a@0 ASC] + HashJoinExec: mode=Partitioned, join_type=RightMark, on=[(col_a@0, c@2)] + RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1 + DataSourceExec: partitions=1, partition_sizes=[0] + RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet + "); + Ok(()) +} + #[tokio::test] async fn test_hash_join_interleaved_projection_preserves_parent_sort() -> Result<()> { let left_schema = create_test_schema()?; diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index 5c17ffbd1e7db..03a28e5f647fc 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -954,7 +954,7 @@ fn handle_custom_pushdown( } // For hash join we only maintain the input order for the right child -// for join type: Inner, Right, RightSemi, RightAnti +// for join type: Inner, Right, RightSemi, RightAnti, RightMark fn handle_hash_join( plan: &HashJoinExec, parent_required: OrderingRequirements, @@ -1045,7 +1045,7 @@ fn build_join_column_index(plan: &HashJoinExec) -> Vec { .chain(map_fields(plan.right().schema(), JoinSide::Right)) .collect::>() } - JoinType::RightSemi | JoinType::RightAnti => { + JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { map_fields(plan.right().schema(), JoinSide::Right) } _ => unreachable!("unexpected join type: {}", plan.join_type()),