diff --git a/datafusion/core/tests/physical_optimizer/sanity_checker.rs b/datafusion/core/tests/physical_optimizer/sanity_checker.rs index e759156282306..893d1aaeffe69 100644 --- a/datafusion/core/tests/physical_optimizer/sanity_checker.rs +++ b/datafusion/core/tests/physical_optimizer/sanity_checker.rs @@ -20,8 +20,9 @@ use std::sync::Arc; use crate::physical_optimizer::test_utils::{ bounded_window_exec, global_limit_exec, hash_join_exec, local_limit_exec, - memory_exec, projection_exec, repartition_exec, sort_exec, sort_expr, - sort_expr_options, sort_merge_join_exec, sort_preserving_merge_exec, union_exec, + memory_exec, projection_exec, repartition_exec, sort_exec, + sort_exec_with_preserve_partitioning, sort_expr, sort_expr_options, + sort_merge_join_exec, sort_preserving_merge_exec, union_exec, }; use arrow::compute::SortOptions; @@ -29,12 +30,13 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTable}; use datafusion::prelude::{CsvReadOptions, SessionContext}; use datafusion_common::config::ConfigOptions; -use datafusion_common::{JoinType, Result, ScalarValue}; +use datafusion_common::{JoinType, NullEquality, Result, ScalarValue}; use datafusion_physical_expr::expressions::{Literal, col}; use datafusion_physical_expr::{Partitioning, RangePartitioning, SplitPoint}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::sanity_checker::SanityCheckPlan; +use datafusion_physical_plan::joins::{StreamJoinPartitionMode, SymmetricHashJoinExec}; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::{ExecutionPlan, displayable}; @@ -417,6 +419,26 @@ fn range_partitioned_exec( .map(|exec| Arc::new(exec) as Arc) } +fn hash_partitioned_exec( + schema: &SchemaRef, + key: &str, + partition_count: usize, +) -> Result> { + let partitioning = Partitioning::Hash(vec![col(key, schema)?], partition_count); + RepartitionExec::try_new(memory_exec(schema), partitioning) + .map(|exec| Arc::new(exec) as Arc) +} + +fn sorted_hash_partitioned_exec( + schema: &SchemaRef, + key: &str, + partition_count: usize, +) -> Result> { + let input = hash_partitioned_exec(schema, key, partition_count)?; + let ordering: LexOrdering = [sort_expr(key, schema)].into(); + Ok(sort_exec_with_preserve_partitioning(ordering, input)) +} + #[test] fn test_partitioned_hash_join_requires_co_partitioned_children() -> Result<()> { let schema = create_test_schema2(); @@ -443,6 +465,64 @@ fn test_partitioned_hash_join_requires_co_partitioned_children() -> Result<()> { Ok(()) } +#[test] +fn test_sort_merge_join_requires_co_partitioned_children() -> Result<()> { + let schema = create_test_schema2(); + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; + + let compatible_join = sort_merge_join_exec( + sorted_hash_partitioned_exec(&schema, "a", 4)?, + sorted_hash_partitioned_exec(&schema, "a", 4)?, + &join_on, + &JoinType::Inner, + ); + assert_sanity_check(&compatible_join, true); + + let incompatible_join = sort_merge_join_exec( + sorted_hash_partitioned_exec(&schema, "a", 4)?, + sorted_hash_partitioned_exec(&schema, "a", 5)?, + &join_on, + &JoinType::Inner, + ); + assert_sanity_check(&incompatible_join, false); + + Ok(()) +} + +#[test] +fn test_symmetric_hash_join_requires_co_partitioned_children() -> Result<()> { + let schema = create_test_schema2(); + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; + + let compatible_join = Arc::new(SymmetricHashJoinExec::try_new( + hash_partitioned_exec(&schema, "a", 4)?, + hash_partitioned_exec(&schema, "a", 4)?, + join_on.clone(), + None, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + None, + None, + StreamJoinPartitionMode::Partitioned, + )?) as Arc; + assert_sanity_check(&compatible_join, true); + + let incompatible_join = Arc::new(SymmetricHashJoinExec::try_new( + hash_partitioned_exec(&schema, "a", 4)?, + hash_partitioned_exec(&schema, "a", 5)?, + join_on, + None, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + None, + None, + StreamJoinPartitionMode::Partitioned, + )?) as Arc; + assert_sanity_check(&incompatible_join, false); + + Ok(()) +} + #[tokio::test] /// Tests that plan is valid when the sort requirements are satisfied. async fn test_bounded_window_agg_sort_requirement() -> Result<()> { 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 82d9c900e85fc..84405a562100a 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -41,7 +41,8 @@ use crate::spill::spill_manager::SpillManager; use crate::statistics::StatisticsArgs; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - PlanProperties, SendableRecordBatchStream, Statistics, check_if_same_properties, + InputDistributionRequirements, PlanProperties, SendableRecordBatchStream, Statistics, + check_if_same_properties, }; use arrow::compute::SortOptions; @@ -413,13 +414,13 @@ impl ExecutionPlan for SortMergeJoinExec { self.input_distribution_requirements().into_per_child() } - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + fn input_distribution_requirements(&self) -> InputDistributionRequirements { let (left_expr, right_expr) = self .on .iter() .map(|(l, r)| (Arc::clone(l), Arc::clone(r))) .unzip(); - crate::InputDistributionRequirements::new(vec![ + InputDistributionRequirements::co_partitioned(vec![ Distribution::KeyPartitioned(left_expr), Distribution::KeyPartitioned(right_expr), ]) diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index f33d9b1d07e70..62f842c671de8 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -53,7 +53,8 @@ use crate::projection::{ use crate::stream::EmptyRecordBatchStream; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - PlanProperties, RecordBatchStream, SendableRecordBatchStream, + InputDistributionRequirements, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, joins::StreamJoinPartitionMode, metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; @@ -415,23 +416,26 @@ impl ExecutionPlan for SymmetricHashJoinExec { self.input_distribution_requirements().into_per_child() } - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(match self.mode { + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + match self.mode { StreamJoinPartitionMode::Partitioned => { let (left_expr, right_expr) = self .on .iter() .map(|(l, r)| (Arc::clone(l) as _, Arc::clone(r) as _)) .unzip(); - vec![ + InputDistributionRequirements::co_partitioned(vec![ Distribution::KeyPartitioned(left_expr), Distribution::KeyPartitioned(right_expr), - ] + ]) } StreamJoinPartitionMode::SinglePartition => { - vec![Distribution::SinglePartition, Distribution::SinglePartition] + InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + Distribution::SinglePartition, + ]) } - }) + } } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index aa741dded77be..156d91875e8ec 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -25,7 +25,8 @@ use datafusion::datasource::file_format::csv::CsvFormat; use datafusion::datasource::listing::{ ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, }; -use datafusion::logical_expr::{Partitioning, RangePartitioning, col}; +use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTable}; +use datafusion::logical_expr::{Partitioning, RangePartitioning, SortExpr, col}; use datafusion::prelude::SessionContext; // ============================================================================== @@ -52,11 +53,13 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { .expect("range partitioning should be valid"), ); + let range_table_dir = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/range_partitioned"); + register_csv_listing_table( ctx, "range_partitioned", - Path::new(env!("CARGO_MANIFEST_DIR")) - .join("test_files/scratch_range_partitioning/range_partitioned"), + &range_table_dir, Arc::clone(&schema), [ "1,1,10\n5,2,50\n", @@ -67,6 +70,14 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { Some(output_partitioning), ); + register_unbounded_stream_table( + ctx, + "unbounded_range_like", + range_table_dir.join("part-0.csv"), + Arc::clone(&schema), + vec![vec![col("range_key").sort(true, false)]], + ); + let shifted_output_partitioning = Partitioning::Range( RangePartitioning::try_new( vec![col("range_key").sort(true, true)], @@ -133,3 +144,17 @@ fn register_csv_listing_table( ctx.register_table(name, Arc::new(table)) .expect("test listing table registration should succeed"); } + +fn register_unbounded_stream_table( + ctx: &SessionContext, + name: &str, + file_path: impl Into, + schema: Arc, + order: Vec>, +) { + let source = FileStreamProvider::new_file(schema, file_path.into()); + let config = StreamConfig::new(Arc::new(source)).with_order(order); + + ctx.register_table(name, Arc::new(StreamTable::new(Arc::new(config)))) + .expect("test stream table registration should succeed"); +} diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index 5d004ca7fdfa5..62d4ee6c2cbf8 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -634,6 +634,80 @@ ORDER BY l.range_key; 30 600 35 700 +########## +# TEST 18: Sort Merge Join Repartitions Compatible Range Inputs +# SortMergeJoinExec does not opt into Range satisfying KeyPartitioned. These +# compatible Range inputs receive Hash repartitions, a future +# Range satisfaction implementation should remove them. +########## + +statement ok +set datafusion.optimizer.prefer_hash_join = false; + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] +02)--SortMergeJoinExec: join_type=Inner, on=[(range_key@0, range_key@0)] +03)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +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=csv, has_header=false +06)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +07)------RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +08)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +statement ok +reset datafusion.optimizer.prefer_hash_join; + +########## +# TEST 19: Symmetric Hash Join Repartitions Inputs +# SymmetricHashJoinExec does not opt into Range satisfying KeyPartitioned. +# The unbounded stream fixture does not yet expose Range partitioning, +# so the join still receives Hash repartitions. +########## + +statement ok +set datafusion.optimizer.prefer_hash_join = true; + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM unbounded_range_like l +FULL JOIN unbounded_range_like r ON l.range_key = r.range_key; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] +02)--SymmetricHashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)] +03)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=1, maintains_sort_order=true +04)------StreamingTableExec: partition_sizes=1, projection=[range_key, value], infinite_source=true, output_ordering=[range_key@0 ASC NULLS LAST] +05)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=1, maintains_sort_order=true +06)------StreamingTableExec: partition_sizes=1, projection=[range_key, value], infinite_source=true, output_ordering=[range_key@0 ASC NULLS LAST] + +query III rowsort +SELECT l.range_key, l.value, r.value +FROM unbounded_range_like l +FULL JOIN unbounded_range_like r ON l.range_key = r.range_key; +---- +1 10 10 +5 50 50 + statement ok reset datafusion.optimizer.prefer_hash_join; @@ -644,7 +718,7 @@ statement ok reset datafusion.optimizer.preserve_file_partitions; ########## -# TEST 18: Union of Range Partitioned Inputs +# TEST 20: Union of Range Partitioned Inputs # Each input exposes Range partitioning on range_key. These changes do not add a # cross-child Range relationship for UNION ALL. ##########