From bff65d5cbf5e17863802d8cc032ed869847dd8b6 Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Sun, 20 Sep 2026 15:29:40 +0200 Subject: [PATCH 1/3] feat: preserve partitioning through co-partitioned Full joins --- datafusion/core/tests/dataframe/mod.rs | 23 +- .../enforce_distribution.rs | 184 +++++- .../physical-plan/src/joins/hash_join/exec.rs | 14 +- .../src/joins/sort_merge_join/exec.rs | 15 +- .../src/joins/symmetric_hash_join.rs | 14 +- datafusion/physical-plan/src/joins/utils.rs | 544 +++++++++++++++++- .../test_files/range_partitioning.slt | 211 +++++++ 7 files changed, 985 insertions(+), 20 deletions(-) diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 6457deeda98ae..1718c8fdbcd2f 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -85,7 +85,7 @@ use datafusion_expr::{ }; use datafusion_physical_expr::Partitioning; use datafusion_physical_expr::aggregate::AggregateExprBuilder; -use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::expressions::{Column, case, is_not_null}; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_physical_plan::aggregates::{ @@ -2778,9 +2778,24 @@ async fn verify_join_output_partitioning() -> Result<()> { ); } JoinType::Full => { - assert!(matches!( - out_partitioning, - &Partitioning::UnknownPartitioning(partition_count) if partition_count == default_partition_count)); + let coalesced_exprs = [("c1", "c2_c1"), ("c2", "c2_c2")] + .into_iter() + .map(|(left, right)| { + let left: Arc = + Arc::new(Column::new_with_schema(left, &join_schema)?); + let right: Arc = + Arc::new(Column::new_with_schema(right, &join_schema)?); + case( + None, + vec![(is_not_null(Arc::clone(&left))?, left)], + Some(right), + ) + }) + .collect::>>()?; + assert_eq!( + out_partitioning, + &Partitioning::Hash(coalesced_exprs, default_partition_count) + ); } } } diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 7c9d4984a6a68..fb310d5e988b4 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -52,7 +52,9 @@ use datafusion_datasource::memory::MemorySourceConfig; use datafusion_expr::{JoinType, Operator}; use datafusion_functions_aggregate::count::count_udaf; use datafusion_physical_expr::aggregate::AggregateExprBuilder; -use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal, binary, lit}; +use datafusion_physical_expr::expressions::{ + BinaryExpr, Column, Literal, binary, case, is_not_null, lit, +}; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::{ LexOrdering, OrderingRequirements, PhysicalSortExpr, @@ -79,6 +81,7 @@ use datafusion_physical_plan::execution_plan::ExecutionPlan; use datafusion_physical_plan::expressions::col; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::JoinOn; +use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::repartition::RepartitionExec; @@ -1414,6 +1417,185 @@ fn range_left_anti_hash_join_rehashes_incompatible_null_options() -> Result<()> Ok(()) } +/// Builds a Full `HashJoinExec` over two scans laid out by `partitioning`, with the right +/// side aliased to `a1`, `b1`. +fn co_partitioned_full_join( + partitioning: Partitioning, +) -> Result> { + let left = parquet_exec_with_output_partitioning(partitioning.clone()); + let right = projection_exec_with_alias( + parquet_exec_with_output_partitioning(partitioning), + vec![ + ("a".to_string(), "a1".to_string()), + ("b".to_string(), "b1".to_string()), + ], + ); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + )]; + Ok(hash_join_exec(left, right, &join_on, &JoinType::Full)) +} + +/// Builds `CASE WHEN first IS NOT NULL THEN first ELSE second END` over two columns of +/// `join`, which is the physical form `coalesce` takes. +fn coalesced_key( + join: &Arc, + first: &str, + second: &str, +) -> Result> { + let first = Arc::new(Column::new_with_schema(first, &join.schema())?) as _; + let second = Arc::new(Column::new_with_schema(second, &join.schema())?) as _; + case( + None, + vec![(is_not_null(Arc::clone(&first))?, first)], + Some(second), + ) +} + +/// Joins `join` on `key` against a scan laid out by `partitioning`, then enforces +/// distribution over the result. +fn plan_join_on_key( + join: Arc, + key: Arc, + partitioning: Partitioning, +) -> Result> { + let right = parquet_exec_with_output_partitioning(partitioning); + let on = vec![( + key, + Arc::new(Column::new_with_schema("c", &right.schema())?) as _, + )]; + let top = hash_join_exec(join, right, &on, &JoinType::Inner); + Ok(TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(top, &DISTRIB_DISTRIB_SORT)) +} + +#[test] +fn full_hash_join_keeps_hash_partitioning_on_coalesced_key() -> Result<()> { + let join = + co_partitioned_full_join(Partitioning::Hash(vec![col("a", &schema())?], 4))?; + let key = coalesced_key(&join, "a", "a1")?; + let plan = + plan_join_on_key(join, key, Partitioning::Hash(vec![col("c", &schema())?], 4))?; + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(CASE WHEN a@0 IS NOT NULL THEN a@0 ELSE a1@5 END, c@2)] + HashJoinExec: mode=Partitioned, join_type=Full, on=[(a@0, a1@0)] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Hash([a@0], 4), file_type=parquet + ProjectionExec: expr=[a@0 as a1, b@1 as b1] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Hash([a@0], 4), file_type=parquet + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Hash([c@2], 4), file_type=parquet + " + ); + Ok(()) +} + +#[test] +fn full_hash_join_keeps_range_partitioning_on_coalesced_key() -> Result<()> { + let join = co_partitioned_full_join(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?)?; + let key = coalesced_key(&join, "a", "a1")?; + let plan = plan_join_on_key( + join, + key, + range_partitioning("c", [10, 20, 30], SortOptions::default())?, + )?; + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(CASE WHEN a@0 IS NOT NULL THEN a@0 ELSE a1@5 END, c@2)] + HashJoinExec: mode=Partitioned, join_type=Full, on=[(a@0, a1@0)] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + ProjectionExec: expr=[a@0 as a1, b@1 as b1] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([c@2 ASC], [(10), (20), (30)], 4), file_type=parquet + " + ); + Ok(()) +} + +#[test] +fn full_hash_join_rehashes_plain_key_after_join() -> Result<()> { + let join = + co_partitioned_full_join(Partitioning::Hash(vec![col("a", &schema())?], 4))?; + let key = Arc::new(Column::new_with_schema("a", &join.schema())?) as _; + let plan = + plan_join_on_key(join, key, Partitioning::Hash(vec![col("c", &schema())?], 4))?; + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, c@2)] + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 + HashJoinExec: mode=Partitioned, join_type=Full, on=[(a@0, a1@0)] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Hash([a@0], 4), file_type=parquet + ProjectionExec: expr=[a@0 as a1, b@1 as b1] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Hash([a@0], 4), file_type=parquet + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Hash([c@2], 4), file_type=parquet + " + ); + Ok(()) +} + +#[test] +fn full_hash_join_keeps_partitioning_on_reversed_coalesced_key() -> Result<()> { + let join = + co_partitioned_full_join(Partitioning::Hash(vec![col("a", &schema())?], 4))?; + let key = coalesced_key(&join, "a1", "a")?; + let plan = + plan_join_on_key(join, key, Partitioning::Hash(vec![col("c", &schema())?], 4))?; + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(CASE WHEN a1@5 IS NOT NULL THEN a1@5 ELSE a@0 END, c@2)] + HashJoinExec: mode=Partitioned, join_type=Full, on=[(a@0, a1@0)] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Hash([a@0], 4), file_type=parquet + ProjectionExec: expr=[a@0 as a1, b@1 as b1] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Hash([a@0], 4), file_type=parquet + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Hash([c@2], 4), file_type=parquet + " + ); + Ok(()) +} + +#[test] +fn swapped_full_hash_join_keeps_partitioning_on_coalesced_key() -> Result<()> { + let join = + co_partitioned_full_join(Partitioning::Hash(vec![col("a", &schema())?], 4))?; + let swapped = join + .downcast_ref::() + .expect("hash_join_exec builds a HashJoinExec") + .swap_inputs(PartitionMode::Partitioned)?; + let key = coalesced_key(&swapped, "a", "a1")?; + let plan = plan_join_on_key( + swapped, + key, + Partitioning::Hash(vec![col("c", &schema())?], 4), + )?; + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(CASE WHEN a@0 IS NOT NULL THEN a@0 ELSE a1@5 END, c@2)] + ProjectionExec: expr=[a@2 as a, b@3 as b, c@4 as c, d@5 as d, e@6 as e, a1@0 as a1, b1@1 as b1] + HashJoinExec: mode=Partitioned, join_type=Full, on=[(a1@0, a@0)] + ProjectionExec: expr=[a@0 as a1, b@1 as b1] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Hash([a@0], 4), file_type=parquet + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Hash([a@0], 4), file_type=parquet + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Hash([c@2], 4), file_type=parquet + " + ); + Ok(()) +} + #[test] fn multi_hash_joins() -> Result<()> { let left = parquet_exec(); diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 7b9e701119ef4..8613b879ed50a 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -63,8 +63,9 @@ use crate::{ common::can_project, joins::utils::{ BuildProbeJoinMetrics, ColumnIndex, JoinFilter, JoinHashMapType, - build_join_schema, check_join_is_valid, estimate_join_statistics, - need_produce_result_in_final, symmetric_join_output_partitioning, + add_full_join_key_equivalences, build_join_schema, check_join_is_valid, + estimate_join_statistics, need_produce_result_in_final, + symmetric_join_output_partitioning, }, metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; @@ -1295,9 +1296,16 @@ impl HashJoinExec { right.output_partitioning().partition_count(), ), PartitionMode::Partitioned => { - symmetric_join_output_partitioning(left, right, &join_type)? + symmetric_join_output_partitioning(left, right, &join_type, on)? } }; + add_full_join_key_equivalences( + &mut eq_properties, + &output_partitioning, + join_type, + on, + left.schema().fields().len(), + )?; let emission_type = // LeftSemi does not emit rows during probing. It records matching build-side 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 b0433250c0c51..e0baa9736768f 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -28,9 +28,9 @@ use super::metrics::SortMergeJoinMetrics; use crate::execution_plan::{EmissionType, boundedness_from_children}; use crate::expressions::PhysicalSortExpr; use crate::joins::utils::{ - JoinFilter, JoinOn, JoinOnRef, build_join_schema, check_join_is_valid, - estimate_join_statistics, reorder_output_after_swap, swap_join_projection, - symmetric_join_output_partitioning, + JoinFilter, JoinOn, JoinOnRef, add_full_join_key_equivalences, build_join_schema, + check_join_is_valid, estimate_join_statistics, reorder_output_after_swap, + swap_join_projection, symmetric_join_output_partitioning, }; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet, SpillMetrics}; use crate::projection::{ @@ -327,7 +327,14 @@ impl SortMergeJoinExec { )?; let mut output_partitioning = - symmetric_join_output_partitioning(left, right, &join_type)?; + symmetric_join_output_partitioning(left, right, &join_type, join_on)?; + add_full_join_key_equivalences( + &mut eq_properties, + &output_partitioning, + join_type, + join_on, + left.schema().fields().len(), + )?; if let Some(projection) = projection { let mapping = ProjectionMapping::from_indices(projection, schema)?; diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 99a12796c688e..b1ca1d5065d3a 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -41,7 +41,8 @@ use crate::joins::stream_join_utils::{ }; use crate::joins::utils::{ BatchSplitter, BatchTransformer, ColumnIndex, JoinFilter, JoinHashMapType, JoinOn, - JoinOnRef, NoopBatchTransformer, StatefulStreamResult, apply_join_filter_to_indices, + JoinOnRef, NoopBatchTransformer, StatefulStreamResult, + add_full_join_key_equivalences, apply_join_filter_to_indices, build_batch_from_indices, build_join_schema, check_join_is_valid, equal_rows_arr, matchable_join_keys, symmetric_join_output_partitioning, update_hash, }; @@ -268,7 +269,7 @@ impl SymmetricHashJoinExec { join_on: JoinOnRef, ) -> Result { // Calculate equivalence properties: - let eq_properties = join_equivalence_properties( + let mut eq_properties = join_equivalence_properties( left.equivalence_properties().clone(), right.equivalence_properties().clone(), &join_type, @@ -280,7 +281,14 @@ impl SymmetricHashJoinExec { )?; let output_partitioning = - symmetric_join_output_partitioning(left, right, &join_type)?; + symmetric_join_output_partitioning(left, right, &join_type, join_on)?; + add_full_join_key_equivalences( + &mut eq_properties, + &output_partitioning, + join_type, + join_on, + left.schema().fields().len(), + )?; Ok(PlanProperties::new( eq_properties, diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 30b2972312f8c..dbc68a7714f89 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -26,6 +26,7 @@ use std::ops::Range; use std::sync::Arc; use std::task::{Context, Poll}; +use crate::distribution_requirements::InputDistributionRequirements; use crate::metrics::{ self, BaselineMetrics, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, MetricType, @@ -70,11 +71,11 @@ use datafusion_common::{ internal_datafusion_err, not_impl_err, plan_err, }; use datafusion_expr::interval_arithmetic::Interval; -use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::expressions::{Column, case, is_not_null}; use datafusion_physical_expr::utils::collect_columns; use datafusion_physical_expr::{ - LexOrdering, PhysicalExpr, PhysicalExprRef, add_offset_to_expr, - add_offset_to_physical_sort_exprs, + Distribution, EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalExprRef, + PhysicalSortExpr, add_offset_to_expr, add_offset_to_physical_sort_exprs, }; use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; @@ -1958,6 +1959,7 @@ pub(crate) fn symmetric_join_output_partitioning( left: &Arc, right: &Arc, join_type: &JoinType, + on: JoinOnRef, ) -> Result { let left_columns_len = left.schema().fields.len(); let left_partitioning = left.output_partitioning(); @@ -1973,13 +1975,114 @@ pub(crate) fn symmetric_join_output_partitioning( adjust_right_output_partitioning(right_partitioning, left_columns_len)? } JoinType::Full => { - // We could also use left partition count as they are necessarily equal. - Partitioning::UnknownPartitioning(right_partitioning.partition_count()) + full_join_output_partitioning(left, right, on, left_columns_len)? } }; Ok(result) } +/// Output partitioning of a Full join whose inputs are co-partitioned on the join keys. +fn full_join_output_partitioning( + left: &Arc, + right: &Arc, + on: JoinOnRef, + left_columns_len: usize, +) -> Result { + let left_partitioning = left.output_partitioning(); + let unknown = + Partitioning::UnknownPartitioning(right.output_partitioning().partition_count()); + if on.is_empty() || left_partitioning.partition_count() < 2 { + return Ok(unknown); + } + let (left_keys, right_keys) = on + .iter() + .map(|(l, r)| (Arc::clone(l), Arc::clone(r))) + .unzip(); + let requirements = InputDistributionRequirements::co_partitioned(vec![ + Distribution::KeyPartitioned(left_keys), + Distribution::KeyPartitioned(right_keys), + ]); + let children = [left.as_ref(), right.as_ref()]; + if !requirements + .unsatisfied_co_partitioned_children("Full join", &children)? + .is_empty() + { + return Ok(unknown); + } + // Either side of a row may be null, so each key becomes the physical form of + // `coalesce(left, right)`. + let coalesced = on + .iter() + .map(|(l, r)| { + let r = add_offset_to_expr(Arc::clone(r), left_columns_len as _)?; + coalesce_keys(Arc::clone(l), r) + }) + .collect::>>()?; + let result = match left_partitioning { + Partitioning::Hash(_, count) => Partitioning::Hash(coalesced, *count), + Partitioning::Range(range) => { + let sort_exprs = coalesced + .into_iter() + .zip(range.ordering().iter()) + .map(|(expr, sort_expr)| PhysicalSortExpr::new(expr, sort_expr.options)); + let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { + internal_datafusion_err!( + "Coalescing join keys produced an empty ordering" + ) + })?; + if ordering.len() == range.ordering().len() { + Partitioning::Range(RangePartitioning::new( + ordering, + range.split_points().to_vec(), + )) + } else { + // Duplicate keys collapse in the ordering, leaving the split points + // misaligned with the remaining keys. + unknown + } + } + _ => unknown, + }; + Ok(result) +} + +/// Builds `CASE WHEN first IS NOT NULL THEN first ELSE second END`. +fn coalesce_keys( + first: PhysicalExprRef, + second: PhysicalExprRef, +) -> Result { + case( + None, + vec![(is_not_null(Arc::clone(&first))?, first)], + Some(second), + ) +} + +/// Records that the two coalesce orders of each Full join key pair are equal, which lets +/// a parent keyed on either order match the partitioning the join reports. Does nothing +/// unless the join is a Full join that kept a key based partitioning. +pub(crate) fn add_full_join_key_equivalences( + eq_properties: &mut EquivalenceProperties, + output_partitioning: &Partitioning, + join_type: JoinType, + on: JoinOnRef, + left_columns_len: usize, +) -> Result<()> { + if join_type != JoinType::Full + || matches!(output_partitioning, Partitioning::UnknownPartitioning(_)) + { + return Ok(()); + } + for (l, r) in on { + let r = add_offset_to_expr(Arc::clone(r), left_columns_len as _)?; + eq_properties.add_equal_conditions( + coalesce_keys(Arc::clone(l), Arc::clone(&r))?, + coalesce_keys(r, Arc::clone(l))?, + )?; + } + Ok(()) +} + /// Convert a boolean filter array into a unified mask bitmap. /// /// Caution: The filter result is NOT a bitmap; it contains true/false/null values. @@ -2648,13 +2751,20 @@ mod tests { use std::time::Duration; use super::*; + use crate::joins::{ + HashJoinExec, PartitionMode, SortMergeJoinExec, StreamJoinPartitionMode, + SymmetricHashJoinExec, + }; use crate::metrics::MetricValue; + use crate::repartition::RepartitionExec; + use crate::test::TestMemoryExec; use arrow::datatypes::{DataType, Fields}; use arrow::error::{ArrowError, Result as ArrowResult}; use datafusion_common::stats::Precision::{Absent, Exact, Inexact}; use datafusion_common::{ScalarValue, SplitPoint, arrow_datafusion_err, arrow_err}; use datafusion_physical_expr::PhysicalSortExpr; + use datafusion_physical_expr::expressions::UnKnownColumn; use rstest::rstest; @@ -4445,6 +4555,430 @@ mod tests { Ok(()) } + fn int32_schema(first: &str, second: &str) -> Arc { + Arc::new(Schema::new(vec![ + Field::new(first, DataType::Int32, true), + Field::new(second, DataType::Int32, true), + ])) + } + + fn unknown_partitioned( + schema: Arc, + partitions: usize, + ) -> Result> { + Ok(TestMemoryExec::try_new_exec( + &vec![vec![]; partitions], + schema, + None, + )?) + } + + fn repartitioned( + schema: Arc, + partitioning: Partitioning, + ) -> Result> { + let source = unknown_partitioned(schema, 1)?; + Ok(Arc::new(RepartitionExec::try_new(source, partitioning)?)) + } + + fn col_at(name: &str, index: usize) -> PhysicalExprRef { + Arc::new(Column::new(name, index)) + } + + fn range_on( + expr: PhysicalExprRef, + options: SortOptions, + split_points: &[i32], + ) -> Result { + let ordering = + LexOrdering::new([PhysicalSortExpr::new(expr, options)]).expect("one key"); + let split_points = split_points + .iter() + .map(|value| SplitPoint::new(vec![ScalarValue::Int32(Some(*value))])) + .collect(); + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points, + )?)) + } + + fn a_equals_c() -> JoinOn { + vec![(col_at("a", 0), col_at("c", 0))] + } + + fn full_join_partitioning( + left: &Arc, + right: &Arc, + on: JoinOnRef, + ) -> Result { + symmetric_join_output_partitioning(left, right, &JoinType::Full, on) + } + + #[test] + fn full_join_output_partitioning_coalesces_hash_keys() -> Result<()> { + let left = repartitioned( + int32_schema("a", "b"), + Partitioning::Hash(vec![col_at("a", 0)], 4), + )?; + let right = repartitioned( + int32_schema("c", "d"), + Partitioning::Hash(vec![col_at("c", 0)], 4), + )?; + let on = a_equals_c(); + + assert_eq!( + full_join_partitioning(&left, &right, &on)?, + Partitioning::Hash(vec![coalesce_keys(col_at("a", 0), col_at("c", 2))?], 4) + ); + assert_eq!( + symmetric_join_output_partitioning(&left, &right, &JoinType::Inner, &on)?, + Partitioning::Hash(vec![col_at("c", 2)], 4) + ); + assert_eq!( + symmetric_join_output_partitioning(&left, &right, &JoinType::Left, &on)?, + Partitioning::Hash(vec![col_at("a", 0)], 4) + ); + Ok(()) + } + + #[test] + fn full_join_output_partitioning_coalesces_range_keys() -> Result<()> { + let options = SortOptions::new(false, true); + let left = repartitioned( + int32_schema("a", "b"), + range_on(col_at("a", 0), options, &[10, 20, 30])?, + )?; + let right = repartitioned( + int32_schema("c", "d"), + range_on(col_at("c", 0), options, &[10, 20, 30])?, + )?; + + let expected = range_on( + coalesce_keys(col_at("a", 0), col_at("c", 2))?, + options, + &[10, 20, 30], + )?; + assert_eq!( + full_join_partitioning(&left, &right, &a_equals_c())?, + expected + ); + Ok(()) + } + + #[test] + fn full_join_output_partitioning_coalesces_every_key_pair() -> Result<()> { + let left = repartitioned( + int32_schema("a", "b"), + Partitioning::Hash(vec![col_at("a", 0), col_at("b", 1)], 4), + )?; + let right = repartitioned( + int32_schema("c", "d"), + Partitioning::Hash(vec![col_at("c", 0), col_at("d", 1)], 4), + )?; + let on = vec![ + (col_at("a", 0), col_at("c", 0)), + (col_at("b", 1), col_at("d", 1)), + ]; + + let expected = Partitioning::Hash( + vec![ + coalesce_keys(col_at("a", 0), col_at("c", 2))?, + coalesce_keys(col_at("b", 1), col_at("d", 3))?, + ], + 4, + ); + assert_eq!(full_join_partitioning(&left, &right, &on)?, expected); + Ok(()) + } + + #[test] + fn full_join_output_partitioning_is_unknown_for_different_partition_counts() + -> Result<()> { + let left = repartitioned( + int32_schema("a", "b"), + Partitioning::Hash(vec![col_at("a", 0)], 4), + )?; + let right = repartitioned( + int32_schema("c", "d"), + Partitioning::Hash(vec![col_at("c", 0)], 2), + )?; + + assert!(matches!( + full_join_partitioning(&left, &right, &a_equals_c())?, + Partitioning::UnknownPartitioning(2) + )); + Ok(()) + } + + #[test] + fn full_join_output_partitioning_is_unknown_for_different_range_split_points() + -> Result<()> { + let options = SortOptions::new(false, true); + let left = repartitioned( + int32_schema("a", "b"), + range_on(col_at("a", 0), options, &[10, 20, 30])?, + )?; + let right = repartitioned( + int32_schema("c", "d"), + range_on(col_at("c", 0), options, &[10, 20, 40])?, + )?; + + assert!(matches!( + full_join_partitioning(&left, &right, &a_equals_c())?, + Partitioning::UnknownPartitioning(4) + )); + Ok(()) + } + + #[test] + fn full_join_output_partitioning_is_unknown_for_different_range_sort_options() + -> Result<()> { + let left = repartitioned( + int32_schema("a", "b"), + range_on(col_at("a", 0), SortOptions::new(false, true), &[10, 20, 30])?, + )?; + let right = repartitioned( + int32_schema("c", "d"), + range_on( + col_at("c", 0), + SortOptions::new(false, false), + &[10, 20, 30], + )?, + )?; + + assert!(matches!( + full_join_partitioning(&left, &right, &a_equals_c())?, + Partitioning::UnknownPartitioning(4) + )); + Ok(()) + } + + #[test] + fn full_join_output_partitioning_is_unknown_for_hash_left_and_range_right() + -> Result<()> { + let left = repartitioned( + int32_schema("a", "b"), + Partitioning::Hash(vec![col_at("a", 0)], 4), + )?; + let right = repartitioned( + int32_schema("c", "d"), + range_on(col_at("c", 0), SortOptions::new(false, true), &[10, 20, 30])?, + )?; + + assert!(matches!( + full_join_partitioning(&left, &right, &a_equals_c())?, + Partitioning::UnknownPartitioning(4) + )); + Ok(()) + } + + #[test] + fn full_join_output_partitioning_is_unknown_for_unknown_input_partitioning() + -> Result<()> { + let left = unknown_partitioned(int32_schema("a", "b"), 4)?; + let right = repartitioned( + int32_schema("c", "d"), + Partitioning::Hash(vec![col_at("c", 0)], 4), + )?; + + assert!(matches!( + full_join_partitioning(&left, &right, &a_equals_c())?, + Partitioning::UnknownPartitioning(4) + )); + Ok(()) + } + + #[test] + fn full_join_output_partitioning_is_unknown_for_single_partition_inputs() -> Result<()> + { + let left = repartitioned( + int32_schema("a", "b"), + Partitioning::Hash(vec![col_at("a", 0)], 1), + )?; + let right = repartitioned( + int32_schema("c", "d"), + Partitioning::Hash(vec![col_at("c", 0)], 1), + )?; + + assert!(matches!( + full_join_partitioning(&left, &right, &a_equals_c())?, + Partitioning::UnknownPartitioning(1) + )); + Ok(()) + } + + #[test] + fn full_join_output_partitioning_is_unknown_when_keys_are_not_the_join_keys() + -> Result<()> { + let left = repartitioned( + int32_schema("a", "b"), + Partitioning::Hash(vec![col_at("b", 1)], 4), + )?; + let right = repartitioned( + int32_schema("c", "d"), + Partitioning::Hash(vec![col_at("c", 0)], 4), + )?; + + assert!(matches!( + full_join_partitioning(&left, &right, &a_equals_c())?, + Partitioning::UnknownPartitioning(4) + )); + Ok(()) + } + + #[test] + fn full_join_output_partitioning_is_unknown_when_right_keys_are_not_the_join_keys() + -> Result<()> { + let left = repartitioned( + int32_schema("a", "b"), + Partitioning::Hash(vec![col_at("a", 0)], 4), + )?; + let right = repartitioned( + int32_schema("c", "d"), + Partitioning::Hash(vec![col_at("d", 1)], 4), + )?; + + assert!(matches!( + full_join_partitioning(&left, &right, &a_equals_c())?, + Partitioning::UnknownPartitioning(4) + )); + Ok(()) + } + + #[test] + fn full_join_output_partitioning_is_unknown_without_join_keys() -> Result<()> { + let left = repartitioned( + int32_schema("a", "b"), + Partitioning::Hash(vec![col_at("a", 0)], 4), + )?; + let right = repartitioned( + int32_schema("c", "d"), + Partitioning::Hash(vec![col_at("c", 0)], 4), + )?; + + assert!(matches!( + full_join_partitioning(&left, &right, &[])?, + Partitioning::UnknownPartitioning(4) + )); + Ok(()) + } + + #[test] + fn full_sort_merge_join_reports_coalesced_partitioning() -> Result<()> { + let left = repartitioned( + int32_schema("a", "b"), + Partitioning::Hash(vec![col_at("a", 0)], 4), + )?; + let right = repartitioned( + int32_schema("c", "d"), + Partitioning::Hash(vec![col_at("c", 0)], 4), + )?; + let join = SortMergeJoinExec::try_new( + left, + right, + a_equals_c(), + None, + JoinType::Full, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + )?; + let expected = + Partitioning::Hash(vec![coalesce_keys(col_at("a", 0), col_at("c", 2))?], 4); + assert_eq!(join.properties().output_partitioning(), &expected); + Ok(()) + } + + #[test] + fn full_symmetric_hash_join_reports_coalesced_partitioning() -> Result<()> { + let left = repartitioned( + int32_schema("a", "b"), + Partitioning::Hash(vec![col_at("a", 0)], 4), + )?; + let right = repartitioned( + int32_schema("c", "d"), + Partitioning::Hash(vec![col_at("c", 0)], 4), + )?; + let join = SymmetricHashJoinExec::try_new( + left, + right, + a_equals_c(), + None, + &JoinType::Full, + NullEquality::NullEqualsNothing, + None, + None, + StreamJoinPartitionMode::Partitioned, + )?; + let expected = + Partitioning::Hash(vec![coalesce_keys(col_at("a", 0), col_at("c", 2))?], 4); + assert_eq!(join.properties().output_partitioning(), &expected); + Ok(()) + } + + fn projected_full_hash_join(projection: Vec) -> Result { + let left = repartitioned( + int32_schema("a", "b"), + Partitioning::Hash(vec![col_at("a", 0)], 4), + )?; + let right = repartitioned( + int32_schema("c", "d"), + Partitioning::Hash(vec![col_at("c", 0)], 4), + )?; + HashJoinExec::try_new( + left, + right, + a_equals_c(), + None, + &JoinType::Full, + Some(projection), + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + } + + #[test] + fn full_hash_join_keys_coalesce_equally_in_either_order() -> Result<()> { + let join = projected_full_hash_join(vec![0, 1, 2, 3])?; + let group = join.properties().equivalence_properties().eq_group(); + let left_first = coalesce_keys(col_at("a", 0), col_at("c", 2))?; + let right_first = coalesce_keys(col_at("c", 2), col_at("a", 0))?; + assert!( + group + .normalize_expr(left_first) + .eq(&group.normalize_expr(right_first)) + ); + Ok(()) + } + + #[test] + fn full_hash_join_projection_dropping_right_key_loses_the_key() -> Result<()> { + let join = projected_full_hash_join(vec![0, 1])?; + + let Partitioning::Hash(exprs, count) = join.properties().output_partitioning() + else { + panic!( + "expected hash partitioning, got {:?}", + join.properties().output_partitioning() + ); + }; + assert_eq!(*count, 4); + assert_eq!(exprs.len(), 1); + assert!(exprs[0].downcast_ref::().is_some()); + Ok(()) + } + + #[test] + fn full_hash_join_projection_remaps_coalesced_keys() -> Result<()> { + let join = projected_full_hash_join(vec![0, 2])?; + + assert_eq!( + join.properties().output_partitioning(), + &Partitioning::Hash(vec![coalesce_keys(col_at("a", 0), col_at("c", 1))?], 4) + ); + Ok(()) + } + #[test] fn test_calculate_join_output_ordering() -> Result<()> { let left_ordering = LexOrdering::new(vec![ diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index ec374b3d62a28..7c45bae15b689 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -1893,6 +1893,217 @@ ORDER BY b.range_key; 5 50 50 20 200 200 +########## +# TEST 49: Full Join Output Partitioning Feeds Aggregate on Coalesced Key +# Co-partitioned Full join inputs keep every row in the partition of its key, so +# an aggregate on COALESCE(l.range_key, r.range_key) needs no repartitioning. +########## + +statement ok +set datafusion.optimizer.enable_join_dynamic_filter_pushdown = false; + +statement ok +set datafusion.execution.parquet.pushdown_filters = false; + +statement ok +set datafusion.optimizer.prefer_hash_join = true; + +statement ok +set datafusion.optimizer.repartition_joins = true; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +statement ok +set datafusion.execution.target_partitions = 4; + +query TT +EXPLAIN SELECT COALESCE(l.range_key, r.range_key) AS k, count(*), sum(l.value) +FROM range_partitioned l +FULL JOIN range_partitioned_sparse r ON l.range_key = r.range_key +GROUP BY COALESCE(l.range_key, r.range_key); +---- +physical_plan +01)ProjectionExec: expr=[coalesce(l.range_key,r.range_key)@0 as k, count(Int64(1))@1 as count(*), sum(l.value)@2 as sum(l.value)] +02)--AggregateExec: mode=SinglePartitioned, gby=[CASE WHEN range_key@0 IS NOT NULL THEN range_key@0 ELSE range_key@2 END as coalesce(l.range_key,r.range_key)], aggr=[count(Int64(1)), sum(l.value)] +03)----HashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +05)------DataSourceExec: file_groups=, projection=[range_key], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query III +SELECT COALESCE(l.range_key, r.range_key) AS k, count(*), sum(l.value) +FROM range_partitioned l +FULL JOIN range_partitioned_sparse r ON l.range_key = r.range_key +GROUP BY COALESCE(l.range_key, r.range_key) +ORDER BY k; +---- +1 1 10 +5 1 50 +8 1 NULL +10 1 100 +15 1 150 +20 1 200 +25 1 250 +30 1 300 +35 1 350 +40 1 NULL + +########## +# TEST 50: Full Join Output Partitioning Feeds a Further Join on Coalesced Key +# The Range partitioning on the coalesced key flows through the projection the +# planner inserts, so the outer Full join needs no repartitioning on either side. +########## + +query TT +EXPLAIN SELECT l.range_key, r.range_key, t.range_key, t.value +FROM range_partitioned l +FULL JOIN range_partitioned r ON l.range_key = r.range_key +FULL JOIN range_partitioned_sparse t ON COALESCE(l.range_key, r.range_key) = t.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Full, on=[(CASE WHEN l.range_key IS NOT NULL THEN l.range_key ELSE r.range_key END@2, range_key@0)], projection=[range_key@0, range_key@1, range_key@3, value@4] +02)--ProjectionExec: expr=[range_key@0 as range_key, range_key@1 as range_key, CASE WHEN range_key@0 IS NOT NULL THEN range_key@0 ELSE range_key@1 END as CASE WHEN l.range_key IS NOT NULL THEN l.range_key ELSE r.range_key END] +03)----HashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)] +04)------DataSourceExec: file_groups=, projection=[range_key], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +05)------DataSourceExec: file_groups=, projection=[range_key], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +06)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query IIII +SELECT l.range_key, r.range_key, t.range_key, t.value +FROM range_partitioned l +FULL JOIN range_partitioned r ON l.range_key = r.range_key +FULL JOIN range_partitioned_sparse t ON COALESCE(l.range_key, r.range_key) = t.range_key +ORDER BY l.range_key, t.range_key; +---- +1 1 NULL NULL +5 5 5 50 +10 10 10 100 +15 15 NULL NULL +20 20 20 200 +25 25 NULL NULL +30 30 30 300 +35 35 NULL NULL +NULL NULL 8 80 +NULL NULL 40 400 + +########## +# TEST 51: Hash Repartitioned Full Join Feeds Aggregate on Coalesced Key +# Differing split points force Hash repartitioning below the join, and the +# resulting Hash partitioning on the coalesced key satisfies the aggregate above. +########## + +query TT +EXPLAIN SELECT COALESCE(l.range_key, r.range_key) AS k, count(*), sum(l.value) +FROM range_partitioned l +FULL JOIN range_partitioned_shifted r ON l.range_key = r.range_key +GROUP BY COALESCE(l.range_key, r.range_key); +---- +physical_plan +01)ProjectionExec: expr=[coalesce(l.range_key,r.range_key)@0 as k, count(Int64(1))@1 as count(*), sum(l.value)@2 as sum(l.value)] +02)--AggregateExec: mode=SinglePartitioned, gby=[CASE WHEN range_key@0 IS NOT NULL THEN range_key@0 ELSE range_key@2 END as coalesce(l.range_key,r.range_key)], aggr=[count(Int64(1)), sum(l.value)] +03)----HashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)] +04)------RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +06)------RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +07)--------DataSourceExec: file_groups=, projection=[range_key], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=parquet + +query III +SELECT COALESCE(l.range_key, r.range_key) AS k, count(*), sum(l.value) +FROM range_partitioned l +FULL JOIN range_partitioned_shifted r ON l.range_key = r.range_key +GROUP BY COALESCE(l.range_key, r.range_key) +ORDER BY k; +---- +1 1 10 +5 1 50 +10 1 100 +15 1 150 +20 1 200 +25 1 250 +30 1 300 +35 1 350 + +########## +# TEST 52: Full Join Output Partitioning Does Not Cover a Single Side Key +# Right only rows carry NULL in l.range_key across every partition, so an +# aggregate on l.range_key alone still needs Hash repartitioning above the join. +########## + +query TT +EXPLAIN SELECT l.range_key, count(*) +FROM range_partitioned l +FULL JOIN range_partitioned_sparse r ON l.range_key = r.range_key +GROUP BY l.range_key; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, count(Int64(1))@1 as count(*)] +02)--AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key], aggr=[count(Int64(1))] +03)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +04)------AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[count(Int64(1))] +05)--------HashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)], projection=[range_key@0] +06)----------DataSourceExec: file_groups=, projection=[range_key], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +07)----------DataSourceExec: file_groups=, projection=[range_key], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT l.range_key, count(*) +FROM range_partitioned l +FULL JOIN range_partitioned_sparse r ON l.range_key = r.range_key +GROUP BY l.range_key +ORDER BY l.range_key; +---- +1 1 +5 1 +10 1 +15 1 +20 1 +25 1 +30 1 +35 1 +NULL 2 + +########## +# TEST 53: Full Join Output Partitioning Matches Coalesce in Either Key Order +# Both coalesce orders are equal on every Full join row, so grouping by +# COALESCE(r.range_key, l.range_key) needs no repartitioning either. +########## + +query TT +EXPLAIN SELECT COALESCE(r.range_key, l.range_key) AS k, count(*) +FROM range_partitioned l +FULL JOIN range_partitioned_sparse r ON l.range_key = r.range_key +GROUP BY COALESCE(r.range_key, l.range_key); +---- +physical_plan +01)ProjectionExec: expr=[coalesce(r.range_key,l.range_key)@0 as k, count(Int64(1))@1 as count(*)] +02)--AggregateExec: mode=SinglePartitioned, gby=[CASE WHEN range_key@1 IS NOT NULL THEN range_key@1 ELSE range_key@0 END as coalesce(r.range_key,l.range_key)], aggr=[count(Int64(1))] +03)----HashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)] +04)------DataSourceExec: file_groups=, projection=[range_key], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +05)------DataSourceExec: file_groups=, projection=[range_key], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT COALESCE(r.range_key, l.range_key) AS k, count(*) +FROM range_partitioned l +FULL JOIN range_partitioned_sparse r ON l.range_key = r.range_key +GROUP BY COALESCE(r.range_key, l.range_key) +ORDER BY k; +---- +1 1 +5 1 +8 1 +10 1 +15 1 +20 1 +25 1 +30 1 +35 1 +40 1 + +statement ok +reset datafusion.optimizer.repartition_joins; + +statement ok +reset datafusion.optimizer.prefer_hash_join; + statement ok reset datafusion.optimizer.preserve_file_partitions; From bc652ae470b8f99531a4b75040f210b986d836ad Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Mon, 21 Sep 2026 00:14:04 +0200 Subject: [PATCH 2/3] refactor: reuse Partitioning::adapt for Full join output partitioning. --- .../enforce_distribution.rs | 11 +- .../physical-plan/src/joins/hash_join/exec.rs | 22 +- .../src/joins/sort_merge_join/exec.rs | 18 +- .../src/joins/symmetric_hash_join.rs | 15 +- datafusion/physical-plan/src/joins/utils.rs | 213 ++++-------------- .../test_files/range_partitioning.slt | 15 +- 6 files changed, 81 insertions(+), 213 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index fb310d5e988b4..9cfaeddb34667 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -1417,8 +1417,8 @@ fn range_left_anti_hash_join_rehashes_incompatible_null_options() -> Result<()> Ok(()) } -/// Builds a Full `HashJoinExec` over two scans laid out by `partitioning`, with the right -/// side aliased to `a1`, `b1`. +/// Builds a Full `HashJoinExec` over two scans laid out by `partitioning`. The right side +/// columns are aliased to `a1` and `b1`. fn co_partitioned_full_join( partitioning: Partitioning, ) -> Result> { @@ -1437,8 +1437,7 @@ fn co_partitioned_full_join( Ok(hash_join_exec(left, right, &join_on, &JoinType::Full)) } -/// Builds `CASE WHEN first IS NOT NULL THEN first ELSE second END` over two columns of -/// `join`, which is the physical form `coalesce` takes. +/// Builds `coalesce(first, second)` over two columns of `join` in its physical CASE form. fn coalesced_key( join: &Arc, first: &str, @@ -1453,8 +1452,8 @@ fn coalesced_key( ) } -/// Joins `join` on `key` against a scan laid out by `partitioning`, then enforces -/// distribution over the result. +/// Inner joins `join` on `key` with a scan laid out by `partitioning`, then runs the +/// distribution rule over the result. fn plan_join_on_key( join: Arc, key: Arc, diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 8613b879ed50a..5587e72d31412 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -63,9 +63,8 @@ use crate::{ common::can_project, joins::utils::{ BuildProbeJoinMetrics, ColumnIndex, JoinFilter, JoinHashMapType, - add_full_join_key_equivalences, build_join_schema, check_join_is_valid, - estimate_join_statistics, need_produce_result_in_final, - symmetric_join_output_partitioning, + build_join_schema, check_join_is_valid, estimate_join_statistics, + need_produce_result_in_final, symmetric_join_output_partitioning, }, metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; @@ -1295,17 +1294,14 @@ impl HashJoinExec { PartitionMode::Auto => Partitioning::UnknownPartitioning( right.output_partitioning().partition_count(), ), - PartitionMode::Partitioned => { - symmetric_join_output_partitioning(left, right, &join_type, on)? - } + PartitionMode::Partitioned => symmetric_join_output_partitioning( + left, + right, + &join_type, + on, + &mut eq_properties, + )?, }; - add_full_join_key_equivalences( - &mut eq_properties, - &output_partitioning, - join_type, - on, - left.schema().fields().len(), - )?; let emission_type = // LeftSemi does not emit rows during probing. It records matching build-side 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 e0baa9736768f..3803a66e1d637 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -28,9 +28,9 @@ use super::metrics::SortMergeJoinMetrics; use crate::execution_plan::{EmissionType, boundedness_from_children}; use crate::expressions::PhysicalSortExpr; use crate::joins::utils::{ - JoinFilter, JoinOn, JoinOnRef, add_full_join_key_equivalences, build_join_schema, - check_join_is_valid, estimate_join_statistics, reorder_output_after_swap, - swap_join_projection, symmetric_join_output_partitioning, + JoinFilter, JoinOn, JoinOnRef, build_join_schema, check_join_is_valid, + estimate_join_statistics, reorder_output_after_swap, swap_join_projection, + symmetric_join_output_partitioning, }; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet, SpillMetrics}; use crate::projection::{ @@ -326,14 +326,12 @@ impl SortMergeJoinExec { join_on, )?; - let mut output_partitioning = - symmetric_join_output_partitioning(left, right, &join_type, join_on)?; - add_full_join_key_equivalences( - &mut eq_properties, - &output_partitioning, - join_type, + let mut output_partitioning = symmetric_join_output_partitioning( + left, + right, + &join_type, join_on, - left.schema().fields().len(), + &mut eq_properties, )?; if let Some(projection) = projection { diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index b1ca1d5065d3a..5a718abdaa72d 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -41,8 +41,7 @@ use crate::joins::stream_join_utils::{ }; use crate::joins::utils::{ BatchSplitter, BatchTransformer, ColumnIndex, JoinFilter, JoinHashMapType, JoinOn, - JoinOnRef, NoopBatchTransformer, StatefulStreamResult, - add_full_join_key_equivalences, apply_join_filter_to_indices, + JoinOnRef, NoopBatchTransformer, StatefulStreamResult, apply_join_filter_to_indices, build_batch_from_indices, build_join_schema, check_join_is_valid, equal_rows_arr, matchable_join_keys, symmetric_join_output_partitioning, update_hash, }; @@ -280,14 +279,12 @@ impl SymmetricHashJoinExec { join_on, )?; - let output_partitioning = - symmetric_join_output_partitioning(left, right, &join_type, join_on)?; - add_full_join_key_equivalences( - &mut eq_properties, - &output_partitioning, - join_type, + let output_partitioning = symmetric_join_output_partitioning( + left, + right, + &join_type, join_on, - left.schema().fields().len(), + &mut eq_properties, )?; Ok(PlanProperties::new( diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index dbc68a7714f89..8ab23868cc5e3 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -75,7 +75,7 @@ use datafusion_physical_expr::expressions::{Column, case, is_not_null}; use datafusion_physical_expr::utils::collect_columns; use datafusion_physical_expr::{ Distribution, EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalExprRef, - PhysicalSortExpr, add_offset_to_expr, add_offset_to_physical_sort_exprs, + add_offset_to_expr, add_offset_to_physical_sort_exprs, }; use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; @@ -1955,11 +1955,15 @@ pub enum StatefulStreamResult { Continue, } +/// Output partitioning of a partitioned join. A Full join over co-partitioned inputs +/// keeps its partitioning on the coalesced keys and records in `eq_properties` that both +/// coalesce orders of each key pair are equal, so a parent keyed on either order matches. pub(crate) fn symmetric_join_output_partitioning( left: &Arc, right: &Arc, join_type: &JoinType, on: JoinOnRef, + eq_properties: &mut EquivalenceProperties, ) -> Result { let left_columns_len = left.schema().fields.len(); let left_partitioning = left.output_partitioning(); @@ -1974,19 +1978,23 @@ pub(crate) fn symmetric_join_output_partitioning( JoinType::Inner | JoinType::Right => { adjust_right_output_partitioning(right_partitioning, left_columns_len)? } - JoinType::Full => { - full_join_output_partitioning(left, right, on, left_columns_len)? - } + JoinType::Full => full_join_output_partitioning( + left, + right, + on, + left_columns_len, + eq_properties, + )?, }; Ok(result) } -/// Output partitioning of a Full join whose inputs are co-partitioned on the join keys. fn full_join_output_partitioning( left: &Arc, right: &Arc, on: JoinOnRef, left_columns_len: usize, + eq_properties: &mut EquivalenceProperties, ) -> Result { let left_partitioning = left.output_partitioning(); let unknown = @@ -2009,44 +2017,23 @@ fn full_join_output_partitioning( { return Ok(unknown); } - // Either side of a row may be null, so each key becomes the physical form of - // `coalesce(left, right)`. - let coalesced = on - .iter() - .map(|(l, r)| { - let r = add_offset_to_expr(Arc::clone(r), left_columns_len as _)?; - coalesce_keys(Arc::clone(l), r) - }) - .collect::>>()?; - let result = match left_partitioning { - Partitioning::Hash(_, count) => Partitioning::Hash(coalesced, *count), - Partitioning::Range(range) => { - let sort_exprs = coalesced - .into_iter() - .zip(range.ordering().iter()) - .map(|(expr, sort_expr)| PhysicalSortExpr::new(expr, sort_expr.options)); - let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { - internal_datafusion_err!( - "Coalescing join keys produced an empty ordering" - ) - })?; - if ordering.len() == range.ordering().len() { - Partitioning::Range(RangePartitioning::new( - ordering, - range.split_points().to_vec(), - )) - } else { - // Duplicate keys collapse in the ordering, leaving the split points - // misaligned with the remaining keys. - unknown - } - } - _ => unknown, - }; - Ok(result) + // Either side of a row may be null, so each key becomes `coalesce(left, right)`. + let mut coalesced = Vec::with_capacity(on.len()); + for (l, r) in on { + let r = add_offset_to_expr(Arc::clone(r), left_columns_len as _)?; + let left_first = coalesce_keys(Arc::clone(l), Arc::clone(&r))?; + let right_first = coalesce_keys(r, Arc::clone(l))?; + eq_properties.add_equal_conditions(Arc::clone(&left_first), right_first)?; + coalesced.push(left_first); + } + let keys = Distribution::KeyPartitioned(coalesced); + Ok(left_partitioning + .adapt(&keys, eq_properties.schema()) + .unwrap_or(unknown)) } -/// Builds `CASE WHEN first IS NOT NULL THEN first ELSE second END`. +/// Builds `CASE WHEN first IS NOT NULL THEN first ELSE second END`, the physical form the +/// planner gives `coalesce(first, second)`. fn coalesce_keys( first: PhysicalExprRef, second: PhysicalExprRef, @@ -2058,31 +2045,6 @@ fn coalesce_keys( ) } -/// Records that the two coalesce orders of each Full join key pair are equal, which lets -/// a parent keyed on either order match the partitioning the join reports. Does nothing -/// unless the join is a Full join that kept a key based partitioning. -pub(crate) fn add_full_join_key_equivalences( - eq_properties: &mut EquivalenceProperties, - output_partitioning: &Partitioning, - join_type: JoinType, - on: JoinOnRef, - left_columns_len: usize, -) -> Result<()> { - if join_type != JoinType::Full - || matches!(output_partitioning, Partitioning::UnknownPartitioning(_)) - { - return Ok(()); - } - for (l, r) in on { - let r = add_offset_to_expr(Arc::clone(r), left_columns_len as _)?; - eq_properties.add_equal_conditions( - coalesce_keys(Arc::clone(l), Arc::clone(&r))?, - coalesce_keys(r, Arc::clone(l))?, - )?; - } - Ok(()) -} - /// Convert a boolean filter array into a unified mask bitmap. /// /// Caution: The filter result is NOT a bitmap; it contains true/false/null values. @@ -4606,12 +4568,29 @@ mod tests { vec![(col_at("a", 0), col_at("c", 0))] } + fn join_partitioning( + left: &Arc, + right: &Arc, + join_type: JoinType, + on: JoinOnRef, + ) -> Result { + let (schema, _) = build_join_schema(&left.schema(), &right.schema(), &join_type); + let mut eq_properties = EquivalenceProperties::new(Arc::new(schema)); + symmetric_join_output_partitioning( + left, + right, + &join_type, + on, + &mut eq_properties, + ) + } + fn full_join_partitioning( left: &Arc, right: &Arc, on: JoinOnRef, ) -> Result { - symmetric_join_output_partitioning(left, right, &JoinType::Full, on) + join_partitioning(left, right, JoinType::Full, on) } #[test] @@ -4631,11 +4610,11 @@ mod tests { Partitioning::Hash(vec![coalesce_keys(col_at("a", 0), col_at("c", 2))?], 4) ); assert_eq!( - symmetric_join_output_partitioning(&left, &right, &JoinType::Inner, &on)?, + join_partitioning(&left, &right, JoinType::Inner, &on)?, Partitioning::Hash(vec![col_at("c", 2)], 4) ); assert_eq!( - symmetric_join_output_partitioning(&left, &right, &JoinType::Left, &on)?, + join_partitioning(&left, &right, JoinType::Left, &on)?, Partitioning::Hash(vec![col_at("a", 0)], 4) ); Ok(()) @@ -4691,25 +4670,6 @@ mod tests { Ok(()) } - #[test] - fn full_join_output_partitioning_is_unknown_for_different_partition_counts() - -> Result<()> { - let left = repartitioned( - int32_schema("a", "b"), - Partitioning::Hash(vec![col_at("a", 0)], 4), - )?; - let right = repartitioned( - int32_schema("c", "d"), - Partitioning::Hash(vec![col_at("c", 0)], 2), - )?; - - assert!(matches!( - full_join_partitioning(&left, &right, &a_equals_c())?, - Partitioning::UnknownPartitioning(2) - )); - Ok(()) - } - #[test] fn full_join_output_partitioning_is_unknown_for_different_range_split_points() -> Result<()> { @@ -4730,64 +4690,6 @@ mod tests { Ok(()) } - #[test] - fn full_join_output_partitioning_is_unknown_for_different_range_sort_options() - -> Result<()> { - let left = repartitioned( - int32_schema("a", "b"), - range_on(col_at("a", 0), SortOptions::new(false, true), &[10, 20, 30])?, - )?; - let right = repartitioned( - int32_schema("c", "d"), - range_on( - col_at("c", 0), - SortOptions::new(false, false), - &[10, 20, 30], - )?, - )?; - - assert!(matches!( - full_join_partitioning(&left, &right, &a_equals_c())?, - Partitioning::UnknownPartitioning(4) - )); - Ok(()) - } - - #[test] - fn full_join_output_partitioning_is_unknown_for_hash_left_and_range_right() - -> Result<()> { - let left = repartitioned( - int32_schema("a", "b"), - Partitioning::Hash(vec![col_at("a", 0)], 4), - )?; - let right = repartitioned( - int32_schema("c", "d"), - range_on(col_at("c", 0), SortOptions::new(false, true), &[10, 20, 30])?, - )?; - - assert!(matches!( - full_join_partitioning(&left, &right, &a_equals_c())?, - Partitioning::UnknownPartitioning(4) - )); - Ok(()) - } - - #[test] - fn full_join_output_partitioning_is_unknown_for_unknown_input_partitioning() - -> Result<()> { - let left = unknown_partitioned(int32_schema("a", "b"), 4)?; - let right = repartitioned( - int32_schema("c", "d"), - Partitioning::Hash(vec![col_at("c", 0)], 4), - )?; - - assert!(matches!( - full_join_partitioning(&left, &right, &a_equals_c())?, - Partitioning::UnknownPartitioning(4) - )); - Ok(()) - } - #[test] fn full_join_output_partitioning_is_unknown_for_single_partition_inputs() -> Result<()> { @@ -4826,25 +4728,6 @@ mod tests { Ok(()) } - #[test] - fn full_join_output_partitioning_is_unknown_when_right_keys_are_not_the_join_keys() - -> Result<()> { - let left = repartitioned( - int32_schema("a", "b"), - Partitioning::Hash(vec![col_at("a", 0)], 4), - )?; - let right = repartitioned( - int32_schema("c", "d"), - Partitioning::Hash(vec![col_at("d", 1)], 4), - )?; - - assert!(matches!( - full_join_partitioning(&left, &right, &a_equals_c())?, - Partitioning::UnknownPartitioning(4) - )); - Ok(()) - } - #[test] fn full_join_output_partitioning_is_unknown_without_join_keys() -> Result<()> { let left = repartitioned( diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index 7c45bae15b689..1656b31aab184 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -1895,8 +1895,7 @@ ORDER BY b.range_key; ########## # TEST 49: Full Join Output Partitioning Feeds Aggregate on Coalesced Key -# Co-partitioned Full join inputs keep every row in the partition of its key, so -# an aggregate on COALESCE(l.range_key, r.range_key) needs no repartitioning. +# Co-partitioned inputs keep every row in the partition of its key, so no repartition. ########## statement ok @@ -1950,8 +1949,7 @@ ORDER BY k; ########## # TEST 50: Full Join Output Partitioning Feeds a Further Join on Coalesced Key -# The Range partitioning on the coalesced key flows through the projection the -# planner inserts, so the outer Full join needs no repartitioning on either side. +# The Range partitioning flows through the planner's projection, so no repartition. ########## query TT @@ -1988,8 +1986,7 @@ NULL NULL 40 400 ########## # TEST 51: Hash Repartitioned Full Join Feeds Aggregate on Coalesced Key -# Differing split points force Hash repartitioning below the join, and the -# resulting Hash partitioning on the coalesced key satisfies the aggregate above. +# Differing split points force Hash repartitioning below the join, not above it. ########## query TT @@ -2025,8 +2022,7 @@ ORDER BY k; ########## # TEST 52: Full Join Output Partitioning Does Not Cover a Single Side Key -# Right only rows carry NULL in l.range_key across every partition, so an -# aggregate on l.range_key alone still needs Hash repartitioning above the join. +# Right only rows carry NULL in l.range_key in every partition, so a repartition stays. ########## query TT @@ -2063,8 +2059,7 @@ NULL 2 ########## # TEST 53: Full Join Output Partitioning Matches Coalesce in Either Key Order -# Both coalesce orders are equal on every Full join row, so grouping by -# COALESCE(r.range_key, l.range_key) needs no repartitioning either. +# Both coalesce orders are equal on every Full join row, so no repartition either. ########## query TT From 1c21f893606010974ceda69c522abfe830a795f1 Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Tue, 22 Sep 2026 23:04:20 +0200 Subject: [PATCH 3/3] fix: updated coverage --- datafusion/physical-plan/src/joins/utils.rs | 222 +++++++++----------- 1 file changed, 98 insertions(+), 124 deletions(-) diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 8ab23868cc5e3..da28b6a3f6902 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -1978,13 +1978,7 @@ pub(crate) fn symmetric_join_output_partitioning( JoinType::Inner | JoinType::Right => { adjust_right_output_partitioning(right_partitioning, left_columns_len)? } - JoinType::Full => full_join_output_partitioning( - left, - right, - on, - left_columns_len, - eq_properties, - )?, + JoinType::Full => full_join_output_partitioning(left, right, on, eq_properties)?, }; Ok(result) } @@ -1993,9 +1987,9 @@ fn full_join_output_partitioning( left: &Arc, right: &Arc, on: JoinOnRef, - left_columns_len: usize, eq_properties: &mut EquivalenceProperties, ) -> Result { + let left_columns_len = left.schema().fields.len(); let left_partitioning = left.output_partitioning(); let unknown = Partitioning::UnknownPartitioning(right.output_partitioning().partition_count()); @@ -2726,7 +2720,6 @@ mod tests { use datafusion_common::stats::Precision::{Absent, Exact, Inexact}; use datafusion_common::{ScalarValue, SplitPoint, arrow_datafusion_err, arrow_err}; use datafusion_physical_expr::PhysicalSortExpr; - use datafusion_physical_expr::expressions::UnKnownColumn; use rstest::rstest; @@ -4527,41 +4520,40 @@ mod tests { fn unknown_partitioned( schema: Arc, partitions: usize, - ) -> Result> { - Ok(TestMemoryExec::try_new_exec( - &vec![vec![]; partitions], - schema, - None, - )?) + ) -> Arc { + TestMemoryExec::try_new_exec(&vec![vec![]; partitions], schema, None) + .expect("memory source") } fn repartitioned( schema: Arc, partitioning: Partitioning, - ) -> Result> { - let source = unknown_partitioned(schema, 1)?; - Ok(Arc::new(RepartitionExec::try_new(source, partitioning)?)) + ) -> Arc { + let source = unknown_partitioned(schema, 1); + Arc::new(RepartitionExec::try_new(source, partitioning).expect("repartition")) } fn col_at(name: &str, index: usize) -> PhysicalExprRef { Arc::new(Column::new(name, index)) } + fn coalesced(first: PhysicalExprRef, second: PhysicalExprRef) -> PhysicalExprRef { + coalesce_keys(first, second).expect("coalesce") + } + fn range_on( expr: PhysicalExprRef, options: SortOptions, split_points: &[i32], - ) -> Result { + ) -> Partitioning { let ordering = LexOrdering::new([PhysicalSortExpr::new(expr, options)]).expect("one key"); let split_points = split_points .iter() .map(|value| SplitPoint::new(vec![ScalarValue::Int32(Some(*value))])) .collect(); - Ok(Partitioning::Range(RangePartitioning::try_new( - ordering, - split_points, - )?)) + let range = RangePartitioning::try_new(ordering, split_points).expect("range"); + Partitioning::Range(range) } fn a_equals_c() -> JoinOn { @@ -4573,7 +4565,7 @@ mod tests { right: &Arc, join_type: JoinType, on: JoinOnRef, - ) -> Result { + ) -> Partitioning { let (schema, _) = build_join_schema(&left.schema(), &right.schema(), &join_type); let mut eq_properties = EquivalenceProperties::new(Arc::new(schema)); symmetric_join_output_partitioning( @@ -4583,77 +4575,76 @@ mod tests { on, &mut eq_properties, ) + .expect("join partitioning") } fn full_join_partitioning( left: &Arc, right: &Arc, on: JoinOnRef, - ) -> Result { + ) -> Partitioning { join_partitioning(left, right, JoinType::Full, on) } #[test] - fn full_join_output_partitioning_coalesces_hash_keys() -> Result<()> { + fn full_join_output_partitioning_coalesces_hash_keys() { let left = repartitioned( int32_schema("a", "b"), Partitioning::Hash(vec![col_at("a", 0)], 4), - )?; + ); let right = repartitioned( int32_schema("c", "d"), Partitioning::Hash(vec![col_at("c", 0)], 4), - )?; + ); let on = a_equals_c(); assert_eq!( - full_join_partitioning(&left, &right, &on)?, - Partitioning::Hash(vec![coalesce_keys(col_at("a", 0), col_at("c", 2))?], 4) + full_join_partitioning(&left, &right, &on), + Partitioning::Hash(vec![coalesced(col_at("a", 0), col_at("c", 2))], 4) ); assert_eq!( - join_partitioning(&left, &right, JoinType::Inner, &on)?, + join_partitioning(&left, &right, JoinType::Inner, &on), Partitioning::Hash(vec![col_at("c", 2)], 4) ); assert_eq!( - join_partitioning(&left, &right, JoinType::Left, &on)?, + join_partitioning(&left, &right, JoinType::Left, &on), Partitioning::Hash(vec![col_at("a", 0)], 4) ); - Ok(()) } #[test] - fn full_join_output_partitioning_coalesces_range_keys() -> Result<()> { + fn full_join_output_partitioning_coalesces_range_keys() { let options = SortOptions::new(false, true); let left = repartitioned( int32_schema("a", "b"), - range_on(col_at("a", 0), options, &[10, 20, 30])?, - )?; + range_on(col_at("a", 0), options, &[10, 20, 30]), + ); let right = repartitioned( int32_schema("c", "d"), - range_on(col_at("c", 0), options, &[10, 20, 30])?, - )?; + range_on(col_at("c", 0), options, &[10, 20, 30]), + ); let expected = range_on( - coalesce_keys(col_at("a", 0), col_at("c", 2))?, + coalesced(col_at("a", 0), col_at("c", 2)), options, &[10, 20, 30], - )?; + ); assert_eq!( - full_join_partitioning(&left, &right, &a_equals_c())?, + full_join_partitioning(&left, &right, &a_equals_c()), expected ); - Ok(()) } #[test] - fn full_join_output_partitioning_coalesces_every_key_pair() -> Result<()> { + fn full_join_output_partitioning_coalesces_every_key_pair() { let left = repartitioned( int32_schema("a", "b"), Partitioning::Hash(vec![col_at("a", 0), col_at("b", 1)], 4), - )?; + ); let right = repartitioned( int32_schema("c", "d"), Partitioning::Hash(vec![col_at("c", 0), col_at("d", 1)], 4), - )?; + ); let on = vec![ (col_at("a", 0), col_at("c", 0)), (col_at("b", 1), col_at("d", 1)), @@ -4661,101 +4652,93 @@ mod tests { let expected = Partitioning::Hash( vec![ - coalesce_keys(col_at("a", 0), col_at("c", 2))?, - coalesce_keys(col_at("b", 1), col_at("d", 3))?, + coalesced(col_at("a", 0), col_at("c", 2)), + coalesced(col_at("b", 1), col_at("d", 3)), ], 4, ); - assert_eq!(full_join_partitioning(&left, &right, &on)?, expected); - Ok(()) + assert_eq!(full_join_partitioning(&left, &right, &on), expected); } #[test] - fn full_join_output_partitioning_is_unknown_for_different_range_split_points() - -> Result<()> { + fn full_join_output_partitioning_is_unknown_for_different_range_split_points() { let options = SortOptions::new(false, true); let left = repartitioned( int32_schema("a", "b"), - range_on(col_at("a", 0), options, &[10, 20, 30])?, - )?; + range_on(col_at("a", 0), options, &[10, 20, 30]), + ); let right = repartitioned( int32_schema("c", "d"), - range_on(col_at("c", 0), options, &[10, 20, 40])?, - )?; + range_on(col_at("c", 0), options, &[10, 20, 40]), + ); - assert!(matches!( - full_join_partitioning(&left, &right, &a_equals_c())?, - Partitioning::UnknownPartitioning(4) - )); - Ok(()) + assert_eq!( + full_join_partitioning(&left, &right, &a_equals_c()).to_string(), + "UnknownPartitioning(4)" + ); } #[test] - fn full_join_output_partitioning_is_unknown_for_single_partition_inputs() -> Result<()> - { + fn full_join_output_partitioning_is_unknown_for_single_partition_inputs() { let left = repartitioned( int32_schema("a", "b"), Partitioning::Hash(vec![col_at("a", 0)], 1), - )?; + ); let right = repartitioned( int32_schema("c", "d"), Partitioning::Hash(vec![col_at("c", 0)], 1), - )?; + ); - assert!(matches!( - full_join_partitioning(&left, &right, &a_equals_c())?, - Partitioning::UnknownPartitioning(1) - )); - Ok(()) + assert_eq!( + full_join_partitioning(&left, &right, &a_equals_c()).to_string(), + "UnknownPartitioning(1)" + ); } #[test] - fn full_join_output_partitioning_is_unknown_when_keys_are_not_the_join_keys() - -> Result<()> { + fn full_join_output_partitioning_is_unknown_when_keys_are_not_the_join_keys() { let left = repartitioned( int32_schema("a", "b"), Partitioning::Hash(vec![col_at("b", 1)], 4), - )?; + ); let right = repartitioned( int32_schema("c", "d"), Partitioning::Hash(vec![col_at("c", 0)], 4), - )?; + ); - assert!(matches!( - full_join_partitioning(&left, &right, &a_equals_c())?, - Partitioning::UnknownPartitioning(4) - )); - Ok(()) + assert_eq!( + full_join_partitioning(&left, &right, &a_equals_c()).to_string(), + "UnknownPartitioning(4)" + ); } #[test] - fn full_join_output_partitioning_is_unknown_without_join_keys() -> Result<()> { + fn full_join_output_partitioning_is_unknown_without_join_keys() { let left = repartitioned( int32_schema("a", "b"), Partitioning::Hash(vec![col_at("a", 0)], 4), - )?; + ); let right = repartitioned( int32_schema("c", "d"), Partitioning::Hash(vec![col_at("c", 0)], 4), - )?; + ); - assert!(matches!( - full_join_partitioning(&left, &right, &[])?, - Partitioning::UnknownPartitioning(4) - )); - Ok(()) + assert_eq!( + full_join_partitioning(&left, &right, &[]).to_string(), + "UnknownPartitioning(4)" + ); } #[test] - fn full_sort_merge_join_reports_coalesced_partitioning() -> Result<()> { + fn full_sort_merge_join_reports_coalesced_partitioning() { let left = repartitioned( int32_schema("a", "b"), Partitioning::Hash(vec![col_at("a", 0)], 4), - )?; + ); let right = repartitioned( int32_schema("c", "d"), Partitioning::Hash(vec![col_at("c", 0)], 4), - )?; + ); let join = SortMergeJoinExec::try_new( left, right, @@ -4764,23 +4747,23 @@ mod tests { JoinType::Full, vec![SortOptions::default()], NullEquality::NullEqualsNothing, - )?; + ) + .expect("sort merge join"); let expected = - Partitioning::Hash(vec![coalesce_keys(col_at("a", 0), col_at("c", 2))?], 4); + Partitioning::Hash(vec![coalesced(col_at("a", 0), col_at("c", 2))], 4); assert_eq!(join.properties().output_partitioning(), &expected); - Ok(()) } #[test] - fn full_symmetric_hash_join_reports_coalesced_partitioning() -> Result<()> { + fn full_symmetric_hash_join_reports_coalesced_partitioning() { let left = repartitioned( int32_schema("a", "b"), Partitioning::Hash(vec![col_at("a", 0)], 4), - )?; + ); let right = repartitioned( int32_schema("c", "d"), Partitioning::Hash(vec![col_at("c", 0)], 4), - )?; + ); let join = SymmetricHashJoinExec::try_new( left, right, @@ -4791,22 +4774,22 @@ mod tests { None, None, StreamJoinPartitionMode::Partitioned, - )?; + ) + .expect("symmetric hash join"); let expected = - Partitioning::Hash(vec![coalesce_keys(col_at("a", 0), col_at("c", 2))?], 4); + Partitioning::Hash(vec![coalesced(col_at("a", 0), col_at("c", 2))], 4); assert_eq!(join.properties().output_partitioning(), &expected); - Ok(()) } - fn projected_full_hash_join(projection: Vec) -> Result { + fn projected_full_hash_join(projection: Vec) -> HashJoinExec { let left = repartitioned( int32_schema("a", "b"), Partitioning::Hash(vec![col_at("a", 0)], 4), - )?; + ); let right = repartitioned( int32_schema("c", "d"), Partitioning::Hash(vec![col_at("c", 0)], 4), - )?; + ); HashJoinExec::try_new( left, right, @@ -4818,48 +4801,39 @@ mod tests { NullEquality::NullEqualsNothing, false, ) + .expect("hash join") } #[test] - fn full_hash_join_keys_coalesce_equally_in_either_order() -> Result<()> { - let join = projected_full_hash_join(vec![0, 1, 2, 3])?; + fn full_hash_join_keys_coalesce_equally_in_either_order() { + let join = projected_full_hash_join(vec![0, 1, 2, 3]); let group = join.properties().equivalence_properties().eq_group(); - let left_first = coalesce_keys(col_at("a", 0), col_at("c", 2))?; - let right_first = coalesce_keys(col_at("c", 2), col_at("a", 0))?; + let left_first = coalesced(col_at("a", 0), col_at("c", 2)); + let right_first = coalesced(col_at("c", 2), col_at("a", 0)); assert!( group .normalize_expr(left_first) .eq(&group.normalize_expr(right_first)) ); - Ok(()) } #[test] - fn full_hash_join_projection_dropping_right_key_loses_the_key() -> Result<()> { - let join = projected_full_hash_join(vec![0, 1])?; - - let Partitioning::Hash(exprs, count) = join.properties().output_partitioning() - else { - panic!( - "expected hash partitioning, got {:?}", - join.properties().output_partitioning() - ); - }; - assert_eq!(*count, 4); - assert_eq!(exprs.len(), 1); - assert!(exprs[0].downcast_ref::().is_some()); - Ok(()) + fn full_hash_join_projection_dropping_right_key_loses_the_key() { + let join = projected_full_hash_join(vec![0, 1]); + let partitioning = join.properties().output_partitioning(); + + assert_eq!(partitioning.partition_count(), 4); + assert!(format!("{partitioning:?}").contains("UnKnownColumn")); } #[test] - fn full_hash_join_projection_remaps_coalesced_keys() -> Result<()> { - let join = projected_full_hash_join(vec![0, 2])?; + fn full_hash_join_projection_remaps_coalesced_keys() { + let join = projected_full_hash_join(vec![0, 2]); assert_eq!( join.properties().output_partitioning(), - &Partitioning::Hash(vec![coalesce_keys(col_at("a", 0), col_at("c", 1))?], 4) + &Partitioning::Hash(vec![coalesced(col_at("a", 0), col_at("c", 1))], 4) ); - Ok(()) } #[test]