From 55e9ae004c12e9b9bc56f63685597d9a6c6941c2 Mon Sep 17 00:00:00 2001 From: Shayan Gh <98089795+ShayanGho@users.noreply.github.com> Date: Sat, 26 Sep 2026 21:36:11 -0500 Subject: [PATCH 1/4] fix: consider input sizes before reusing join partitioning --- .../enforce_distribution.rs | 102 ++++++++++++++++++ .../enforce_distribution.rs | 13 +++ 2 files changed, 115 insertions(+) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 7c9d4984a6a68..5cf3fe328baaa 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -402,6 +402,28 @@ fn parquet_exec_with_output_partitioning( DataSourceExec::from_data_source(config) } +fn parquet_exec_with_output_partitioning_and_statistics( + output_partitioning: Partitioning, + statistics: Statistics, +) -> Arc { + let file_groups = (0..output_partitioning.partition_count()) + .map(|partition| { + FileGroup::new(vec![PartitionedFile::new(format!("p{partition}"), 100)]) + }) + .collect::>(); + + let config = FileScanConfigBuilder::new( + ObjectStoreUrl::parse("test:///").unwrap(), + Arc::new(ParquetSource::new(schema())), + ) + .with_file_groups(file_groups) + .with_output_partitioning(Some(output_partitioning)) + .with_statistics(statistics) + .build(); + + DataSourceExec::from_data_source(config) +} + fn range_partitioning( column: &str, split_values: impl IntoIterator, @@ -1123,6 +1145,86 @@ fn range_hash_join_repartitions_unpartitioned_side_to_match_range() -> Result<() Ok(()) } +#[test] +fn range_hash_join_rejects_smaller_reference() -> Result<()> { + let stats = |rows, bytes| { + let mut statistics = Statistics::new_unknown(&schema()); + statistics.num_rows = Precision::Inexact(rows); + statistics.total_byte_size = Precision::Inexact(bytes); + statistics + }; + + let left = parquet_exec_with_output_partitioning_and_statistics( + Partitioning::UnknownPartitioning(4), + stats(100_000, 8_000_000), + ); + let right = parquet_exec_with_output_partitioning_and_statistics( + range_partitioning("a", [10, 20, 30], SortOptions::default())?, + stats(100, 8_000), + ); + + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + let children = plan.children(); + assert_eq!(children.len(), 2); + for child in children { + assert!( + matches!(child.output_partitioning(), Partitioning::Hash(_, 4)), + "expected hash fallback, got {:?}", + child.output_partitioning(), + ); + } + + Ok(()) +} + +#[test] +fn range_hash_join_preserves_equal_reference() -> Result<()> { + let stats = |rows, bytes| { + let mut statistics = Statistics::new_unknown(&schema()); + statistics.num_rows = Precision::Inexact(rows); + statistics.total_byte_size = Precision::Inexact(bytes); + statistics + }; + + let left = parquet_exec_with_output_partitioning_and_statistics( + Partitioning::UnknownPartitioning(4), + stats(100, 8_000), + ); + let right = parquet_exec_with_output_partitioning_and_statistics( + range_partitioning("a", [10, 20, 30], SortOptions::default())?, + stats(100, 8_000), + ); + + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + let expected = range_partitioning("a", [10, 20, 30], SortOptions::default())?; + + let children = plan.children(); + assert_eq!(children.len(), 2); + for child in children { + assert_eq!(child.output_partitioning(), &expected); + } + + Ok(()) +} + #[test] fn range_hash_join_rehashes_incompatible_data_type() -> Result<()> { let left = parquet_exec(); diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 776e3c3bf14bd..0595f264fd37d 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -1212,6 +1212,19 @@ fn enforce_distribution_relationships( } }; + // Reject a reference when another child in the co-partitioned group + // is strictly larger, allowing standard hash partitioning instead. + let best_satisfied_child = best_satisfied_child.filter(|(ref_idx, _)| { + let reference_size = + PlanSize::from_plan(children[*ref_idx].context.plan.as_ref()); + + !unsatisfied_children.iter().any(|&child_idx| { + child_idx != *ref_idx + && PlanSize::from_plan(children[child_idx].context.plan.as_ref()) + > reference_size + }) + }); + // Validate that best_satisfied_child can be adapted across all unsatisfied children in the group let best_satisfied_child = best_satisfied_child.filter(|(ref_idx, ref_partitioning)| { From 15503e745a2e95755b3ae7e0148c854681aefe48 Mon Sep 17 00:00:00 2001 From: Shayan Gh <98089795+ShayanGho@users.noreply.github.com> Date: Sat, 26 Sep 2026 22:10:11 -0500 Subject: [PATCH 2/4] test: consolidate partitioned scan helpers --- .../enforce_distribution.rs | 61 +++++++------------ 1 file changed, 23 insertions(+), 38 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 5cf3fe328baaa..0acd0c6611252 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -385,27 +385,22 @@ fn parquet_exec_multiple_sorted( fn parquet_exec_with_output_partitioning( output_partitioning: Partitioning, ) -> Arc { - let file_groups = (0..output_partitioning.partition_count()) - .map(|partition| { - FileGroup::new(vec![PartitionedFile::new(format!("p{partition}"), 100)]) - }) - .collect::>(); - - let config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("test:///").unwrap(), - Arc::new(ParquetSource::new(schema())), + parquet_exec_with_partitioning_and_size( + output_partitioning, + Precision::Absent, + Precision::Absent, ) - .with_file_groups(file_groups) - .with_output_partitioning(Some(output_partitioning)) - .build(); - - DataSourceExec::from_data_source(config) } -fn parquet_exec_with_output_partitioning_and_statistics( +fn parquet_exec_with_partitioning_and_size( output_partitioning: Partitioning, - statistics: Statistics, + num_rows: Precision, + total_byte_size: Precision, ) -> Arc { + let mut statistics = Statistics::new_unknown(&schema()); + statistics.num_rows = num_rows; + statistics.total_byte_size = total_byte_size; + let file_groups = (0..output_partitioning.partition_count()) .map(|partition| { FileGroup::new(vec![PartitionedFile::new(format!("p{partition}"), 100)]) @@ -1147,20 +1142,15 @@ fn range_hash_join_repartitions_unpartitioned_side_to_match_range() -> Result<() #[test] fn range_hash_join_rejects_smaller_reference() -> Result<()> { - let stats = |rows, bytes| { - let mut statistics = Statistics::new_unknown(&schema()); - statistics.num_rows = Precision::Inexact(rows); - statistics.total_byte_size = Precision::Inexact(bytes); - statistics - }; - - let left = parquet_exec_with_output_partitioning_and_statistics( + let left = parquet_exec_with_partitioning_and_size( Partitioning::UnknownPartitioning(4), - stats(100_000, 8_000_000), + Precision::Inexact(100_000), + Precision::Inexact(8_000_000), ); - let right = parquet_exec_with_output_partitioning_and_statistics( + let right = parquet_exec_with_partitioning_and_size( range_partitioning("a", [10, 20, 30], SortOptions::default())?, - stats(100, 8_000), + Precision::Inexact(100), + Precision::Inexact(8_000), ); let join_on = vec![( @@ -1188,20 +1178,15 @@ fn range_hash_join_rejects_smaller_reference() -> Result<()> { #[test] fn range_hash_join_preserves_equal_reference() -> Result<()> { - let stats = |rows, bytes| { - let mut statistics = Statistics::new_unknown(&schema()); - statistics.num_rows = Precision::Inexact(rows); - statistics.total_byte_size = Precision::Inexact(bytes); - statistics - }; - - let left = parquet_exec_with_output_partitioning_and_statistics( + let left = parquet_exec_with_partitioning_and_size( Partitioning::UnknownPartitioning(4), - stats(100, 8_000), + Precision::Inexact(100), + Precision::Inexact(8_000), ); - let right = parquet_exec_with_output_partitioning_and_statistics( + let right = parquet_exec_with_partitioning_and_size( range_partitioning("a", [10, 20, 30], SortOptions::default())?, - stats(100, 8_000), + Precision::Inexact(100), + Precision::Inexact(8_000), ); let join_on = vec![( From 1aab317699d85751304d228e4feead0860987f6c Mon Sep 17 00:00:00 2001 From: Shayan Gh <98089795+ShayanGho@users.noreply.github.com> Date: Sun, 11 Oct 2026 22:56:56 -0500 Subject: [PATCH 3/4] fix: reject only small references when a large input would adapt Reject a reference child only when it is below the CollectLeft thresholds and another co-partitioned input is above them, so inputs of comparable size keep the reference and only one side is shuffled. Restores the original assertion of range_fallback_preserves_target_parallelism. Co-Authored-By: Claude Opus 5.5 --- .../enforce_distribution.rs | 42 +++++++++++++++++-- .../enforce_distribution.rs | 31 +++++++++++--- 2 files changed, 64 insertions(+), 9 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 8e4e29cab26c8..974b9da617119 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -1257,6 +1257,42 @@ fn range_hash_join_preserves_equal_reference() -> Result<()> { Ok(()) } +#[test] +fn range_hash_join_preserves_large_reference() -> Result<()> { + // Both inputs are above the `CollectLeft` thresholds, so the slightly + // larger one is adapted to the reference and only it is shuffled. + let left = parquet_exec_with_partitioning_and_size( + Partitioning::UnknownPartitioning(4), + Precision::Inexact(11_000_000), + Precision::Inexact(880_000_000), + ); + let right = parquet_exec_with_partitioning_and_size( + range_partitioning("a", [10, 20, 30], SortOptions::default())?, + Precision::Inexact(10_000_000), + Precision::Inexact(800_000_000), + ); + + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + let expected = range_partitioning("a", [10, 20, 30], SortOptions::default())?; + + let children = plan.children(); + assert_eq!(children.len(), 2); + for child in children { + assert_eq!(child.output_partitioning(), &expected); + } + + Ok(()) +} + #[test] fn range_singleton_scaling_respects_required_keys() -> Result<()> { let ordering = [PhysicalSortExpr::new_default(col("a", &schema())?)].into(); @@ -1496,13 +1532,11 @@ fn range_fallback_preserves_target_parallelism() -> Result<()> { .with_query_execution_partitions(4) .to_plan(join, &DISTRIB_DISTRIB_SORT); let rendered = displayable(plan.as_ref()).indent(true).to_string(); - // The native four-partition input is smaller, so reject it as a - // reference and preserve target parallelism with hash partitioning. for child in plan.children() { - let Partitioning::Hash(_, count) = child.output_partitioning() else { + let Partitioning::Range(range) = child.output_partitioning() else { panic!("{rendered}") }; - assert_eq!(*count, 4, "{rendered}"); + assert_eq!(range.partition_count(), 4, "{rendered}"); } } Ok(()) diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 8b919d38f1c51..ead9463355c1e 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -1078,6 +1078,21 @@ impl PlanSize { fn is_known(&self) -> bool { self.byte_size.is_some() || self.num_rows.is_some() } + + /// Whether the plan is below the `CollectLeft` thresholds of `JoinSelection`, + /// comparing bytes when known and rows otherwise. `None` when the size is unknown. + fn is_small(&self, config: &ConfigOptions) -> Option { + let optimizer = &config.optimizer; + match (self.byte_size, self.num_rows) { + (Some(bytes), _) => { + Some(bytes < optimizer.hash_join_single_partition_threshold) + } + (None, Some(rows)) => { + Some(rows < optimizer.hash_join_single_partition_threshold_rows) + } + (None, None) => None, + } + } } impl PartialOrd for PlanSize { @@ -1123,8 +1138,9 @@ fn enforce_distribution_relationships( plan_name: &str, input_distributions: &InputDistributionRequirements, children: &mut [DistributionChildState], - target_partitions: usize, + config: &ConfigOptions, ) -> Result<()> { + let target_partitions = config.execution.target_partitions; let mut repartitioned_for_relationship = vec![false; children.len()]; loop { @@ -1210,16 +1226,21 @@ fn enforce_distribution_relationships( } }; - // Reject a reference when another child in the co-partitioned group - // is strictly larger, allowing standard hash partitioning instead. + // Do not adapt a child above the `CollectLeft` thresholds to a reference + // below them; fall back to hash partitioning instead. Otherwise keep the + // reference, so that it is not shuffled as well. let best_satisfied_child = best_satisfied_child.filter(|(ref_idx, _)| { let reference_size = PlanSize::from_plan(children[*ref_idx].context.plan.as_ref()); + if reference_size.is_small(config) != Some(true) { + return true; + } !unsatisfied_children.iter().any(|&child_idx| { child_idx != *ref_idx && PlanSize::from_plan(children[child_idx].context.plan.as_ref()) - > reference_size + .is_small(config) + == Some(false) }) }); @@ -1678,7 +1699,7 @@ pub fn ensure_distribution_with_stats( plan.name(), &input_distributions, &mut children, - target_partitions, + config, )?; let children = children From 13a7a19fe7e95c9aada80dd1a0f2f796abf00e57 Mon Sep 17 00:00:00 2001 From: Shayan Gh <98089795+ShayanGho@users.noreply.github.com> Date: Tue, 10 Nov 2026 21:14:26 -0600 Subject: [PATCH 4/4] test: cover the reference size gate and share the CollectLeft rule Replace the three reference-size tests with a table test covering the byte and row thresholds, their boundaries, unknown sizes, and both hash and sort-merge joins. Share the bytes-then-rows rule with JoinSelection through below_collect_thresholds, and document that the thresholds also steer the co-partitioning reference for partitioned joins. Co-Authored-By: Claude Fable 5.2 --- datafusion/common/src/config.rs | 8 +- .../enforce_distribution.rs | 198 +++++++++--------- .../enforce_distribution.rs | 42 ++-- .../physical-optimizer/src/join_selection.rs | 20 +- .../test_files/information_schema.slt | 4 +- docs/source/user-guide/configs.md | 4 +- 6 files changed, 142 insertions(+), 134 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 5106ea51ea6ba..b915c0760d64a 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1849,11 +1849,15 @@ config_namespace! { pub enable_piecewise_merge_join: bool, default = false /// The maximum estimated size in bytes for one input side of a HashJoin - /// will be collected into a single partition + /// will be collected into a single partition. Partitioned joins also use + /// it: an input below it is not used as the co-partitioning reference for + /// an input above it. pub hash_join_single_partition_threshold: usize, default = 4 * 1024 * 1024 /// The maximum estimated size in rows for one input side of a HashJoin - /// will be collected into a single partition + /// will be collected into a single partition. Partitioned joins also use + /// it: an input below it is not used as the co-partitioning reference for + /// an input above it. pub hash_join_single_partition_threshold_rows: usize, default = 1024 * 128 /// Maximum size in bytes for the build side of a hash join to be pushed down as an InList expression for dynamic filtering. diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 321cef39e5e44..3c8654eeed838 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -1994,36 +1994,110 @@ fn range_hash_join_repartitions_unpartitioned_side_to_match_range() -> Result<() } #[test] -fn range_hash_join_rejects_smaller_reference() -> Result<()> { - let left = parquet_exec_with_partitioning_and_size( - Partitioning::UnknownPartitioning(4), - Precision::Inexact(100_000), - Precision::Inexact(8_000_000), - ); - let right = parquet_exec_with_partitioning_and_size( - range_partitioning("a", [10, 20, 30], SortOptions::default())?, - Precision::Inexact(100), - Precision::Inexact(8_000), - ); +fn range_join_reference_size_gate() -> Result<()> { + let optimizer = TestConfig::default().config.optimizer; + let bytes = optimizer.hash_join_single_partition_threshold; + let rows = optimizer.hash_join_single_partition_threshold_rows; + let known = |n: usize| Precision::Inexact(n); + let absent = Precision::Absent; + + // (name, other (rows, bytes), reference (rows, bytes), expect hash fallback) + let cases = [ + ( + "bytes: straddle", + (absent, known(bytes)), + (absent, known(bytes - 1)), + true, + ), + ( + "bytes: both at or above", + (absent, known(bytes + 1)), + (absent, known(bytes)), + false, + ), + ( + "bytes: both below", + (absent, known(100)), + (absent, known(100)), + false, + ), + ( + "rows: straddle", + (known(rows), absent), + (known(rows - 1), absent), + true, + ), + ( + "rows: both at or above", + (known(rows + 1), absent), + (known(rows), absent), + false, + ), + ( + "bytes win over rows", + (known(rows * 10), known(100)), + (known(1), known(100)), + false, + ), + ( + "other unknown", + (absent, absent), + (known(1), known(100)), + false, + ), + ( + "reference unknown", + (known(rows * 10), known(bytes * 10)), + (absent, absent), + false, + ), + ]; - let join_on = vec![( - Arc::new(Column::new_with_schema("a", &left.schema())?) as _, - Arc::new(Column::new_with_schema("a", &right.schema())?) as _, - )]; - let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); + type MakeJoin = fn( + Arc, + Arc, + &JoinOn, + &JoinType, + ) -> Arc; + let joins: [(&str, MakeJoin); 2] = [ + ("hash", hash_join_exec), + ("sort-merge", sort_merge_join_exec), + ]; - let plan = TestConfig::default() - .with_query_execution_partitions(4) - .to_plan(join, &DISTRIB_DISTRIB_SORT); + for (join_name, make_join) in joins { + for (name, other, reference, expect_hash) in cases { + let left = parquet_exec_with_partitioning_and_size( + Partitioning::UnknownPartitioning(4), + other.0, + other.1, + ); + let right = parquet_exec_with_partitioning_and_size( + range_partitioning("a", [10, 20, 30], SortOptions::default())?, + reference.0, + reference.1, + ); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = make_join(left, right, &join_on, &JoinType::Inner); - let children = plan.children(); - assert_eq!(children.len(), 2); - for child in children { - assert!( - matches!(child.output_partitioning(), Partitioning::Hash(_, 4)), - "expected hash fallback, got {:?}", - child.output_partitioning(), - ); + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + for child in plan.children() { + let is_hash = + matches!(child.output_partitioning(), Partitioning::Hash(_, 4)); + let is_range = + matches!(child.output_partitioning(), Partitioning::Range(_)); + assert!( + if expect_hash { is_hash } else { is_range }, + "{join_name}, {name}: got {:?}", + child.output_partitioning(), + ); + } + } } Ok(()) @@ -2076,76 +2150,6 @@ fn range_requirement_scales_within_sample_resolution() -> Result<()> { Ok(()) } -#[test] -fn range_hash_join_preserves_equal_reference() -> Result<()> { - let left = parquet_exec_with_partitioning_and_size( - Partitioning::UnknownPartitioning(4), - Precision::Inexact(100), - Precision::Inexact(8_000), - ); - let right = parquet_exec_with_partitioning_and_size( - range_partitioning("a", [10, 20, 30], SortOptions::default())?, - Precision::Inexact(100), - Precision::Inexact(8_000), - ); - - let join_on = vec![( - Arc::new(Column::new_with_schema("a", &left.schema())?) as _, - Arc::new(Column::new_with_schema("a", &right.schema())?) as _, - )]; - let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); - - let plan = TestConfig::default() - .with_query_execution_partitions(4) - .to_plan(join, &DISTRIB_DISTRIB_SORT); - - let expected = range_partitioning("a", [10, 20, 30], SortOptions::default())?; - - let children = plan.children(); - assert_eq!(children.len(), 2); - for child in children { - assert_eq!(child.output_partitioning(), &expected); - } - - Ok(()) -} - -#[test] -fn range_hash_join_preserves_large_reference() -> Result<()> { - // Both inputs are above the `CollectLeft` thresholds, so the slightly - // larger one is adapted to the reference and only it is shuffled. - let left = parquet_exec_with_partitioning_and_size( - Partitioning::UnknownPartitioning(4), - Precision::Inexact(11_000_000), - Precision::Inexact(880_000_000), - ); - let right = parquet_exec_with_partitioning_and_size( - range_partitioning("a", [10, 20, 30], SortOptions::default())?, - Precision::Inexact(10_000_000), - Precision::Inexact(800_000_000), - ); - - let join_on = vec![( - Arc::new(Column::new_with_schema("a", &left.schema())?) as _, - Arc::new(Column::new_with_schema("a", &right.schema())?) as _, - )]; - let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); - - let plan = TestConfig::default() - .with_query_execution_partitions(4) - .to_plan(join, &DISTRIB_DISTRIB_SORT); - - let expected = range_partitioning("a", [10, 20, 30], SortOptions::default())?; - - let children = plan.children(); - assert_eq!(children.len(), 2); - for child in children { - assert_eq!(child.output_partitioning(), &expected); - } - - Ok(()) -} - #[test] fn range_singleton_scaling_respects_required_keys() -> Result<()> { let ordering = [PhysicalSortExpr::new_default(col("a", &schema())?)].into(); diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 1e4cfe9472fc3..89ac19f683793 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -19,6 +19,7 @@ use std::any::Any; use std::fmt::Debug; use std::sync::Arc; +use crate::join_selection::below_collect_thresholds; use crate::output_requirements::OutputRequirementExec; use crate::utils::{ add_sort_above_with_check, is_coalesce_partitions, is_repartition, @@ -1230,21 +1231,6 @@ impl PlanSize { fn is_known(&self) -> bool { self.byte_size.is_some() || self.num_rows.is_some() } - - /// Whether the plan is below the `CollectLeft` thresholds of `JoinSelection`, - /// comparing bytes when known and rows otherwise. `None` when the size is unknown. - fn is_small(&self, config: &ConfigOptions) -> Option { - let optimizer = &config.optimizer; - match (self.byte_size, self.num_rows) { - (Some(bytes), _) => { - Some(bytes < optimizer.hash_join_single_partition_threshold) - } - (None, Some(rows)) => { - Some(rows < optimizer.hash_join_single_partition_threshold_rows) - } - (None, None) => None, - } - } } impl PartialOrd for PlanSize { @@ -1382,18 +1368,20 @@ fn enforce_distribution_relationships( // below them; fall back to hash partitioning instead. Otherwise keep the // reference, so that it is not shuffled as well. let best_satisfied_child = best_satisfied_child.filter(|(ref_idx, _)| { - let reference_size = - PlanSize::from_plan(&children[*ref_idx].context.plan, stats_ctx); - if reference_size.is_small(config) != Some(true) { - return true; - } - - !unsatisfied_children.iter().any(|&child_idx| { - child_idx != *ref_idx - && PlanSize::from_plan(&children[child_idx].context.plan, stats_ctx) - .is_small(config) - == Some(false) - }) + let is_small = |idx: usize| { + let stats = stats_ctx + .compute_arc(&children[idx].context.plan, &StatisticsArgs::new()) + .ok()?; + below_collect_thresholds( + &stats, + config.optimizer.hash_join_single_partition_threshold, + config.optimizer.hash_join_single_partition_threshold_rows, + ) + }; + is_small(*ref_idx) != Some(true) + || !unsatisfied_children + .iter() + .any(|&child_idx| is_small(child_idx) == Some(false)) }); // Validate that best_satisfied_child can be adapted across all unsatisfied children in the group diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs index b8a3aefbe8f6e..6a5df4bfe3963 100644 --- a/datafusion/physical-optimizer/src/join_selection.rs +++ b/datafusion/physical-optimizer/src/join_selection.rs @@ -111,12 +111,24 @@ fn supports_collect_by_thresholds( // `Exact(0)` and `Inexact(0)` are both valid stats, and we should not treat // them as unknown, `Absent` will return None (this is in regards to why // `!=0` is not checked) + below_collect_thresholds(&stats, threshold_byte_size, threshold_num_rows) + .unwrap_or(false) +} + +/// Whether `stats` is below the `CollectLeft` thresholds, comparing bytes when +/// known and rows otherwise. `None` when the size is unknown. +pub(crate) fn below_collect_thresholds( + stats: &Statistics, + threshold_byte_size: usize, + threshold_num_rows: usize, +) -> Option { if let Some(byte_size) = stats.total_byte_size.get_value() { - *byte_size < threshold_byte_size - } else if let Some(num_rows) = stats.num_rows.get_value() { - *num_rows < threshold_num_rows + Some(*byte_size < threshold_byte_size) } else { - false + stats + .num_rows + .get_value() + .map(|num_rows| *num_rows < threshold_num_rows) } } diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 4355051ef0254..e61a8bf36b335 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -488,8 +488,8 @@ datafusion.optimizer.expand_views_at_output false When set to true, if the retur datafusion.optimizer.filter_null_join_keys false When set to true, the optimizer will insert filters before a join between a nullable and non-nullable column to filter out nulls on the nullable side. This filter can add additional overhead when the file format does not fully support predicate push down. datafusion.optimizer.hash_join_inlist_pushdown_max_distinct_values 150 Maximum number of distinct values (rows) in the build side of a hash join to be pushed down as an InList expression for dynamic filtering. Build sides with more rows than this will use hash table lookups instead. Set to 0 to always use hash table lookups. This provides an additional limit beyond `hash_join_inlist_pushdown_max_size` to prevent very large IN lists that might not provide much benefit over hash table lookups. This uses the deduplicated row count once the build side has been evaluated. The default is 150 values per partition. This is inspired by Trino's `max-filter-keys-per-column` setting. See: datafusion.optimizer.hash_join_inlist_pushdown_max_size 131072 Maximum size in bytes for the build side of a hash join to be pushed down as an InList expression for dynamic filtering. Build sides larger than this will use hash table lookups instead. Set to 0 to always use hash table lookups. InList pushdown can be more efficient for small build sides because it can result in better statistics pruning as well as use any bloom filters present on the scan side. InList expressions are also more transparent and easier to serialize over the network in distributed uses of DataFusion. On the other hand InList pushdown requires making a copy of the data and thus adds some overhead to the build side and uses more memory. This setting is per-partition, so we may end up using `hash_join_inlist_pushdown_max_size` * `target_partitions` memory. The default is 128kB per partition. This should allow point lookup joins (e.g. joining on a unique primary key) to use InList pushdown in most cases but avoids excessive memory usage or overhead for larger joins. -datafusion.optimizer.hash_join_single_partition_threshold 4194304 The maximum estimated size in bytes for one input side of a HashJoin will be collected into a single partition -datafusion.optimizer.hash_join_single_partition_threshold_rows 131072 The maximum estimated size in rows for one input side of a HashJoin will be collected into a single partition +datafusion.optimizer.hash_join_single_partition_threshold 4194304 The maximum estimated size in bytes for one input side of a HashJoin will be collected into a single partition. Partitioned joins also use it: an input below it is not used as the co-partitioning reference for an input above it. +datafusion.optimizer.hash_join_single_partition_threshold_rows 131072 The maximum estimated size in rows for one input side of a HashJoin will be collected into a single partition. Partitioned joins also use it: an input below it is not used as the co-partitioning reference for an input above it. datafusion.optimizer.join_reordering true When set to true, the physical plan optimizer may swap join inputs based on statistics. When set to false, statistics-driven join input reordering is disabled and the original join order in the query is used. datafusion.optimizer.max_passes 3 The logical optimizer applies rules in order for up to `max_passes` passes, stopping early if a pass leaves the plan unchanged or repeats an earlier plan. The physical optimizer always runs in one pass. datafusion.optimizer.prefer_existing_sort false When true, DataFusion will opportunistically remove sorts when the data is already sorted, (i.e. setting `preserve_order` to true on `RepartitionExec` and using `SortPreservingMergeExec`) When false, DataFusion will maximize plan parallelism using `RepartitionExec` even if this requires subsequently resorting data using a `SortExec`. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 445c63a05c93f..3e4af2d9a9356 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -167,8 +167,8 @@ The following configuration settings are available: | datafusion.optimizer.use_statistics_registry | false | (Deprecated) Ignored: the physical plan optimizer always consults the session's pluggable `StatisticsRegistry` (register providers on the `SessionState`; with none it is a no-op). | | datafusion.optimizer.prefer_hash_join | true | When set to true, the physical plan optimizer will prefer HashJoin over SortMergeJoin. HashJoin can work more efficiently than SortMergeJoin but consumes more memory | | datafusion.optimizer.enable_piecewise_merge_join | false | When set to true, piecewise merge join is enabled. PiecewiseMergeJoin is currently experimental. Physical planner will opt for PiecewiseMergeJoin when there is only one range filter. | -| datafusion.optimizer.hash_join_single_partition_threshold | 4194304 | The maximum estimated size in bytes for one input side of a HashJoin will be collected into a single partition | -| datafusion.optimizer.hash_join_single_partition_threshold_rows | 131072 | The maximum estimated size in rows for one input side of a HashJoin will be collected into a single partition | +| datafusion.optimizer.hash_join_single_partition_threshold | 4194304 | The maximum estimated size in bytes for one input side of a HashJoin will be collected into a single partition. Partitioned joins also use it: an input below it is not used as the co-partitioning reference for an input above it. | +| datafusion.optimizer.hash_join_single_partition_threshold_rows | 131072 | The maximum estimated size in rows for one input side of a HashJoin will be collected into a single partition. Partitioned joins also use it: an input below it is not used as the co-partitioning reference for an input above it. | | datafusion.optimizer.hash_join_inlist_pushdown_max_size | 131072 | Maximum size in bytes for the build side of a hash join to be pushed down as an InList expression for dynamic filtering. Build sides larger than this will use hash table lookups instead. Set to 0 to always use hash table lookups. InList pushdown can be more efficient for small build sides because it can result in better statistics pruning as well as use any bloom filters present on the scan side. InList expressions are also more transparent and easier to serialize over the network in distributed uses of DataFusion. On the other hand InList pushdown requires making a copy of the data and thus adds some overhead to the build side and uses more memory. This setting is per-partition, so we may end up using `hash_join_inlist_pushdown_max_size` \* `target_partitions` memory. The default is 128kB per partition. This should allow point lookup joins (e.g. joining on a unique primary key) to use InList pushdown in most cases but avoids excessive memory usage or overhead for larger joins. | | datafusion.optimizer.hash_join_inlist_pushdown_max_distinct_values | 150 | Maximum number of distinct values (rows) in the build side of a hash join to be pushed down as an InList expression for dynamic filtering. Build sides with more rows than this will use hash table lookups instead. Set to 0 to always use hash table lookups. This provides an additional limit beyond `hash_join_inlist_pushdown_max_size` to prevent very large IN lists that might not provide much benefit over hash table lookups. This uses the deduplicated row count once the build side has been evaluated. The default is 150 values per partition. This is inspired by Trino's `max-filter-keys-per-column` setting. See: | | datafusion.optimizer.default_filter_selectivity | 20 | The default filter selectivity used by Filter Statistics when an exact selectivity cannot be determined. Valid values are between 0 (no selectivity) and 100 (all rows are selected). |