diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 43ceee3444ced..45b342ee8b5ac 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..9cfaeddb34667 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,184 @@ fn range_left_anti_hash_join_rehashes_incompatible_null_options() -> Result<()> Ok(()) } +/// 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> { + 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 `coalesce(first, second)` over two columns of `join` in its physical CASE form. +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), + ) +} + +/// 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, + 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 9858a4c06cd4a..080c4fe70d046 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1294,9 +1294,13 @@ impl HashJoinExec { PartitionMode::Auto => Partitioning::UnknownPartitioning( right.output_partitioning().partition_count(), ), - PartitionMode::Partitioned => { - symmetric_join_output_partitioning(left, right, &join_type)? - } + PartitionMode::Partitioned => symmetric_join_output_partitioning( + left, + right, + &join_type, + on, + &mut eq_properties, + )?, }; let emission_type = 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 c0e57a93a49a1..fe5c739d174cf 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -326,8 +326,13 @@ impl SortMergeJoinExec { join_on, )?; - let mut output_partitioning = - symmetric_join_output_partitioning(left, right, &join_type)?; + let mut output_partitioning = symmetric_join_output_partitioning( + left, + right, + &join_type, + join_on, + &mut eq_properties, + )?; 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..5a718abdaa72d 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -268,7 +268,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, @@ -279,8 +279,13 @@ impl SymmetricHashJoinExec { join_on, )?; - let output_partitioning = - symmetric_join_output_partitioning(left, right, &join_type)?; + let output_partitioning = symmetric_join_output_partitioning( + left, + right, + &join_type, + join_on, + &mut eq_properties, + )?; 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..da28b6a3f6902 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, + add_offset_to_expr, add_offset_to_physical_sort_exprs, }; use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; @@ -1954,10 +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(); @@ -1972,14 +1978,67 @@ pub(crate) fn symmetric_join_output_partitioning( JoinType::Inner | JoinType::Right => { 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()) - } + JoinType::Full => full_join_output_partitioning(left, right, on, eq_properties)?, }; Ok(result) } +fn full_join_output_partitioning( + left: &Arc, + right: &Arc, + on: JoinOnRef, + 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()); + 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 `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`, the physical form the +/// planner gives `coalesce(first, second)`. +fn coalesce_keys( + first: PhysicalExprRef, + second: PhysicalExprRef, +) -> Result { + case( + None, + vec![(is_not_null(Arc::clone(&first))?, first)], + Some(second), + ) +} + /// 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,7 +2707,13 @@ 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}; @@ -4445,6 +4510,332 @@ 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, + ) -> Arc { + TestMemoryExec::try_new_exec(&vec![vec![]; partitions], schema, None) + .expect("memory source") + } + + fn repartitioned( + schema: Arc, + partitioning: 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], + ) -> 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(); + let range = RangePartitioning::try_new(ordering, split_points).expect("range"); + Partitioning::Range(range) + } + + fn a_equals_c() -> JoinOn { + vec![(col_at("a", 0), col_at("c", 0))] + } + + fn join_partitioning( + left: &Arc, + right: &Arc, + join_type: JoinType, + on: JoinOnRef, + ) -> 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( + left, + right, + &join_type, + on, + &mut eq_properties, + ) + .expect("join partitioning") + } + + fn full_join_partitioning( + left: &Arc, + right: &Arc, + on: JoinOnRef, + ) -> Partitioning { + join_partitioning(left, right, JoinType::Full, on) + } + + #[test] + 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![coalesced(col_at("a", 0), col_at("c", 2))], 4) + ); + assert_eq!( + join_partitioning(&left, &right, JoinType::Inner, &on), + Partitioning::Hash(vec![col_at("c", 2)], 4) + ); + assert_eq!( + join_partitioning(&left, &right, JoinType::Left, &on), + Partitioning::Hash(vec![col_at("a", 0)], 4) + ); + } + + #[test] + 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]), + ); + let right = repartitioned( + int32_schema("c", "d"), + range_on(col_at("c", 0), options, &[10, 20, 30]), + ); + + let expected = range_on( + coalesced(col_at("a", 0), col_at("c", 2)), + options, + &[10, 20, 30], + ); + assert_eq!( + full_join_partitioning(&left, &right, &a_equals_c()), + expected + ); + } + + #[test] + 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)), + ]; + + let expected = Partitioning::Hash( + vec![ + 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); + } + + #[test] + 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]), + ); + let right = repartitioned( + int32_schema("c", "d"), + range_on(col_at("c", 0), options, &[10, 20, 40]), + ); + + 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() { + 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_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() { + 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_eq!( + full_join_partitioning(&left, &right, &a_equals_c()).to_string(), + "UnknownPartitioning(4)" + ); + } + + #[test] + 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_eq!( + full_join_partitioning(&left, &right, &[]).to_string(), + "UnknownPartitioning(4)" + ); + } + + #[test] + 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, + a_equals_c(), + None, + JoinType::Full, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + ) + .expect("sort merge join"); + let expected = + Partitioning::Hash(vec![coalesced(col_at("a", 0), col_at("c", 2))], 4); + assert_eq!(join.properties().output_partitioning(), &expected); + } + + #[test] + 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, + a_equals_c(), + None, + &JoinType::Full, + NullEquality::NullEqualsNothing, + None, + None, + StreamJoinPartitionMode::Partitioned, + ) + .expect("symmetric hash join"); + let expected = + Partitioning::Hash(vec![coalesced(col_at("a", 0), col_at("c", 2))], 4); + assert_eq!(join.properties().output_partitioning(), &expected); + } + + 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, + a_equals_c(), + None, + &JoinType::Full, + Some(projection), + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + .expect("hash join") + } + + #[test] + 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 = 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)) + ); + } + + #[test] + 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() { + let join = projected_full_hash_join(vec![0, 2]); + + assert_eq!( + join.properties().output_partitioning(), + &Partitioning::Hash(vec![coalesced(col_at("a", 0), col_at("c", 1))], 4) + ); + } + #[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..1656b31aab184 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -1893,6 +1893,212 @@ ORDER BY b.range_key; 5 50 50 20 200 200 +########## +# TEST 49: Full Join Output Partitioning Feeds Aggregate on Coalesced Key +# Co-partitioned inputs keep every row in the partition of its key, so no repartition. +########## + +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 flows through the planner's projection, so no repartition. +########## + +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, not above it. +########## + +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 in every partition, so a repartition stays. +########## + +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 no repartition 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;