From 80edd7c34ec779facc82cfe4dc019464d639b3d4 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Sat, 11 Jul 2026 07:30:55 -0400 Subject: [PATCH] allow range partitioning on symmetric and sort merge joins --- datafusion/catalog/src/streaming.rs | 41 ++++++++- .../physical_optimizer/sanity_checker.rs | 79 +++++++++++++++- .../src/joins/sort_merge_join/exec.rs | 8 +- .../src/joins/symmetric_hash_join.rs | 19 ++-- datafusion/physical-plan/src/streaming.rs | 48 +++++++++- .../src/test_context/range_partitioning.rs | 89 ++++++++++++++++++- .../test_files/range_partitioning.slt | 76 +++++++++++++++- 7 files changed, 336 insertions(+), 24 deletions(-) diff --git a/datafusion/catalog/src/streaming.rs b/datafusion/catalog/src/streaming.rs index e609877c2b778..5bfecef1fb2ed 100644 --- a/datafusion/catalog/src/streaming.rs +++ b/datafusion/catalog/src/streaming.rs @@ -24,7 +24,10 @@ use async_trait::async_trait; use datafusion_common::{DFSchema, Result, plan_err}; use datafusion_expr::{Expr, SortExpr, TableType}; use datafusion_physical_expr::equivalence::project_ordering; -use datafusion_physical_expr::{LexOrdering, create_physical_sort_exprs}; +use datafusion_physical_expr::projection::ProjectionMapping; +use datafusion_physical_expr::{ + EquivalenceProperties, LexOrdering, Partitioning, create_physical_sort_exprs, +}; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::streaming::{PartitionStream, StreamingTableExec}; use log::debug; @@ -38,6 +41,7 @@ pub struct StreamingTable { partitions: Vec>, infinite: bool, sort_order: Vec, + output_partitioning: Option, } impl StreamingTable { @@ -62,6 +66,7 @@ impl StreamingTable { partitions, infinite: false, sort_order: vec![], + output_partitioning: None, }) } @@ -76,6 +81,33 @@ impl StreamingTable { self.sort_order = sort_order; self } + + /// Declares the output partitioning of this streaming table. + /// + /// The partitioning expressions refer to the table schema before scan + /// projection. If a scan projection removes a partitioning expression, the + /// physical plan reports unknown partitioning. + pub fn with_output_partitioning(mut self, output_partitioning: Partitioning) -> Self { + self.output_partitioning = Some(output_partitioning); + self + } + + fn output_partitioning( + &self, + projection: Option<&Vec>, + ) -> Result { + let Some(output_partitioning) = &self.output_partitioning else { + return Ok(Partitioning::UnknownPartitioning(self.partitions.len())); + }; + let Some(projection) = projection else { + return Ok(output_partitioning.clone()); + }; + + let projection_mapping = + ProjectionMapping::from_indices(projection, &self.schema)?; + let eq_properties = EquivalenceProperties::new(Arc::clone(&self.schema)); + Ok(output_partitioning.project(&projection_mapping, &eq_properties)) + } } #[async_trait] @@ -119,13 +151,16 @@ impl TableProvider for StreamingTable { vec![] }; - Ok(Arc::new(StreamingTableExec::try_new( + let exec = StreamingTableExec::try_new( Arc::clone(&self.schema), self.partitions.clone(), projection, LexOrdering::new(physical_sort), self.infinite, limit, - )?)) + )? + .with_output_partitioning(self.output_partitioning(projection)?)?; + + Ok(Arc::new(exec)) } } diff --git a/datafusion/core/tests/physical_optimizer/sanity_checker.rs b/datafusion/core/tests/physical_optimizer/sanity_checker.rs index e759156282306..5ffae658a40c8 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}; @@ -443,6 +445,77 @@ 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 ordering: LexOrdering = [sort_expr("a", &schema)].into(); + + let compatible_join = sort_merge_join_exec( + sort_exec_with_preserve_partitioning( + ordering.clone(), + range_partitioned_exec(&schema, "a", [10])?, + ), + sort_exec_with_preserve_partitioning( + ordering.clone(), + range_partitioned_exec(&schema, "a", [10])?, + ), + &join_on, + &JoinType::Inner, + ); + assert_sanity_check(&compatible_join, true); + + let incompatible_join = sort_merge_join_exec( + sort_exec_with_preserve_partitioning( + ordering.clone(), + range_partitioned_exec(&schema, "a", [10])?, + ), + sort_exec_with_preserve_partitioning( + ordering, + range_partitioned_exec(&schema, "a", [20])?, + ), + &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( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [10])?, + 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( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [20])?, + 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 46c37696ece14..5d3621c49219b 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::{ChildStats, 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,16 +414,17 @@ 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), ]) + .allow_range_satisfaction_for_key_partitioning() } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index f33d9b1d07e70..0aaf8ac608b0b 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,27 @@ 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), - ] + ]) + .allow_range_satisfaction_for_key_partitioning() } StreamJoinPartitionMode::SinglePartition => { - vec![Distribution::SinglePartition, Distribution::SinglePartition] + InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + Distribution::SinglePartition, + ]) } - }) + } } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/physical-plan/src/streaming.rs b/datafusion/physical-plan/src/streaming.rs index cdf4b08f718c6..61a9b9cc6d0de 100644 --- a/datafusion/physical-plan/src/streaming.rs +++ b/datafusion/physical-plan/src/streaming.rs @@ -35,6 +35,7 @@ use crate::{ExecutionPlan, Partitioning, SendableRecordBatchStream}; use arrow::datatypes::{Schema, SchemaRef}; use datafusion_common::{Result, internal_err, plan_err}; use datafusion_execution::TaskContext; +use datafusion_physical_expr::projection::ProjectionMapping; use datafusion_physical_expr::{EquivalenceProperties, LexOrdering}; use async_trait::async_trait; @@ -100,7 +101,7 @@ impl StreamingTableExec { let cache = Self::compute_properties( Arc::clone(&projected_schema), projected_output_ordering.clone(), - &partitions, + Partitioning::UnknownPartitioning(partitions.len()), infinite, ); Ok(Self { @@ -115,6 +116,25 @@ impl StreamingTableExec { }) } + /// Declares the output partitioning of this stream. + /// + /// `output_partitioning` must describe this plan's current output and have + /// the same number of partitions as the stream. + pub fn with_output_partitioning( + mut self, + output_partitioning: Partitioning, + ) -> Result { + if output_partitioning.partition_count() != self.partitions.len() { + return plan_err!( + "Output partitioning has {} partitions but stream has {} partitions", + output_partitioning.partition_count(), + self.partitions.len() + ); + } + Arc::make_mut(&mut self.cache).partitioning = output_partitioning; + Ok(self) + } + pub fn partitions(&self) -> &Vec> { &self.partitions } @@ -147,14 +167,12 @@ impl StreamingTableExec { fn compute_properties( schema: SchemaRef, orderings: Vec, - partitions: &[Arc], + output_partitioning: Partitioning, infinite: bool, ) -> PlanProperties { // Calculate equivalence properties: let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings); - // Get output partitioning: - let output_partitioning = Partitioning::UnknownPartitioning(partitions.len()); let boundedness = if infinite { Boundedness::Unbounded { requires_infinite_memory: false, @@ -204,6 +222,16 @@ impl DisplayAs for StreamingTableExec { if let Some(fetch) = self.limit { write!(f, ", fetch={fetch}")?; } + if !matches!( + self.cache.output_partitioning(), + Partitioning::UnknownPartitioning(_) + ) { + write!( + f, + ", output_partitioning={}", + self.cache.output_partitioning() + )?; + } display_orderings(f, &self.projected_output_ordering)?; @@ -306,6 +334,17 @@ impl ExecutionPlan for StreamingTableExec { }; lex_orderings.push(ordering); } + let projection_mapping = ProjectionMapping::try_new( + projection + .expr() + .iter() + .map(|expr| (Arc::clone(&expr.expr), expr.alias.clone())), + &self.schema(), + )?; + let output_partitioning = self + .cache + .output_partitioning() + .project(&projection_mapping, self.cache.equivalence_properties()); StreamingTableExec::try_new( Arc::clone(self.partition_schema()), @@ -315,6 +354,7 @@ impl ExecutionPlan for StreamingTableExec { self.is_infinite(), self.limit(), ) + .and_then(|exec| exec.with_output_partitioning(output_partitioning)) .map(|e| Some(Arc::new(e) as _)) } diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index aa741dded77be..3cab5fd30b14c 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -19,13 +19,23 @@ use std::fs::{create_dir_all, remove_dir_all, write}; use std::path::Path; use std::sync::Arc; -use arrow::datatypes::{DataType, Field, Schema}; +use arrow::array::{ArrayRef, Int32Array}; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion::catalog::streaming::StreamingTable; use datafusion::common::{ScalarValue, SplitPoint}; use datafusion::datasource::file_format::csv::CsvFormat; use datafusion::datasource::listing::{ ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, }; use datafusion::logical_expr::{Partitioning, RangePartitioning, col}; +use datafusion::physical_expr::{ + Partitioning as PhysicalPartitioning, PhysicalSortExpr, + RangePartitioning as PhysicalRangePartitioning, expressions::col as physical_col, +}; +use datafusion::physical_plan::streaming::PartitionStream; +use datafusion::physical_plan::test::TestPartitionStream; use datafusion::prelude::SessionContext; // ============================================================================== @@ -52,11 +62,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 +79,12 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { Some(output_partitioning), ); + register_unbounded_range_stream_table( + ctx, + "unbounded_range_like", + Arc::clone(&schema), + ); + let shifted_output_partitioning = Partitioning::Range( RangePartitioning::try_new( vec![col("range_key").sort(true, true)], @@ -133,3 +151,68 @@ fn register_csv_listing_table( ctx.register_table(name, Arc::new(table)) .expect("test listing table registration should succeed"); } + +fn register_unbounded_range_stream_table( + ctx: &SessionContext, + name: &str, + schema: Arc, +) { + let output_partitioning = PhysicalPartitioning::Range( + PhysicalRangePartitioning::try_new( + [PhysicalSortExpr { + expr: physical_col("range_key", &schema) + .expect("range key should exist in stream schema"), + options: SortOptions::default(), + }] + .into(), + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(30))]), + ], + ) + .expect("range partitioning should be valid"), + ); + let partitions = [ + vec![(1, 1, 10), (5, 2, 50)], + vec![(10, 1, 100), (15, 2, 150)], + vec![(20, 1, 200), (25, 2, 250)], + vec![(30, 1, 300), (35, 2, 350)], + ] + .into_iter() + .map(|rows| range_stream_partition(Arc::clone(&schema), &rows)) + .collect(); + + ctx.register_table( + name, + Arc::new( + StreamingTable::try_new(schema, partitions) + .expect("range stream table should be valid") + .with_infinite_table(true) + .with_output_partitioning(output_partitioning), + ), + ) + .expect("test stream table registration should succeed"); +} + +fn range_stream_partition( + schema: SchemaRef, + rows: &[(i32, i32, i32)], +) -> Arc { + let range_key: Vec = rows.iter().map(|(range_key, _, _)| *range_key).collect(); + let non_range_key: Vec = rows + .iter() + .map(|(_, non_range_key, _)| *non_range_key) + .collect(); + let value: Vec = rows.iter().map(|(_, _, value)| *value).collect(); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(range_key)) as ArrayRef, + Arc::new(Int32Array::from(non_range_key)) as ArrayRef, + Arc::new(Int32Array::from(value)) as ArrayRef, + ], + ) + .expect("range stream batch should be valid"); + Arc::new(TestPartitionStream::new_with_batches(vec![batch])) +} diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index 5d004ca7fdfa5..5af70921a9fa5 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 Avoids Repartition for Compatible Range Inputs +# Compatible Range inputs satisfy SortMergeJoinExec's co-partitioned +# KeyPartitioned requirements. +########## + +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)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +05)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +06)------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 Avoids Repartition for Compatible Range Inputs +# Compatible Range streams satisfy SymmetricHashJoinExec's co-partitioned +# KeyPartitioned requirements. +########## + +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)----StreamingTableExec: partition_sizes=4, projection=[range_key, value], infinite_source=true, output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) +04)----StreamingTableExec: partition_sizes=4, projection=[range_key, value], infinite_source=true, output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) + +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 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 +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. ##########