diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index e01311e25be8b..d20ff3f4267b7 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -870,6 +870,222 @@ fn range_inner_hash_join_rehashes_incompatible_range_partitioning() -> Result<() Ok(()) } +#[test] +fn left_hash_join_uses_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Left); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + println!("{}", displayable(plan.as_ref()).indent(true)); + + Ok(()) +} + +#[test] +fn left_semi_hash_join_uses_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::LeftSemi); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + println!("{}", displayable(plan.as_ref()).indent(true)); + + Ok(()) +} + +#[test] +fn left_anti_hash_join_uses_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::LeftAnti); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + println!("{}", displayable(plan.as_ref()).indent(true)); + + Ok(()) +} + +#[test] +fn left_mark_hash_join_uses_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::LeftMark); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + println!("{}", displayable(plan.as_ref()).indent(true)); + + Ok(()) +} + +#[test] +fn left_hash_join_range_rehash_incompatible_split_points() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 15, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Left); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + println!("{}", displayable(plan.as_ref()).indent(true)); + + Ok(()) +} + +#[test] +fn left_hash_join_range_rehash_incompatible_partition_counts() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20], + SortOptions::default(), + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::LeftSemi); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + println!("{}", displayable(plan.as_ref()).indent(true)); + + Ok(()) +} + +#[test] +fn left_hash_join_range_rehash_incompatible_sort_options() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [30, 20, 10], + SortOptions::default().with_descending(true), + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::LeftAnti); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + println!("{}", displayable(plan.as_ref()).indent(true)); + + Ok(()) +} + +#[test] +fn left_hash_join_range_rehash_incompatible_key_expressions() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("b", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::LeftMark); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + println!("{}", displayable(plan.as_ref()).indent(true)); + + Ok(()) +} + #[test] fn multi_hash_joins() -> Result<()> { let left = parquet_exec(); diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index 5d004ca7fdfa5..9f597f28bf660 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -395,9 +395,10 @@ ORDER BY l.non_range_key, l.value, r.value; 2 350 350 ########## -# TEST 12: Non-Inner Range Join Repartitions -# Only inner partitioned hash joins opt in to Range satisfying KeyPartitioned -# requirements. Non-inner joins keep using Hash repartitioning. +# TEST 12: Left Range Join Repartitions +# Only inner and right-side (Right/RightSemi/RightAnti/RightMark) partitioned +# hash joins opt in to Range satisfying KeyPartitioned requirements. Other +# join types, such as Left, keep using Hash repartitioning. ########## query TT