diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index f5742f09f9b08..ea5db1dea9c82 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1672,6 +1672,22 @@ config_namespace! { /// query is used. pub join_reordering: bool, default = true + /// When set to true, the physical plan optimizer enumerates join orders for + /// subtrees of joins and picks the cheapest from cardinality estimates, + /// considering bushy shapes as well as left-deep ones. Subtrees whose inputs + /// lack row count statistics are left untouched. + pub join_enumeration: bool, default = true + + /// How much cheaper an enumerated join order must be, in percent, before it + /// replaces the order the planner produced. Estimates are often too close to tell + /// two orders apart, so a small gain is not worth acting on. + pub join_enumeration_min_improvement: u8, default = 10 + + /// Maximum inputs in a join subtree for which `join_enumeration` searches, at a + /// cost of `O(3^n)`. Larger subtrees keep the planner's order, as do subtrees of + /// more than 16 inputs regardless of this setting. + pub join_enumeration_limit: usize, default = 12 + /// When set to true, the physical plan optimizer uses the pluggable /// `StatisticsRegistry` for statistics propagation across operators. /// This enables more accurate cardinality estimates compared to each @@ -1689,7 +1705,7 @@ config_namespace! { /// The maximum estimated size in bytes for one input side of a HashJoin /// will be collected into a single partition - pub hash_join_single_partition_threshold: usize, default = 1024 * 1024 + 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 diff --git a/datafusion/core/src/optimizer_rule_reference.md b/datafusion/core/src/optimizer_rule_reference.md index 1367ed0843c59..e4bb2350958e3 100644 --- a/datafusion/core/src/optimizer_rule_reference.md +++ b/datafusion/core/src/optimizer_rule_reference.md @@ -72,22 +72,23 @@ in multiple phases. | ----- | ------------------------------ | ----------------------- | ------------------------------------------------------------------------------------------------------------ | | 1 | `OutputRequirements` | add phase | Adds helper nodes so output requirements survive later physical rewrites. | | 2 | `aggregate_statistics` | - | Uses exact source statistics to answer some aggregates without scanning data. | -| 3 | `join_selection` | - | Chooses join implementation, build side, and partition mode from statistics and stream properties. | -| 4 | `LimitedDistinctAggregation` | - | Pushes limit hints into grouped distinct-style aggregations when only a small result is needed. | -| 5 | `FilterPushdown` | pre-optimization phase | Pushes supported physical filters down toward data sources before distribution and sorting are enforced. | -| 6 | `WindowTopN` | - | Replaces eligible row-number window and filter patterns with per-partition TopK execution. | -| 7 | `EnsureRequirements` | - | Enforces both distribution and sorting requirements in a single idempotent rule. | -| 8 | `CombinePartialFinalAggregate` | - | Collapses adjacent partial and final aggregates when the distributed shape makes them redundant. | -| 9 | `OptimizeAggregateOrder` | - | Updates aggregate expressions to use the best ordering once sort requirements are known. | -| 10 | `ProjectionPushdown` | early pass | Pushes projections toward inputs before later physical rewrites add more limit and TopK structure. | -| 11 | `OutputRequirements` | remove phase | Removes the temporary output-requirement helper nodes after requirement-sensitive planning is done. | -| 12 | `LimitAggregation` | - | Passes a limit hint into eligible aggregations so they can keep fewer accumulator buckets. | -| 13 | `LimitPushPastWindows` | - | Pushes fetch limits through bounded window operators when doing so keeps the result correct. | -| 14 | `HashJoinBuffering` | - | Adds buffering on the probe side of hash joins so probing can start before build completion. | -| 15 | `LimitPushdown` | - | Moves physical limits into child operators or fetch-enabled variants to cut data early. | -| 16 | `TopKRepartition` | - | Pushes TopK below hash repartition when the partition key is a prefix of the sort key. | -| 17 | `ProjectionPushdown` | late pass | Runs projection pushdown again after limit and TopK rewrites expose new pruning opportunities. | -| 18 | `PushdownSort` | - | Pushes sort requirements into data sources that can already return sorted output. | -| 19 | `EnsureCooperative` | - | Wraps non-cooperative plan parts so long-running tasks yield fairly. | -| 20 | `FilterPushdown(Post)` | post-optimization phase | Pushes dynamic filters at the end of optimization, after plan references stop moving. | -| 21 | `SanityCheckPlan` | - | Validates that the final physical plan meets ordering, distribution, and infinite-input safety requirements. | +| 3 | `join_enumeration` | - | Chooses the join tree shape by costing alternative orders from statistics. | +| 4 | `join_selection` | - | Chooses join implementation, build side, and partition mode from statistics and stream properties. | +| 5 | `LimitedDistinctAggregation` | - | Pushes limit hints into grouped distinct-style aggregations when only a small result is needed. | +| 6 | `FilterPushdown` | pre-optimization phase | Pushes supported physical filters down toward data sources before distribution and sorting are enforced. | +| 7 | `WindowTopN` | - | Replaces eligible row-number window and filter patterns with per-partition TopK execution. | +| 8 | `EnsureRequirements` | - | Enforces both distribution and sorting requirements in a single idempotent rule. | +| 9 | `CombinePartialFinalAggregate` | - | Collapses adjacent partial and final aggregates when the distributed shape makes them redundant. | +| 10 | `OptimizeAggregateOrder` | - | Updates aggregate expressions to use the best ordering once sort requirements are known. | +| 11 | `ProjectionPushdown` | early pass | Pushes projections toward inputs before later physical rewrites add more limit and TopK structure. | +| 12 | `OutputRequirements` | remove phase | Removes the temporary output-requirement helper nodes after requirement-sensitive planning is done. | +| 13 | `LimitAggregation` | - | Passes a limit hint into eligible aggregations so they can keep fewer accumulator buckets. | +| 14 | `LimitPushPastWindows` | - | Pushes fetch limits through bounded window operators when doing so keeps the result correct. | +| 15 | `HashJoinBuffering` | - | Adds buffering on the probe side of hash joins so probing can start before build completion. | +| 16 | `LimitPushdown` | - | Moves physical limits into child operators or fetch-enabled variants to cut data early. | +| 17 | `TopKRepartition` | - | Pushes TopK below hash repartition when the partition key is a prefix of the sort key. | +| 18 | `ProjectionPushdown` | late pass | Runs projection pushdown again after limit and TopK rewrites expose new pruning opportunities. | +| 19 | `PushdownSort` | - | Pushes sort requirements into data sources that can already return sorted output. | +| 20 | `EnsureCooperative` | - | Wraps non-cooperative plan parts so long-running tasks yield fairly. | +| 21 | `FilterPushdown(Post)` | post-optimization phase | Pushes dynamic filters at the end of optimization, after plan references stop moving. | +| 22 | `SanityCheckPlan` | - | Validates that the final physical plan meets ordering, distribution, and infinite-input safety requirements. | diff --git a/datafusion/core/tests/physical_optimizer/join_enumeration.rs b/datafusion/core/tests/physical_optimizer/join_enumeration.rs new file mode 100644 index 0000000000000..231fd6b4ccd44 --- /dev/null +++ b/datafusion/core/tests/physical_optimizer/join_enumeration.rs @@ -0,0 +1,621 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Tests for the cost-based join order enumeration in [`JoinSelection`]. + +use std::sync::Arc; + +use arrow::array::{Int32Array, RecordBatch}; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::util::pretty::pretty_format_batches; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::config::ConfigOptions; +use datafusion_common::{ + ColumnStatistics, JoinSide, JoinType, NullEquality, Result, ScalarValue, +}; +use datafusion_common::{Statistics, stats::Precision}; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::BinaryExpr; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_optimizer::PhysicalOptimizerRule; +use datafusion_physical_optimizer::join_enumeration::graph::{ + JoinGraph, RelSet, iter_rels, +}; +use datafusion_physical_optimizer::join_enumeration::{ + Combine, DefaultJoinCostModel, Exchange, JoinCostModel, JoinCostModelFactory, + JoinEnumeration, PartSet, +}; +use datafusion_physical_optimizer::join_selection::JoinSelection; +use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter}; +use datafusion_physical_plan::joins::{ + CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, +}; +use datafusion_physical_plan::{ExecutionPlan, displayable}; +use insta::assert_snapshot; + +use crate::physical_optimizer::join_selection::StatisticsExec; + +/// A table of `rows` rows with `(name, distinct_count)` columns. Each gets a +/// `[0, distinct_count)` range, as a real scan with min/max statistics would. +fn table(rows: usize, columns: &[(&str, usize)]) -> (Statistics, Schema) { + let column_statistics = columns + .iter() + .map(|(_, distinct)| ColumnStatistics { + distinct_count: Precision::Inexact(*distinct), + min_value: Precision::Inexact(ScalarValue::Int32(Some(0))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(*distinct as i32 - 1))), + ..Default::default() + }) + .collect(); + let schema = Schema::new( + columns + .iter() + .map(|(name, _)| Field::new(*name, DataType::Int32, false)) + .collect::>(), + ); + ( + Statistics { + num_rows: Precision::Inexact(rows), + total_byte_size: Precision::Absent, + column_statistics, + }, + schema, + ) +} + +fn scan(rows: usize, columns: &[(&str, usize)]) -> Arc { + let (statistics, schema) = table(rows, columns); + Arc::new(StatisticsExec::new(statistics, schema)) +} + +fn scan_without_statistics(columns: &[&str]) -> Arc { + let schema = Schema::new( + columns + .iter() + .map(|name| Field::new(*name, DataType::Int32, false)) + .collect::>(), + ); + let statistics = Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Absent, + column_statistics: vec![ColumnStatistics::new_unknown(); columns.len()], + }; + Arc::new(StatisticsExec::new(statistics, schema)) +} + +fn join( + left: Arc, + right: Arc, + on: &[(&str, &str)], +) -> Result> { + join_of_type(left, right, on, JoinType::Inner, None) +} + +/// The same shape as [`late_reducer_plan`], joined by sort merge instead of hash. +fn sort_merge_late_reducer_plan() -> Result> { + let fact = scan(1_000_000, &[("f_id", 1_000_000), ("f_type", 1_000)]); + let other = scan(1_000_000, &[("o_id", 1_000_000)]); + let types = scan(10, &[("t_type", 10)]); + + let joined = sort_merge_join(fact, other, &[("f_id", "o_id")])?; + sort_merge_join(joined, types, &[("f_type", "t_type")]) +} + +fn sort_merge_join( + left: Arc, + right: Arc, + on: &[(&str, &str)], +) -> Result> { + let keys = on + .iter() + .map(|(left_key, right_key)| { + Ok(( + Arc::new(Column::new_with_schema(left_key, &left.schema())?) as _, + Arc::new(Column::new_with_schema(right_key, &right.schema())?) as _, + )) + }) + .collect::>>()?; + Ok(Arc::new(SortMergeJoinExec::try_new( + left, + right, + keys, + None, + JoinType::Inner, + vec![SortOptions::default(); on.len()], + NullEquality::NullEqualsNothing, + )?)) +} + +fn join_of_type( + left: Arc, + right: Arc, + on: &[(&str, &str)], + join_type: JoinType, + filter: Option, +) -> Result> { + let keys = on + .iter() + .map(|(left_key, right_key)| { + Ok(( + Arc::new(Column::new_with_schema(left_key, &left.schema())?) as _, + Arc::new(Column::new_with_schema(right_key, &right.schema())?) as _, + )) + }) + .collect::>>()?; + Ok(Arc::new(HashJoinExec::try_new( + left, + right, + keys, + filter, + &join_type, + None, + PartitionMode::Auto, + NullEquality::NullEqualsNothing, + false, + )?)) +} + +/// A `left_col > right_col` filter over one column of each side. +fn greater_than_filter( + left_col: (&str, usize), + right_col: (&str, usize), +) -> Result { + let schema = Schema::new(vec![ + Field::new(left_col.0, DataType::Int32, false), + Field::new(right_col.0, DataType::Int32, false), + ]); + let expression = Arc::new(BinaryExpr::new( + Arc::new(Column::new(left_col.0, 0)), + Operator::Gt, + Arc::new(Column::new(right_col.0, 1)), + )); + Ok(JoinFilter::new( + expression, + vec![ + ColumnIndex { + index: left_col.1, + side: JoinSide::Left, + }, + ColumnIndex { + index: right_col.1, + side: JoinSide::Right, + }, + ], + Arc::new(schema), + )) +} + +/// A three way join in its expensive `FROM` order: the two large tables first produce a +/// million rows, where either of them with the tiny table first gives ten thousand. +fn late_reducer_plan() -> Result> { + let fact = scan(1_000_000, &[("f_id", 1_000_000), ("f_type", 1_000)]); + let other = scan(1_000_000, &[("o_id", 1_000_000)]); + let types = scan(10, &[("t_type", 10)]); + + let fact_other = join(fact, other, &[("f_id", "o_id")])?; + join(fact_other, types, &[("f_type", "t_type")]) +} + +/// Both rules in pipeline order: shape first, then build side and partition mode. +fn optimize( + plan: Arc, + config: &ConfigOptions, +) -> Result> { + let plan = JoinEnumeration::new().optimize(plan, config)?; + JoinSelection::new().optimize(plan, config) +} + +fn formatted(plan: &Arc) -> String { + displayable(plan.as_ref()).indent(true).to_string() +} + +#[test] +fn reorders_a_late_reducer() -> Result<()> { + let plan = late_reducer_plan()?; + // The planner's order: the two million row tables are joined first. + assert_snapshot!(formatted(&plan), @r" + HashJoinExec: mode=Auto, join_type=Inner, on=[(f_type@1, t_type@0)] + HashJoinExec: mode=Auto, join_type=Inner, on=[(f_id@0, o_id@0)] + StatisticsExec: col_count=2, row_count=Inexact(1000000) + StatisticsExec: col_count=1, row_count=Inexact(1000000) + StatisticsExec: col_count=1, row_count=Inexact(10) + "); + + // The reducing join moves down so the large tables never join directly. + assert_snapshot!(formatted(&optimize(plan, &ConfigOptions::new())?), @r" + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(f_id@0, o_id@0)], projection=[f_id@0, f_type@1, o_id@3, t_type@2] + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(t_type@0, f_type@1)], projection=[f_id@1, f_type@2, t_type@0] + StatisticsExec: col_count=1, row_count=Inexact(10) + StatisticsExec: col_count=2, row_count=Inexact(1000000) + StatisticsExec: col_count=1, row_count=Inexact(1000000) + "); + Ok(()) +} + +#[test] +fn respects_the_config_flag() -> Result<()> { + let mut config = ConfigOptions::new(); + config.optimizer.join_enumeration = false; + let optimized = optimize(late_reducer_plan()?, &config)?; + // Without enumeration the large tables still join first, for a million rows. + assert_snapshot!(formatted(&optimized), @r" + ProjectionExec: expr=[f_id@1 as f_id, f_type@2 as f_type, o_id@3 as o_id, t_type@0 as t_type] + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(t_type@0, f_type@1)] + StatisticsExec: col_count=1, row_count=Inexact(10) + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(f_id@0, o_id@0)] + StatisticsExec: col_count=2, row_count=Inexact(1000000) + StatisticsExec: col_count=1, row_count=Inexact(1000000) + "); + Ok(()) +} + +#[test] +fn leaves_plans_without_statistics_alone() -> Result<()> { + let fact = scan_without_statistics(&["f_id", "f_type"]); + let other = scan_without_statistics(&["o_id"]); + let types = scan_without_statistics(&["t_type"]); + let plan = join( + join(fact, other, &[("f_id", "o_id")])?, + types, + &[("f_type", "t_type")], + )?; + + let optimized = optimize(Arc::clone(&plan), &ConfigOptions::new())?; + assert_snapshot!(formatted(&optimized), @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(f_type@1, t_type@0)] + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(f_id@0, o_id@0)] + StatisticsExec: col_count=2, row_count=Absent + StatisticsExec: col_count=1, row_count=Absent + StatisticsExec: col_count=1, row_count=Absent + "); + Ok(()) +} + +#[test] +fn keeps_an_already_optimal_order() -> Result<()> { + // Already in the cheap order, so the enumerator must not churn the plan. + let fact = scan(1_000_000, &[("f_id", 1_000_000), ("f_type", 1_000)]); + let other = scan(1_000_000, &[("o_id", 1_000_000)]); + let types = scan(10, &[("t_type", 10)]); + + let reduced = join(fact, types, &[("f_type", "t_type")])?; + let plan = join(reduced, other, &[("f_id", "o_id")])?; + + let mut disabled = ConfigOptions::new(); + disabled.optimizer.join_enumeration = false; + assert_eq!( + formatted(&optimize(Arc::clone(&plan), &ConfigOptions::new())?), + formatted(&optimize(plan, &disabled)?), + ); + Ok(()) +} + +/// Row counts from outside the plan, keyed by the first column a relation emits. The +/// tiny dimension table is claimed to be the largest of the three. +fn external_rows(column: &str) -> f64 { + match column { + "t_type" => 1e9, + _ => 1e6, + } +} + +/// Hands out [`ExternalStatisticsCostModel`]. +#[derive(Debug)] +struct ExternalStatistics {} + +impl JoinCostModelFactory for ExternalStatistics { + fn create<'graph>( + &self, + graph: &'graph JoinGraph, + config: &ConfigOptions, + ) -> Result> { + Ok(Box::new(ExternalStatisticsCostModel { + graph, + inner: DefaultJoinCostModel::new(graph, config), + })) + } +} + +/// Counts rows from statistics the plan does not carry, and leaves the rest of the model +/// to the built-in one. +struct ExternalStatisticsCostModel<'a> { + graph: &'a JoinGraph, + inner: DefaultJoinCostModel<'a>, +} + +impl JoinCostModel for ExternalStatisticsCostModel<'_> { + fn cardinality(&self, mask: RelSet) -> f64 { + iter_rels(mask) + .map(|rel| { + let schema = self.graph.relations[rel].plan.schema(); + external_rows(schema.field(0).name()) + }) + .product::() + .max(1.0) + } + + fn combine(&self, left: RelSet, right: RelSet) -> Option { + self.inner.combine(left, right) + } + + fn exchanges( + &self, + left: RelSet, + right: RelSet, + left_part: PartSet, + right_part: PartSet, + collect_only: Option, + ) -> Vec { + self.inner + .exchanges(left, right, left_part, right_part, collect_only) + } +} + +#[test] +fn searches_under_a_plugged_in_cost_model() -> Result<()> { + let plan = late_reducer_plan()?; + // Under the plan's own statistics this order is the expensive one, and + // `reorders_a_late_reducer` shows the built-in model rewriting it. Told the dimension + // table is the largest of the three, the search keeps it. + let enumerated = JoinEnumeration::new() + .with_cost_model(Arc::new(ExternalStatistics {})) + .optimize(Arc::clone(&plan), &ConfigOptions::new())?; + assert_eq!(formatted(&enumerated), formatted(&plan)); + Ok(()) +} + +/// A session over four in-memory tables shaped like a small star schema. +fn star_schema_context( + join_enumeration: bool, + prefer_hash_join: bool, +) -> Result { + let mut config = SessionConfig::new(); + config.options_mut().optimizer.join_enumeration = join_enumeration; + config.options_mut().optimizer.prefer_hash_join = prefer_hash_join; + let ctx = SessionContext::new_with_config(config); + + let ints = |name: &str, values: Vec| -> Result { + Ok(RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new(name, DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(values))], + )?) + }; + + // fact(f_id, f_type, f_region), 240 rows. + let fact = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("f_id", DataType::Int32, false), + Field::new("f_type", DataType::Int32, false), + Field::new("f_region", DataType::Int32, false), + ])), + vec![ + Arc::new(Int32Array::from((0..240).collect::>())), + Arc::new(Int32Array::from( + (0..240).map(|i| i % 12).collect::>(), + )), + Arc::new(Int32Array::from( + (0..240).map(|i| i % 5).collect::>(), + )), + ], + )?; + ctx.register_batch("fact", fact)?; + ctx.register_batch("ids", ints("o_id", (0..240).rev().collect())?)?; + ctx.register_batch("types", ints("t_type", vec![3, 7])?)?; + ctx.register_batch("regions", ints("r_region", vec![1, 2, 4])?)?; + Ok(ctx) +} + +/// A plain join tree, `EXISTS`, `NOT EXISTS`, and a non-equi predicate. +const STAR_QUERIES: [&str; 4] = [ + "select f_id, f_type, f_region from fact, ids, types, regions \ + where f_id = o_id and f_type = t_type and f_region = r_region order by f_id", + "select f_id, f_type from fact where exists \ + (select 1 from types where t_type = f_type) \ + and f_id in (select o_id from ids) order by f_id", + "select f_id, f_type from fact where not exists \ + (select 1 from types where t_type = f_type) \ + and f_id in (select o_id from ids) order by f_id", + "select f_id, f_type, t_type from fact, ids, types \ + where f_id = o_id and f_type = t_type and f_id > t_type order by f_id", +]; + +#[tokio::test] +async fn reordering_returns_the_same_rows() -> Result<()> { + for prefer_hash_join in [true, false] { + reordering_returns_the_same_rows_with(prefer_hash_join).await?; + } + Ok(()) +} + +async fn reordering_returns_the_same_rows_with(prefer_hash_join: bool) -> Result<()> { + let enumerated = star_schema_context(true, prefer_hash_join)?; + let baseline = star_schema_context(false, prefer_hash_join)?; + let mut reordered_any = false; + for query in STAR_QUERIES { + let enumerated_plan = enumerated.sql(query).await?.create_physical_plan().await?; + let baseline_plan = baseline.sql(query).await?.create_physical_plan().await?; + reordered_any |= formatted(&enumerated_plan) != formatted(&baseline_plan); + + let enumerated_rows = enumerated.sql(query).await?.collect().await?; + let baseline_rows = baseline.sql(query).await?.collect().await?; + assert_eq!( + pretty_format_batches(&enumerated_rows)?.to_string(), + pretty_format_batches(&baseline_rows)?.to_string(), + "rows differ for: {query}" + ); + assert!(enumerated_rows.iter().map(|b| b.num_rows()).sum::() > 0); + } + // Rows matching would prove nothing if no plan had changed. + assert!(reordered_any); + Ok(()) +} + +/// A selective semi join sitting above a join of two large tables, as TPC-H q18 +/// has it. +fn late_semi_join_plan(anti: bool) -> Result> { + let fact = scan(1_000_000, &[("f_id", 1_000_000), ("f_type", 1_000)]); + let other = scan(1_000_000, &[("o_id", 1_000_000)]); + // Sized so the reducer keeps one percent either way round: ten of the thousand types + // match for the semi join, all but ten for the anti join. + let wanted = if anti { + scan(990, &[("w_type", 990)]) + } else { + scan(10, &[("w_type", 10)]) + }; + + let joined = join(fact, other, &[("f_id", "o_id")])?; + let join_type = if anti { + JoinType::LeftAnti + } else { + JoinType::LeftSemi + }; + join_of_type(joined, wanted, &[("f_type", "w_type")], join_type, None) +} + +#[test] +fn applies_a_selective_semi_join_first() -> Result<()> { + let plan = late_semi_join_plan(false)?; + assert_snapshot!(formatted(&plan), @r" + HashJoinExec: mode=Auto, join_type=LeftSemi, on=[(f_type@1, w_type@0)] + HashJoinExec: mode=Auto, join_type=Inner, on=[(f_id@0, o_id@0)] + StatisticsExec: col_count=2, row_count=Inexact(1000000) + StatisticsExec: col_count=1, row_count=Inexact(1000000) + StatisticsExec: col_count=1, row_count=Inexact(10) + "); + + // A `RightSemi` filtering the fact table before the inner join. + assert_snapshot!(formatted(&optimize(plan, &ConfigOptions::new())?), @r" + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(f_id@0, o_id@0)] + HashJoinExec: mode=CollectLeft, join_type=RightSemi, on=[(w_type@0, f_type@1)] + StatisticsExec: col_count=1, row_count=Inexact(10) + StatisticsExec: col_count=2, row_count=Inexact(1000000) + StatisticsExec: col_count=1, row_count=Inexact(1000000) + "); + Ok(()) +} + +#[test] +fn applies_an_anti_join_first() -> Result<()> { + let optimized = optimize(late_semi_join_plan(true)?, &ConfigOptions::new())?; + // The anti join is pushed down the same way. + assert_snapshot!(formatted(&optimized), @r" + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(f_id@0, o_id@0)] + HashJoinExec: mode=CollectLeft, join_type=RightAnti, on=[(w_type@0, f_type@1)] + StatisticsExec: col_count=1, row_count=Inexact(990) + StatisticsExec: col_count=2, row_count=Inexact(1000000) + StatisticsExec: col_count=1, row_count=Inexact(1000000) + "); + Ok(()) +} + +#[test] +fn moves_a_non_equi_filter_with_its_join() -> Result<()> { + // `f_type > t_type` rides on the fact/types join, which moves below the join + // with the second large table. + let fact = scan(1_000_000, &[("f_id", 1_000_000), ("f_type", 1_000)]); + let other = scan(1_000_000, &[("o_id", 1_000_000)]); + let types = scan(10, &[("t_type", 10)]); + + let joined = join(fact, other, &[("f_id", "o_id")])?; + let plan = join_of_type( + joined, + types, + &[("f_type", "t_type")], + JoinType::Inner, + Some(greater_than_filter(("f_type", 1), ("t_type", 0))?), + )?; + + // Re-attached to the join that now brings its two columns together. + assert_snapshot!(formatted(&optimize(plan, &ConfigOptions::new())?), @r" + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(f_id@0, o_id@0)], projection=[f_id@0, f_type@1, o_id@3, t_type@2] + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(t_type@0, f_type@1)], filter=f_type@1 > t_type@0, projection=[f_id@1, f_type@2, t_type@0] + StatisticsExec: col_count=1, row_count=Inexact(10) + StatisticsExec: col_count=2, row_count=Inexact(1000000) + StatisticsExec: col_count=1, row_count=Inexact(1000000) + "); + Ok(()) +} + +#[test] +fn reorders_sort_merge_joins() -> Result<()> { + // A sort merge join carries no projection, so the columns the subtree used to + // emit are restored by one projection above it. + assert_snapshot!( + formatted(&optimize(sort_merge_late_reducer_plan()?, &ConfigOptions::new())?), + @r" + ProjectionExec: expr=[f_id@1 as f_id, f_type@2 as f_type, o_id@3 as o_id, t_type@0 as t_type] + SortMergeJoinExec: join_type=Inner, on=[(f_id@1, o_id@0)] + SortMergeJoinExec: join_type=Inner, on=[(t_type@0, f_type@1)] + StatisticsExec: col_count=1, row_count=Inexact(10) + StatisticsExec: col_count=2, row_count=Inexact(1000000) + StatisticsExec: col_count=1, row_count=Inexact(1000000) + " + ); + Ok(()) +} + +#[test] +fn reorders_around_a_cross_join() -> Result<()> { + // Only `types` has a predicate, so `other` can only join by cross product, + // which belongs on the small pair rather than under the reducing join. + let fact = scan(1_000_000, &[("f_id", 1_000_000), ("f_type", 1_000)]); + let other = scan(10, &[("o_id", 10)]); + let types = scan(10, &[("t_type", 10)]); + + let crossed: Arc = Arc::new(CrossJoinExec::new(fact, other)); + let plan = join(crossed, types, &[("f_type", "t_type")])?; + + assert_snapshot!(formatted(&optimize(plan, &ConfigOptions::new())?), @r" + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(t_type@1, f_type@1)], projection=[f_id@2, f_type@3, o_id@0, t_type@1] + CrossJoinExec + StatisticsExec: col_count=1, row_count=Inexact(10) + StatisticsExec: col_count=1, row_count=Inexact(10) + StatisticsExec: col_count=2, row_count=Inexact(1000000) + "); + Ok(()) +} + +#[test] +fn reorders_a_nested_loop_join() -> Result<()> { + // The equi join inflates its inputs a hundredfold, so the non-equi predicate is + // cheaper applied first, even at the default selectivity a filter gets. + let a = scan(1_000, &[("a_k", 10), ("a_t", 1_000)]); + let b = scan(1_000, &[("b_k", 10)]); + let t = scan(10, &[("t_t", 10)]); + + let joined = join(a, b, &[("a_k", "b_k")])?; + let plan: Arc = Arc::new(NestedLoopJoinExec::try_new( + joined, + t, + Some(greater_than_filter(("a_t", 1), ("t_t", 0))?), + &JoinType::Inner, + None, + )?); + + // Enumeration alone: `JoinSelection` swaps nested loop inputs itself, which + // would otherwise look like the reordering under test. + let enumerated = JoinEnumeration::new().optimize(plan, &ConfigOptions::new())?; + assert_snapshot!(formatted(&enumerated), @r" + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(b_k@0, a_k@0)], projection=[a_k@1, a_t@2, b_k@0, t_t@3] + StatisticsExec: col_count=1, row_count=Inexact(1000) + NestedLoopJoinExec: join_type=Inner, filter=a_t@1 > t_t@0, projection=[a_k@1, a_t@2, t_t@0] + StatisticsExec: col_count=1, row_count=Inexact(10) + StatisticsExec: col_count=2, row_count=Inexact(1000) + "); + Ok(()) +} diff --git a/datafusion/core/tests/physical_optimizer/join_selection.rs b/datafusion/core/tests/physical_optimizer/join_selection.rs index 63654ae048863..5bf457ed5d038 100644 --- a/datafusion/core/tests/physical_optimizer/join_selection.rs +++ b/datafusion/core/tests/physical_optimizer/join_selection.rs @@ -75,6 +75,16 @@ fn get_thresholds() -> (usize, usize) { ) } +/// The byte size [`small_statistics`] derives from the configured threshold. +fn small_byte_size() -> usize { + get_thresholds().1 / 128 +} + +/// The byte size [`big_statistics`] derives from the configured threshold. +fn big_byte_size() -> usize { + get_thresholds().1 * 2 +} + /// Return statistics for small table fn small_statistics() -> Statistics { let (threshold_num_rows, threshold_byte_size) = get_thresholds(); @@ -258,14 +268,14 @@ async fn test_join_with_swap() { .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, - Precision::Inexact(8192) + Precision::Inexact(small_byte_size()) ); assert_eq!( StatisticsContext::new() .compute(swapped_join.right().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, - Precision::Inexact(2097152) + Precision::Inexact(big_byte_size()) ); } @@ -382,14 +392,14 @@ async fn test_left_join_no_swap() { .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, - Precision::Inexact(8192) + Precision::Inexact(small_byte_size()) ); assert_eq!( StatisticsContext::new() .compute(swapped_join.right().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, - Precision::Inexact(2097152) + Precision::Inexact(big_byte_size()) ); } @@ -431,14 +441,14 @@ async fn test_join_with_swap_semi() { .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, - Precision::Inexact(8192) + Precision::Inexact(small_byte_size()) ); assert_eq!( StatisticsContext::new() .compute(swapped_join.right().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, - Precision::Inexact(2097152) + Precision::Inexact(big_byte_size()) ); assert_eq!(original_schema, swapped_join.schema()); } @@ -668,14 +678,14 @@ async fn test_join_with_swap_mark() { .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, - Precision::Inexact(8192) + Precision::Inexact(small_byte_size()) ); assert_eq!( StatisticsContext::new() .compute(swapped_join.right().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, - Precision::Inexact(2097152) + Precision::Inexact(big_byte_size()) ); assert_eq!(original_schema, swapped_join.schema()); } @@ -794,14 +804,14 @@ async fn test_join_no_swap() { .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, - Precision::Inexact(8192) + Precision::Inexact(small_byte_size()) ); assert_eq!( StatisticsContext::new() .compute(swapped_join.right().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, - Precision::Inexact(2097152) + Precision::Inexact(big_byte_size()) ); } @@ -867,14 +877,14 @@ async fn test_nl_join_with_swap(join_type: JoinType) { .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, - Precision::Inexact(8192) + Precision::Inexact(small_byte_size()) ); assert_eq!( StatisticsContext::new() .compute(swapped_join.right().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, - Precision::Inexact(2097152) + Precision::Inexact(big_byte_size()) ); } @@ -938,14 +948,14 @@ async fn test_nl_join_with_swap_no_proj(join_type: JoinType) { .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, - Precision::Inexact(8192) + Precision::Inexact(small_byte_size()) ); assert_eq!( StatisticsContext::new() .compute(swapped_join.right().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, - Precision::Inexact(2097152) + Precision::Inexact(big_byte_size()) ); } diff --git a/datafusion/core/tests/physical_optimizer/mod.rs b/datafusion/core/tests/physical_optimizer/mod.rs index f3b2884dab188..0ce5990201d15 100644 --- a/datafusion/core/tests/physical_optimizer/mod.rs +++ b/datafusion/core/tests/physical_optimizer/mod.rs @@ -26,6 +26,7 @@ mod enforce_sorting; mod enforce_sorting_monotonicity; mod ensure_requirements; mod filter_pushdown; +mod join_enumeration; mod join_selection; #[expect(clippy::needless_pass_by_value)] mod limit_pushdown; diff --git a/datafusion/core/tests/physical_optimizer/partition_statistics.rs b/datafusion/core/tests/physical_optimizer/partition_statistics.rs index 6cabcdb710393..6decb7b9092f3 100644 --- a/datafusion/core/tests/physical_optimizer/partition_statistics.rs +++ b/datafusion/core/tests/physical_optimizer/partition_statistics.rs @@ -1596,7 +1596,7 @@ mod test { // For collect left mode, the min/max values are from the entire left table and the specific partition of the right table. let expected_p0_statistics = Statistics { num_rows: Precision::Inexact(2), - total_byte_size: Precision::Absent, + total_byte_size: Precision::Inexact(32), column_statistics: vec![ // Left id column: all partitions (id 1..4) ColumnStatistics { @@ -1677,7 +1677,7 @@ mod test { // For partitioned mode, the min/max values are from the specific partition for each side. let expected_p0_statistics = Statistics { num_rows: Precision::Inexact(2), - total_byte_size: Precision::Absent, + total_byte_size: Precision::Inexact(32), column_statistics: vec![ // Left id column: partition 0 only (id 3..4) ColumnStatistics { @@ -1756,7 +1756,7 @@ mod test { // For auto mode, the min/max values are from the entire left and right tables. let expected_p0_statistics = Statistics { num_rows: Precision::Inexact(4), - total_byte_size: Precision::Absent, + total_byte_size: Precision::Inexact(64), column_statistics: vec![ // Left id column: all partitions (id 1..4) ColumnStatistics { diff --git a/datafusion/physical-optimizer/src/join_enumeration/graph.rs b/datafusion/physical-optimizer/src/join_enumeration/graph.rs new file mode 100644 index 0000000000000..58793ce770d09 --- /dev/null +++ b/datafusion/physical-optimizer/src/join_enumeration/graph.rs @@ -0,0 +1,593 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The join graph: a subtree of joins as relations plus the predicates between them. +//! +//! Extraction flattens a subtree into a [`JoinGraph`], which is what +//! [`JoinCostModel`](super::JoinCostModel) estimates over and what the rebuilt tree is +//! assembled from. + +use std::sync::Arc; + +use datafusion_common::error::Result; +use datafusion_common::{JoinSide, JoinType, NullEquality, Statistics, internal_err}; +use datafusion_physical_expr::PhysicalExprRef; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_plan::joins::utils::{ + ColumnIndex, JoinFilter, max_distinct_count, +}; +use datafusion_physical_plan::joins::{ + CrossJoinExec, HashJoinExec, NestedLoopJoinExec, SortMergeJoinExec, +}; +use datafusion_physical_plan::projection::{ProjectionExec, all_alias_free_columns}; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; + +/// Hard upper bound on the relations in one join graph. The search allocates `2^n` and +/// visits `3^n`, so larger graphs keep the planner's order regardless of the limit. +pub(crate) const MAX_RELATIONS: usize = 16; + +/// Computes a plan node's statistics, shared with the rest of `JoinSelection`. +pub(crate) type StatsFn<'a> = + dyn FnMut(&dyn ExecutionPlan) -> Result> + 'a; + +/// A bitmask over relation indices. +pub type RelSet = u64; + +/// The set holding `rel` alone. +pub fn bit(rel: usize) -> RelSet { + 1u64 << rel +} + +/// The relations in `mask`, lowest index first. +pub fn iter_rels(mask: RelSet) -> impl Iterator { + std::iter::successors(Some(mask), |m| Some(m & m.wrapping_sub(1))) + .take_while(|m| *m != 0) + .map(|m| m.trailing_zeros() as usize) +} + +/// Whether `mask` holds every relation in `required`. +pub fn covers(mask: RelSet, required: RelSet) -> bool { + required & !mask == 0 +} + +/// One column of one relation, tracked instead of a plain index because reordering +/// moves columns to other positions. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub struct ColRef { + /// Index into [`JoinGraph::relations`]. + pub rel: usize, + /// Index into that relation's schema. + pub col: usize, +} + +/// What a relation contributes to the join. +#[derive(Debug)] +pub enum Role { + /// An ordinary input, contributing its columns. + Output, + /// The quantified side of a semi or anti join, which filters instead of + /// contributing columns. + Reducer(Reducer), +} + +/// The quantified side of a semi or anti join, as what it does to the side it filters. +#[derive(Debug)] +pub struct Reducer { + /// `true` for an anti join, which keeps the rows that do *not* match. + pub anti: bool, + /// Keys, as `(column of the filtered side, column index here)`. + pub keys: Vec<(ColRef, usize)>, + /// Relations the keys reference; this reducer applies only to a set covering them. + pub required: RelSet, +} + +/// One leaf of the join graph: a subplan the enumerator does not look inside. +#[derive(Debug)] +pub struct Relation { + /// The subplan this relation stands for. + pub plan: Arc, + /// Estimated row count, clamped to at least 1. + pub rows: f64, + /// Estimated bytes per row, when the input reports a size. + pub width: Option, + /// Per-column distinct value estimate, clamped to `[1, rows]`. + pub ndv: Vec, + /// What it contributes to the join. + pub role: Role, +} + +/// An equi-join predicate `left = right` between two distinct relations. +#[derive(Clone, Copy, Debug)] +pub struct Edge { + /// The column on one side. + pub left: ColRef, + /// The column it is compared against. + pub right: ColRef, +} + +/// A non-equi join predicate, moved along with its column references rewritten. +#[derive(Debug)] +pub struct Filter { + /// The predicate itself. + pub filter: JoinFilter, + /// The column each entry of the filter's intermediate schema comes from. + pub columns: Vec, + /// The relations those columns belong to. + pub required: RelSet, +} + +/// A connected set of joins as relations plus the predicates between them. +#[derive(Debug)] +pub struct JoinGraph { + /// The leaves, addressed by index throughout. + pub relations: Vec, + /// The equi-join predicates between them. + pub edges: Vec, + /// The non-equi predicates between them. + pub filters: Vec, + /// Columns the original subtree emitted; the rebuilt one reproduces this exactly. + pub output: Vec, + /// Null handling shared by the subtree's joins. A join that differs becomes a + /// relation instead. + pub null_equality: Option, + /// The original tree's internal nodes as `(node, one child)`, children first, so + /// the planner's shape can be scored under the same formula as the alternatives. + pub original_nodes: Vec<(RelSet, RelSet)>, + /// The relations that are reducers rather than ordinary inputs. + pub reducers: RelSet, + /// Which join operator the subtree used, and which the rebuild emits. + pub kind: Option, +} + +impl JoinGraph { + /// Distinct values estimated for one column. + pub fn ndv(&self, col: ColRef) -> f64 { + self.relations[col.rel].ndv[col.col] + } + + /// Every relation in the graph. + pub fn all(&self) -> RelSet { + (0..self.relations.len()).fold(0, |mask, rel| mask | bit(rel)) + } + + /// How `rel` reduces the side it filters, if it is a reducer at all. + pub fn reducer(&self, rel: usize) -> Option<&Reducer> { + match &self.relations[rel].role { + Role::Reducer(reducer) => Some(reducer), + Role::Output => None, + } + } + + /// The operator the rebuild emits, defaulting to a hash join. + pub fn kind(&self) -> JoinKind { + self.kind.unwrap_or(JoinKind::Hash) + } + + /// The null handling the rebuild emits, defaulting to `NullEqualsNothing`. + pub fn null_equality(&self) -> NullEquality { + self.null_equality + .unwrap_or(NullEquality::NullEqualsNothing) + } +} + +/// The join operators the rule can flatten and rebuild. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum JoinKind { + /// Rebuilt as a [`HashJoinExec`]. + Hash, + /// Rebuilt as a [`SortMergeJoinExec`], which has no built-in projection, so the + /// subtree gets one projection on top instead. + SortMerge, +} + +/// One join, seen the same way whichever operator implements it. `kind` and +/// `null_equality` are `None` for operators without equi-join keys. +struct JoinView<'a> { + kind: Option, + role: JoinRole, + left: &'a Arc, + right: &'a Arc, + on: &'a [(PhysicalExprRef, PhysicalExprRef)], + filter: Option<&'a JoinFilter>, + null_equality: Option, + projection: Option<&'a [usize]>, +} + +fn join_view(plan: &Arc) -> Option> { + if let Some(join) = plan.downcast_ref::() { + if join.null_aware || join.fetch().is_some() || join.on().is_empty() { + return None; + } + return Some(JoinView { + kind: Some(JoinKind::Hash), + role: join_role(join.join_type(), join.filter().is_some())?, + left: join.left(), + right: join.right(), + on: join.on(), + filter: join.filter(), + null_equality: Some(join.null_equality), + projection: join.projection.as_deref(), + }); + } + if let Some(join) = plan.downcast_ref::() { + if join.on().is_empty() { + return None; + } + return Some(JoinView { + kind: Some(JoinKind::SortMerge), + role: join_role(&join.join_type(), join.filter().is_some())?, + left: join.left(), + right: join.right(), + on: join.on(), + filter: join.filter().as_ref(), + null_equality: Some(join.null_equality()), + projection: None, + }); + } + if let Some(join) = plan.downcast_ref::() { + // A semi or anti variant has no keys to model as a reducer. + if *join.join_type() != JoinType::Inner { + return None; + } + return Some(JoinView { + kind: None, + role: JoinRole::Inner, + left: join.left(), + right: join.right(), + on: &[], + filter: join.filter(), + null_equality: None, + projection: join.projection().as_deref(), + }); + } + let join = plan.downcast_ref::()?; + Some(JoinView { + kind: None, + role: JoinRole::Inner, + left: join.left(), + right: join.right(), + on: &[], + filter: None, + null_equality: None, + projection: None, + }) +} + +/// How a join takes part in enumeration, if at all. +#[derive(Clone, Copy, Debug)] +enum JoinRole { + Inner, + /// A semi or anti join; `output` names the side whose rows survive. + Reducing { + anti: bool, + output: JoinSide, + }, +} + +/// Classifies a join. Outer and mark joins are excluded because they do not just filter +/// their inputs, and semi and anti joins with a filter because it is part of their test. +fn join_role(join_type: &JoinType, has_filter: bool) -> Option { + let role = match join_type { + JoinType::Inner => JoinRole::Inner, + JoinType::LeftSemi => JoinRole::Reducing { + anti: false, + output: JoinSide::Left, + }, + JoinType::RightSemi => JoinRole::Reducing { + anti: false, + output: JoinSide::Right, + }, + JoinType::LeftAnti => JoinRole::Reducing { + anti: true, + output: JoinSide::Left, + }, + JoinType::RightAnti => JoinRole::Reducing { + anti: true, + output: JoinSide::Right, + }, + _ => return None, + }; + if has_filter && !matches!(role, JoinRole::Inner) { + return None; + } + Some(role) +} + +/// Whether a graph-wide choice and one join's are compatible. +fn agree(graph: Option, join: Option) -> bool { + match (graph, join) { + (Some(graph), Some(join)) => graph == join, + _ => true, + } +} + +fn as_column(expr: &PhysicalExprRef) -> Option { + expr.downcast_ref::().map(|col| col.index()) +} + +/// Extracts the maximal reorderable subtree at `plan`. `None` covers every bail-out: an +/// unmodelled join feature, a non-column key, missing row counts, too few or many inputs. +pub(crate) fn extract( + plan: &Arc, + stats: &mut StatsFn, +) -> Result> { + // Start at a join, or at the pruning projection usually above one, so its column list + // becomes the top join's projection instead of a `ProjectionExec` above a wider join. + let is_root = join_view(plan).is_some() + || plan + .downcast_ref::() + .is_some_and(|projection| all_alias_free_columns(projection.expr())); + if !is_root { + return Ok(None); + } + Extractor::new(stats).extract(plan) +} + +struct Extractor<'a, 's> { + graph: JoinGraph, + stats: &'s mut StatsFn<'a>, +} + +impl<'a, 's> Extractor<'a, 's> { + fn new(stats: &'s mut StatsFn<'a>) -> Self { + Self { + graph: JoinGraph { + relations: vec![], + edges: vec![], + filters: vec![], + output: vec![], + null_equality: None, + original_nodes: vec![], + reducers: 0, + kind: None, + }, + stats, + } + } + + fn extract(mut self, plan: &Arc) -> Result> { + let Some((output, _)) = self.visit(plan)? else { + return Ok(None); + }; + let mut graph = self.graph; + graph.output = output; + + if graph.relations.len() < 3 { + return Ok(None); + } + // A filter over one relation has no join to sit at; the node would be a leaf. + if graph + .filters + .iter() + .any(|filter| filter.required.count_ones() < 2) + { + return Ok(None); + } + Ok(Some(graph)) + } + + fn visit( + &mut self, + plan: &Arc, + ) -> Result, RelSet)>> { + if let Some(view) = join_view(plan) + && agree(self.graph.null_equality, view.null_equality) + && agree(self.graph.kind, view.kind) + { + self.graph.null_equality = self.graph.null_equality.or(view.null_equality); + self.graph.kind = self.graph.kind.or(view.kind); + let visited = match view.role { + JoinRole::Inner => self.visit_inner(&view)?, + JoinRole::Reducing { anti, output } => { + self.visit_reducing(&view, anti, output)? + } + }; + let Some((columns, mask)) = visited else { + return Ok(None); + }; + + // For a semi or anti join the projection selects from the output side alone. + let columns = match view.projection { + Some(projection) => projection.iter().map(|idx| columns[*idx]).collect(), + None => columns, + }; + Ok(Some((columns, mask))) + } else if let Some(projection) = plan.downcast_ref::() + && all_alias_free_columns(projection.expr()) + { + // Looking through pruning projections lets the enumerator see a whole chain, + // since `ProjectionPushdown` has not folded them into the joins yet. + let Some((child, mask)) = self.visit(projection.input())? else { + return Ok(None); + }; + let columns = projection + .expr() + .iter() + .map(|proj| { + proj.expr + .downcast_ref::() + .map(|col| child[col.index()]) + }) + .collect::>>(); + Ok(columns.map(|columns| (columns, mask))) + } else { + let Some(rel) = self.push_relation(plan, Role::Output)? else { + return Ok(None); + }; + Ok(Some(( + (0..plan.schema().fields().len()) + .map(|col| ColRef { rel, col }) + .collect(), + bit(rel), + ))) + } + } + + /// Flattens an inner join: both sides join the graph, predicates become edges + /// and filters. + fn visit_inner(&mut self, view: &JoinView) -> Result, RelSet)>> { + let Some((left, left_mask)) = self.visit(view.left)? else { + return Ok(None); + }; + let Some((right, right_mask)) = self.visit(view.right)? else { + return Ok(None); + }; + + for (left_key, right_key) in view.on { + let (Some(left_key), Some(right_key)) = + (as_column(left_key), as_column(right_key)) + else { + // A key like `cast(a) = b` would need re-deriving against a different schema. + return Ok(None); + }; + let edge = Edge { + left: left[left_key], + right: right[right_key], + }; + // Duplicates would be double counted by the cost model. + if !self + .graph + .edges + .iter() + .any(|e| (e.left, e.right) == (edge.left, edge.right)) + { + self.graph.edges.push(edge); + } + } + + if let Some(filter) = view.filter { + let columns = filter + .column_indices() + .iter() + .map(|ColumnIndex { index, side }| match side { + JoinSide::Left => Some(left[*index]), + JoinSide::Right => Some(right[*index]), + JoinSide::None => None, + }) + .collect::>>(); + let Some(columns) = columns else { + return Ok(None); + }; + let required = columns.iter().fold(0, |mask, col| mask | bit(col.rel)); + self.graph.filters.push(Filter { + filter: filter.clone(), + columns, + required, + }); + } + + let mut columns = left; + columns.extend(right); + self.graph + .original_nodes + .push((left_mask | right_mask, left_mask)); + Ok(Some((columns, left_mask | right_mask))) + } + + /// Flattens a semi or anti join: its output side joins the graph, its quantified + /// side becomes a reducer. + fn visit_reducing( + &mut self, + view: &JoinView, + anti: bool, + output: JoinSide, + ) -> Result, RelSet)>> { + let (output_plan, reducer_plan) = match output { + JoinSide::Left => (view.left, view.right), + JoinSide::Right => (view.right, view.left), + JoinSide::None => return internal_err!("semi join with no output side"), + }; + + let Some((columns, mask)) = self.visit(output_plan)? else { + return Ok(None); + }; + + // Keys resolve against the output side, so they precede the reducer relation. + let mut keys = Vec::with_capacity(view.on.len()); + let mut required = 0; + for (left_key, right_key) in view.on { + let (output_key, reducer_key) = match output { + JoinSide::Left => (left_key, right_key), + _ => (right_key, left_key), + }; + let (Some(output_key), Some(reducer_key)) = + (as_column(output_key), as_column(reducer_key)) + else { + return Ok(None); + }; + let column = columns[output_key]; + required |= bit(column.rel); + keys.push((column, reducer_key)); + } + + let role = Role::Reducer(Reducer { + anti, + keys, + required, + }); + let Some(rel) = self.push_relation(reducer_plan, role)? else { + return Ok(None); + }; + self.graph.original_nodes.push((mask | bit(rel), bit(rel))); + Ok(Some((columns, mask | bit(rel)))) + } + + fn push_relation( + &mut self, + plan: &Arc, + role: Role, + ) -> Result> { + if plan.boundedness().is_unbounded() { + // Reordering could break the pipeline properties the other subrules establish. + return Ok(None); + } + if self.graph.relations.len() >= MAX_RELATIONS { + return Ok(None); + } + let statistics = (self.stats)(plan.as_ref())?; + let Some(rows) = statistics.num_rows.get_value().copied() else { + return Ok(None); + }; + let rows = (rows as f64).max(1.0); + let ndv = statistics + .column_statistics + .iter() + .map(|col| { + max_distinct_count(&statistics.num_rows, col) + .get_value() + .map(|ndv| (*ndv as f64).clamp(1.0, rows)) + .unwrap_or(rows) + }) + .collect(); + + let rel = self.graph.relations.len(); + if matches!(role, Role::Reducer(_)) { + self.graph.reducers |= bit(rel); + } + self.graph.relations.push(Relation { + plan: Arc::clone(plan), + rows, + width: statistics + .total_byte_size + .get_value() + .map(|bytes| *bytes as f64 / rows), + ndv, + role, + }); + Ok(Some(rel)) + } +} diff --git a/datafusion/physical-optimizer/src/join_enumeration/mod.rs b/datafusion/physical-optimizer/src/join_enumeration/mod.rs new file mode 100644 index 0000000000000..c670994e8b66f --- /dev/null +++ b/datafusion/physical-optimizer/src/join_enumeration/mod.rs @@ -0,0 +1,1250 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Cost-based join order enumeration. +//! +//! A subtree of joins is flattened into relations plus the predicates between them, a +//! dynamic program searches the orders (bushy as well as left-deep) under a `C_out` cost +//! model, and the subtree is rebuilt if the winner is clearly cheaper. +//! +//! The estimates come from a [`JoinCostModel`], so a different one can be plugged in +//! with [`JoinEnumeration::with_cost_model`]. +//! +//! The graph itself, and how a subtree of joins is flattened into one, is in [`graph`]. +//! +//! Reordering is sound because a tree of inner joins equals the cross product of its +//! relations filtered by all its predicates. Semi and anti joins take part as reducers: +//! they filter their output side rather than contributing columns. + +pub mod graph; + +use std::collections::HashMap; +use std::sync::Arc; + +use crate::PhysicalOptimizerRule; +use crate::optimizer::{ConfigOnlyContext, PhysicalOptimizerContext}; + +use arrow::compute::SortOptions; +use arrow::datatypes::{FieldRef, Schema}; +use datafusion_common::config::ConfigOptions; +use datafusion_common::error::Result; +use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::{JoinSide, JoinType, Statistics, internal_err}; +use datafusion_expr_common::operator::Operator; +use datafusion_physical_expr::PhysicalExprRef; +use datafusion_physical_expr::expressions::{BinaryExpr, Column}; +use datafusion_physical_expr::projection::ProjectionExpr; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::execution_plan::replace_children_if_necessary; +use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter}; +use datafusion_physical_plan::joins::{ + CrossJoinExec, HashJoinExecBuilder, NestedLoopJoinExec, PartitionMode, + SortMergeJoinExec, +}; +use datafusion_physical_plan::operator_statistics::StatisticsRegistry; +use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; + +use graph::{ + ColRef, Edge, Filter, JoinGraph, JoinKind, MAX_RELATIONS, RelSet, StatsFn, bit, + covers, extract, iter_rels, +}; + +/// Chooses the shape of the join tree, before [`JoinSelection`] decides how each +/// join runs. +/// +/// [`JoinSelection`]: crate::join_selection::JoinSelection +#[derive(Debug)] +pub struct JoinEnumeration { + cost_model: Arc, +} + +impl JoinEnumeration { + #[expect(missing_docs)] + pub fn new() -> Self { + Self { + cost_model: Arc::new(DefaultJoinCostModelFactory {}), + } + } + + /// Searches with `cost_model` in place of [`DefaultJoinCostModel`]. + pub fn with_cost_model(mut self, cost_model: Arc) -> Self { + self.cost_model = cost_model; + self + } +} + +impl Default for JoinEnumeration { + fn default() -> Self { + Self::new() + } +} + +impl PhysicalOptimizerRule for JoinEnumeration { + fn optimize( + &self, + plan: Arc, + config: &ConfigOptions, + ) -> Result> { + self.optimize_with_context(plan, &ConfigOnlyContext::new(config)) + } + + fn optimize_with_context( + &self, + plan: Arc, + context: &dyn PhysicalOptimizerContext, + ) -> Result> { + let config = context.config_options(); + if !config.optimizer.join_enumeration { + return Ok(plan); + } + let mut default_registry = None; + let registry: Option<&StatisticsRegistry> = + if config.optimizer.use_statistics_registry { + Some(context.statistics_registry().unwrap_or_else(|| { + default_registry + .insert(StatisticsRegistry::default_with_builtin_providers()) + })) + } else { + None + }; + let mut stats = |plan: &dyn ExecutionPlan| { + if let Some(registry) = registry { + registry + .compute(plan) + .map(|s| Arc::::clone(s.base_arc())) + } else { + StatisticsContext::new().compute(plan, &StatisticsArgs::new()) + } + }; + Ok( + enumerate_join_order(&plan, config, &mut stats, self.cost_model.as_ref())? + .unwrap_or(plan), + ) + } + + fn name(&self) -> &str { + "join_enumeration" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// A hash partitioning, as the set of key classes it is partitioned on. Zero means +/// not hash partitioned, which is where every scan starts. +pub type PartSet = u32; + +/// A valid way of combining two relation sets. +#[derive(Clone, Copy, Debug)] +pub enum Combine { + /// An inner join, both sides contributing their columns. + Inner, + /// A semi or anti join applying `reducer` to the opposite set. + Reducer { + /// The reducing relation, which is the whole of its side. + reducer: usize, + }, +} + +/// One way of exchanging a join's inputs: what the exchange costs, the partitioning it +/// leaves the output in, the side that builds, and the mode the join runs in. +#[derive(Clone, Copy, Debug)] +pub struct Exchange { + /// Cost of moving the inputs, in the unit [`JoinCostModel::cardinality`] counts in. + pub cost: f64, + /// The partitioning the join's output is left in, for a join above to reuse. + pub partitioning: PartSet, + /// The side that builds: what `CollectLeft` collects, or which side is hashed. + pub build: RelSet, + /// The mode the join is costed under, which is the one the rebuild emits. + pub mode: PartitionMode, +} + +/// Cardinality and cost estimates over the subsets of a [`JoinGraph`], which is all the +/// search knows about the data. +/// +/// [`DefaultJoinCostModel`] derives them from the statistics the subtree's inputs report. +/// Implement this trait, and hand it to [`JoinEnumeration::with_cost_model`], to search +/// under other estimates: statistics kept outside the plan, a different cost function, or +/// knowledge of the data the plan has no way to carry. +pub trait JoinCostModel { + /// Estimated rows from joining every relation in `mask`. + /// + /// Must depend on the set alone and not on the order its relations were joined in: + /// that is what makes the dynamic program valid. Should stay at or above `1.0`, since + /// the search multiplies cardinalities. + fn cardinality(&self, mask: RelSet) -> f64; + + /// Whether `left` and `right` can be combined, and how: as an inner join, or with one + /// side applying as a reducer. `None` prunes the pair from the search, which is what + /// keeps a cut no predicate spans, or a reducer whose keys the other side does not + /// supply, out of the plan. + fn combine(&self, left: RelSet, right: RelSet) -> Option; + + /// Cost of joining `left` (partitioned as `left_part`) with `right` (`right_part`), + /// one entry per way of exchanging the inputs. Returning several lets the search carry + /// each partitioning forward and keep whichever pays off at the joins above; an empty + /// result prunes the pair. + /// + /// `collect_only` is the side that must build, when the combination forces one. + fn exchanges( + &self, + left: RelSet, + right: RelSet, + left_part: PartSet, + right_part: PartSet, + collect_only: Option, + ) -> Vec; + + /// The side that must build, when one of them is a reducer: the reducer is hashed so + /// the side it filters can stream. + fn reducer_side(&self, left: RelSet, right: RelSet) -> Option { + match self.combine(left, right) { + Some(Combine::Reducer { reducer }) => Some(bit(reducer)), + _ => None, + } + } + + /// Cost of a whole tree, given its internal nodes as `(node, one child)`, children + /// first. The search scores its own candidates this way, so the shape the planner + /// produced can be compared against them. + fn tree_cost(&self, nodes: &[(RelSet, RelSet)]) -> f64 { + let mut parts: HashMap = HashMap::new(); + let mut total = 0.0; + for (mask, child) in nodes { + if mask.count_ones() < 2 { + continue; + } + let other = mask ^ child; + let (child_part, other_part) = ( + parts.get(child).copied().unwrap_or(0), + parts.get(&other).copied().unwrap_or(0), + ); + let collect_only = self.reducer_side(*child, other); + let best = self + .exchanges(*child, other, child_part, other_part, collect_only) + .into_iter() + .min_by(|a, b| a.cost.total_cmp(&b.cost)); + let (exchange, part) = + best.map_or((0.0, 0), |best| (best.cost, best.partitioning)); + total += self.cardinality(*mask) + exchange; + parts.insert(*mask, part); + } + total + } +} + +/// Builds the [`JoinCostModel`] for one join subtree. +pub trait JoinCostModelFactory: std::fmt::Debug + Send + Sync { + /// Creates a cost model over `graph`. + fn create<'graph>( + &self, + graph: &'graph JoinGraph, + config: &ConfigOptions, + ) -> Result>; +} + +/// Hands out [`DefaultJoinCostModel`], which the rule searches with unless it was given +/// another factory. +#[derive(Debug, Default)] +pub struct DefaultJoinCostModelFactory {} + +impl JoinCostModelFactory for DefaultJoinCostModelFactory { + fn create<'graph>( + &self, + graph: &'graph JoinGraph, + config: &ConfigOptions, + ) -> Result> { + Ok(Box::new(DefaultJoinCostModel::new(graph, config))) + } +} + +/// The built-in [`JoinCostModel`]: a `C_out` model over the row counts, distinct counts +/// and widths the subtree's inputs report. +pub struct DefaultJoinCostModel<'a> { + graph: &'a JoinGraph, + /// Aggregated selectivity per connected relation pair, as + /// `(rel_a, rel_b, selectivity)` with `rel_a < rel_b`. + pair_selectivity: Vec<(usize, usize, f64)>, + /// Neighbours of each relation, over equi-join keys and non-equi filters alike. + /// Reducers neighbour nothing: they are applied. + adjacency: Vec, + /// The relations of each connected component of `adjacency`. + components: Vec, + /// Fraction of its filtered side each reducer keeps; `1.0` for non-reducers. + reducer_selectivity: Vec, + /// Selectivity of each non-equi filter, with the relations it needs. + filter_selectivity: Vec<(RelSet, f64)>, + /// The key class of each edge, as a bit position. Joins on the same class can + /// reuse each other's hash partitioning. + edge_class: Vec, + /// The size a build side must stay under for `JoinSelection` to broadcast it. + broadcast_bytes: f64, + /// The row count it must stay under when no byte estimate is available. + broadcast_rows: f64, +} + +impl<'a> DefaultJoinCostModel<'a> { + /// Precomputes the selectivities and connectivity of `graph`. + pub fn new(graph: &'a JoinGraph, config: &ConfigOptions) -> Self { + // Denominate each relation pair by its most selective key, as + // `estimate_inner_join_cardinality` does, so the two models agree. + let mut denominators: HashMap<(usize, usize), f64> = HashMap::new(); + for edge in &graph.edges { + let (a, b) = (edge.left.rel, edge.right.rel); + let key = if a < b { (a, b) } else { (b, a) }; + let denominator = graph.ndv(edge.left).max(graph.ndv(edge.right)).max(1.0); + denominators + .entry(key) + .and_modify(|current| *current = current.max(denominator)) + .or_insert(denominator); + } + + let mut adjacency = vec![0; graph.relations.len()]; + let mut pair_selectivity = Vec::with_capacity(denominators.len()); + for ((a, b), denominator) in denominators { + adjacency[a] |= bit(b); + adjacency[b] |= bit(a); + pair_selectivity.push((a, b, 1.0 / denominator)); + } + // `HashMap` iteration order is not deterministic, but plans must be. + pair_selectivity.sort_unstable_by_key(|(a, b, _)| (*a, *b)); + + let reducer_selectivity = (0..graph.relations.len()) + .map(|rel| match graph.reducer(rel) { + None => 1.0, + Some(reducer) => { + // Fraction of the filtered side's key values the reducer covers. + let matched = reducer + .keys + .iter() + .map(|(filtered, col)| { + let reducer_ndv = graph.ndv(ColRef { rel, col: *col }); + (reducer_ndv / graph.ndv(*filtered)).clamp(0.0, 1.0) + }) + .fold(1.0f64, f64::min); + if reducer.anti { 1.0 - matched } else { matched } + } + }) + .collect(); + + // A non-equi filter has no statistics, so it takes the optimizer's default. + let default_selectivity = + f64::from(config.optimizer.default_filter_selectivity) / 100.0; + let filter_selectivity = graph + .filters + .iter() + .map(|filter| (filter.required, default_selectivity)) + .collect(); + + // A filter links what it references, so a pair joined only by one counts as + // connected and its cuts prune like any other. + for filter in &graph.filters { + for rel in iter_rels(filter.required) { + adjacency[rel] |= filter.required & !bit(rel); + } + } + + let mut components: Vec = vec![]; + let mut seen: RelSet = 0; + for rel in 0..graph.relations.len() { + if seen & bit(rel) != 0 { + continue; + } + let mut component = bit(rel); + loop { + let grown = iter_rels(component) + .fold(component, |mask, rel| mask | adjacency[rel]); + if grown == component { + break; + } + component = grown; + } + seen |= component; + components.push(component); + } + + Self { + graph, + pair_selectivity, + adjacency, + components, + edge_class: key_classes(&graph.edges), + broadcast_bytes: config.optimizer.hash_join_single_partition_threshold as f64, + broadcast_rows: config.optimizer.hash_join_single_partition_threshold_rows + as f64, + reducer_selectivity, + filter_selectivity, + } + } + + fn connected(&self, left: RelSet, right: RelSet) -> bool { + iter_rels(left).any(|rel| self.adjacency[rel] & right != 0) + } + + /// The key classes joining `left` to `right`, which a partitioned join hashes both + /// sides on. + fn crossing_classes(&self, left: RelSet, right: RelSet) -> PartSet { + let mut classes = 0; + for (index, edge) in self.graph.edges.iter().enumerate() { + let (a, b) = (bit(edge.left.rel), bit(edge.right.rel)); + if (left & a != 0 && right & b != 0) || (left & b != 0 && right & a != 0) { + classes |= 1 << self.edge_class[index]; + } + } + classes + } + + /// Whether `JoinSelection` will broadcast this side rather than partition it, by the + /// same bytes-or-rows test it uses. A sort merge join has no mode that collects. + fn broadcasts(&self, side: RelSet) -> bool { + if self.graph.kind() == JoinKind::SortMerge { + return false; + } + let mut width = 0.0; + for rel in iter_rels(side) { + if self.graph.reducer(rel).is_some() { + continue; + } + match self.graph.relations[rel].width { + Some(bytes) => width += bytes, + None => return self.cardinality(side) < self.broadcast_rows, + } + } + self.cardinality(side) * width < self.broadcast_bytes + } +} + +impl JoinCostModel for DefaultJoinCostModel<'_> { + /// Rows are the product of the relations' row counts, cut by the selectivity of + /// every predicate the set closes over. + fn cardinality(&self, mask: RelSet) -> f64 { + let mut rows = 1.0; + for rel in iter_rels(mask) { + // A reducer contributes a selectivity, not rows. + rows *= match self.graph.reducer(rel) { + Some(_) => self.reducer_selectivity[rel], + None => self.graph.relations[rel].rows, + }; + } + for (a, b, selectivity) in &self.pair_selectivity { + if mask & bit(*a) != 0 && mask & bit(*b) != 0 { + rows *= selectivity; + } + } + for (required, selectivity) in &self.filter_selectivity { + if covers(mask, *required) { + rows *= selectivity; + } + } + rows.max(1.0) + } + + fn combine(&self, left: RelSet, right: RelSet) -> Option { + let reducers = self.graph.reducers; + // A lone reducer is applied to the other side, which must supply its keys. + for (reducer_side, filtered) in [(right, left), (left, right)] { + if reducer_side.is_power_of_two() && reducer_side & reducers != 0 { + let reducer = reducer_side.trailing_zeros() as usize; + let required = self.graph.reducer(reducer)?.required; + return covers(filtered, required) + .then_some(Combine::Reducer { reducer }); + } + } + // Otherwise both sides must contribute columns. An unconnected cut is a cross + // product, allowed only between whole components, which is the only way to join + // a disconnected graph. + if left & !reducers == 0 || right & !reducers == 0 { + return None; + } + let separates_components = self + .components + .iter() + .all(|component| component & left == 0 || component & right == 0); + (self.connected(left, right) || separates_components).then_some(Combine::Inner) + } + + /// A side already hashed on the join key is not moved again, so the cost of an + /// exchange depends on what the joins below left behind. + fn exchanges( + &self, + left: RelSet, + right: RelSet, + left_part: PartSet, + right_part: PartSet, + collect_only: Option, + ) -> Vec { + let classes = self.crossing_classes(left, right); + // Without a key every pair is examined, which the output cardinality does not + // show: a filter estimated to keep few rows still compares all of them. + let pairs = if classes == 0 { + self.cardinality(left) * self.cardinality(right) + } else { + 0.0 + }; + let mut options = vec![]; + for (build, probe_part) in [(left, right_part), (right, left_part)] { + if collect_only.is_some_and(|only| only != build) { + continue; + } + // With no key there is nothing to hash on, so a side must be collected. + if classes == 0 || self.broadcasts(build) { + options.push(Exchange { + cost: pairs + self.cardinality(build), + partitioning: probe_part, + build, + mode: PartitionMode::CollectLeft, + }); + } + } + if classes != 0 { + let mut moved = 0.0; + if left_part != classes { + moved += self.cardinality(left); + } + if right_part != classes { + moved += self.cardinality(right); + } + let build = collect_only.unwrap_or_else(|| { + if self.cardinality(left) <= self.cardinality(right) { + left + } else { + right + } + }); + options.push(Exchange { + cost: moved, + partitioning: classes, + build, + mode: PartitionMode::Partitioned, + }); + } + options + } +} + +/// Groups equi-join predicates that share a column into key classes, so two joins hashing +/// on the same class can reuse one partitioning. Returns each edge's class as a bit +/// position, with classes past the bitmask's width collapsed into the last. +fn key_classes(edges: &[Edge]) -> Vec { + let mut ids: Vec = (0..edges.len()).collect(); + let shares = |a: &Edge, b: &Edge| { + [a.left, a.right] + .iter() + .any(|col| *col == b.left || *col == b.right) + }; + loop { + let mut merged = false; + for i in 0..edges.len() { + for j in (i + 1)..edges.len() { + if ids[i] != ids[j] && shares(&edges[i], &edges[j]) { + let (keep, drop) = (ids[i].min(ids[j]), ids[i].max(ids[j])); + ids.iter_mut() + .filter(|id| **id == drop) + .for_each(|id| *id = keep); + merged = true; + } + } + } + if !merged { + break; + } + } + let mut compact = ids.clone(); + compact.sort_unstable(); + compact.dedup(); + ids.iter() + .map(|id| { + let position = compact.iter().position(|c| c == id).unwrap_or(0); + position.min(PartSet::BITS as usize - 1) as u32 + }) + .collect() +} + +/// The winning join tree: per internal node, which input goes on the left and how the +/// join exchanges its data. +struct Solution { + nodes: HashMap, + cost: f64, +} + +/// Exhaustive dynamic programming over connected relation subsets, each paired with the +/// partitioning its plan leaves behind. Carrying the partitioning lets a later join reuse +/// an earlier one's exchange instead of paying for another. +fn solve_dp(graph: &JoinGraph, model: &dyn JoinCostModel) -> Option { + let n = graph.relations.len(); + let full: RelSet = graph.all(); + + // Per subset, the cheapest plan for each partitioning it can be left in, and the + // choice that got there. A scan arrives hash partitioned on nothing. + let mut best: Vec> = vec![HashMap::new(); 1usize << n]; + let mut choice: Vec> = + vec![HashMap::new(); 1usize << n]; + for rel in 0..n { + best[bit(rel) as usize].insert(0, 0.0); + } + + for mask in 1..=full { + if mask.count_ones() < 2 { + continue; + } + let cardinality = model.cardinality(mask); + // Subsets containing the lowest set bit, so each pair of halves is seen once. + let lowest = mask & mask.wrapping_neg(); + let mut left = mask; + while left != 0 { + left = (left - 1) & mask; + if left & lowest == 0 { + continue; + } + let right = mask ^ left; + if right == 0 || model.combine(left, right).is_none() { + continue; + } + let collect_only = model.reducer_side(left, right); + for (left_part, left_cost) in best[left as usize].clone() { + for (right_part, right_cost) in best[right as usize].clone() { + let below = left_cost + right_cost + cardinality; + for exchange in + model.exchanges(left, right, left_part, right_part, collect_only) + { + let candidate = below + exchange.cost; + let entry = best[mask as usize] + .entry(exchange.partitioning) + .or_insert(f64::MAX); + if candidate < *entry { + *entry = candidate; + choice[mask as usize].insert( + exchange.partitioning, + (left, exchange.build, exchange.mode), + ); + } + } + } + } + } + } + + let (&winning_part, &cost) = best[full as usize] + .iter() + .min_by(|a, b| a.1.total_cmp(b.1))?; + + // Walk the winning tree, keeping the left input and mode of each node it uses. + let mut nodes = HashMap::new(); + let mut stack = vec![(full, winning_part)]; + while let Some((mask, part)) = stack.pop() { + if mask.count_ones() < 2 { + continue; + } + let Some(&(split, build, mode)) = choice[mask as usize].get(&part) else { + continue; + }; + let other = mask ^ split; + // The build side goes on the left: that is the side `CollectLeft` gathers. + nodes.insert(mask, (build, mode)); + for child in [split, other] { + let child_part = best[child as usize] + .iter() + .min_by(|a, b| a.1.total_cmp(b.1)) + .map(|(part, _)| *part) + .unwrap_or(0); + stack.push((child, child_part)); + } + } + + Some(Solution { nodes, cost }) +} + +fn position(columns: &[ColRef], col: ColRef) -> Option { + columns.iter().position(|candidate| *candidate == col) +} + +fn push_unique(columns: &mut Vec, col: ColRef) { + if position(columns, col).is_none() { + columns.push(col); + } +} + +fn extend_required(columns: &mut Vec, wanted: &[ColRef], side: RelSet) { + for col in wanted.iter().filter(|col| side & bit(col.rel) != 0) { + push_unique(columns, *col); + } +} + +struct Rebuilder<'a> { + graph: &'a JoinGraph, + model: &'a dyn JoinCostModel, + solution: &'a Solution, + /// Each relation's plan, already rewritten if it held a join subtree. + relations: &'a [Arc], +} + +impl Rebuilder<'_> { + fn node( + &self, + mask: RelSet, + required: &[ColRef], + ) -> Result<(Arc, Vec)> { + if mask.is_power_of_two() { + // A single relation, opaque here; narrowing is `ProjectionPushdown`'s job. + let rel = mask.trailing_zeros() as usize; + let plan = Arc::clone(&self.relations[rel]); + let columns = (0..plan.schema().fields().len()) + .map(|col| ColRef { rel, col }) + .collect(); + return Ok((plan, columns)); + } + + let Some(&(left, mode)) = self.solution.nodes.get(&mask) else { + return internal_err!("join enumeration produced no split for {mask:b}"); + }; + let right = mask ^ left; + match self.model.combine(left, right) { + None => internal_err!("join enumeration produced an invalid join"), + Some(Combine::Reducer { reducer }) => { + self.reducing(mask, required, reducer, mode) + } + // The search chose the build side and the mode along with the order, so both + // are emitted as decided rather than left to `JoinSelection`. + Some(Combine::Inner) => self.inner(required, left, right, mode), + } + } + + /// Builds one join of the kind the subtree used. + fn build_join( + &self, + left: Built, + right: Built, + spec: JoinSpec, + required: &[ColRef], + ) -> Result<(Arc, Vec)> { + let JoinSpec { + on, + join_type, + filter, + mode, + } = spec; + // A semi or anti join emits one side, and a reducer's own columns are never + // emitted, so its `Built` carries none. + let natural = match join_type { + JoinType::RightSemi | JoinType::RightAnti => right.columns.clone(), + JoinType::LeftSemi | JoinType::LeftAnti => left.columns.clone(), + _ => { + let mut natural = left.columns.clone(); + natural.extend(right.columns.clone()); + natural + } + }; + let keys = on.len(); + if keys == 0 { + // No keys: with a filter this is a nested loop join, without one a cross + // join. + let projection = projection_for(required, &natural)?; + return match filter { + Some(filter) => { + let join = NestedLoopJoinExec::try_new( + left.plan, + right.plan, + Some(filter), + &join_type, + projection.clone(), + )?; + // `projection_for` is `None` only for the identity. + let emitted = match projection { + Some(_) => required.to_vec(), + None => natural, + }; + Ok((Arc::new(join), emitted)) + } + // A cross join has no projection of its own. + None => { + Ok((Arc::new(CrossJoinExec::new(left.plan, right.plan)), natural)) + } + }; + } + match self.graph.kind() { + JoinKind::Hash => { + let join = HashJoinExecBuilder::new(left.plan, right.plan, on, join_type) + .with_filter(filter) + .with_null_equality(self.graph.null_equality()) + .with_partition_mode(mode) + .with_projection(projection_for(required, &natural)?) + .build()?; + Ok((Arc::new(join), required.to_vec())) + } + JoinKind::SortMerge => { + let join: Arc = Arc::new(SortMergeJoinExec::try_new( + left.plan, + right.plan, + on, + filter, + join_type, + vec![SortOptions::default(); keys], + self.graph.null_equality(), + )?); + // A sort merge join emits every column, so drop the ones nothing above + // needs here rather than once above the whole subtree: what this join + // emits is what the next one sorts. `ProjectionPushdown` cannot do it + // afterwards, since it only pushes through a join whose columns stay + // left-then-right, and reordering interleaves them. Reordering alone is + // left to the parent, which addresses columns by position anyway. + if required.len() == natural.len() { + return Ok((join, natural)); + } + let Some(projection) = projection_for(required, &natural)? else { + return Ok((join, natural)); + }; + Ok((narrow(join, &projection)?, required.to_vec())) + } + } + } + + fn inner( + &self, + required: &[ColRef], + left_mask: RelSet, + right_mask: RelSet, + mode: PartitionMode, + ) -> Result<(Arc, Vec)> { + // Each edge crosses exactly one cut, at the lowest node holding both endpoints. + let mut keys: Vec<(ColRef, ColRef)> = vec![]; + for edge in &self.graph.edges { + let (left, right) = (edge.left, edge.right); + if left_mask & bit(left.rel) != 0 && right_mask & bit(right.rel) != 0 { + keys.push((left, right)); + } else if left_mask & bit(right.rel) != 0 && right_mask & bit(left.rel) != 0 { + keys.push((right, left)); + } + } + // Filters land at their lowest common ancestor: covered here, by neither input. + let mask = left_mask | right_mask; + let filters: Vec<&Filter> = self + .graph + .filters + .iter() + .filter(|filter| { + covers(mask, filter.required) + && !covers(left_mask, filter.required) + && !covers(right_mask, filter.required) + }) + .collect(); + + // Each side emits this join's keys, its filters' columns, and what is needed + // above. + let child_required = |side: RelSet, take_left: bool| { + let mut columns: Vec = vec![]; + for (left, right) in &keys { + push_unique(&mut columns, if take_left { *left } else { *right }); + } + for filter in &filters { + extend_required(&mut columns, &filter.columns, side); + } + extend_required(&mut columns, required, side); + columns + }; + + let (left_plan, left_columns) = + self.node(left_mask, &child_required(left_mask, true))?; + let (right_plan, right_columns) = + self.node(right_mask, &child_required(right_mask, false))?; + + let on = keys + .iter() + .map(|(left, right)| { + Ok(( + key_expr(&left_columns, &left_plan, *left)?, + key_expr(&right_columns, &right_plan, *right)?, + )) + }) + .collect::>>()?; + let filter = rebuild_filters(&filters, &left_columns, &right_columns)?; + + self.build_join( + Built { + plan: left_plan, + columns: left_columns, + }, + Built { + plan: right_plan, + columns: right_columns, + }, + JoinSpec { + on, + join_type: JoinType::Inner, + filter, + mode, + }, + required, + ) + } + + fn reducing( + &self, + mask: RelSet, + required: &[ColRef], + reducer: usize, + mode: PartitionMode, + ) -> Result<(Arc, Vec)> { + let Some(info) = self.graph.reducer(reducer) else { + return internal_err!("relation {reducer} is not a reducer"); + }; + let filtered_mask = mask & !bit(reducer); + + // Filters never land here: they reference output relations only, so their + // lowest covering node is always inside the filtered side. + let mut filtered_required: Vec = vec![]; + for (column, _) in &info.keys { + push_unique(&mut filtered_required, *column); + } + extend_required(&mut filtered_required, required, filtered_mask); + + let (filtered_plan, filtered_columns) = + self.node(filtered_mask, &filtered_required)?; + let reducer_plan = Arc::clone(&self.relations[reducer]); + + // Reducer on the build side so the filtered side streams: `RightSemi` and + // `RightAnti` emit rows of their right input. + let reducer_schema = reducer_plan.schema(); + let on = info + .keys + .iter() + .map(|(column, reducer_col)| { + let reducer_key = Arc::new(Column::new( + reducer_schema.field(*reducer_col).name(), + *reducer_col, + )) as _; + Ok(( + reducer_key, + key_expr(&filtered_columns, &filtered_plan, *column)?, + )) + }) + .collect::>>()?; + + let join_type = if info.anti { + JoinType::RightAnti + } else { + JoinType::RightSemi + }; + + self.build_join( + Built { + plan: reducer_plan, + columns: vec![], + }, + Built { + plan: filtered_plan, + columns: filtered_columns, + }, + JoinSpec { + on, + join_type, + filter: None, + mode, + }, + required, + ) + } +} + +fn key_expr( + columns: &[ColRef], + plan: &Arc, + col: ColRef, +) -> Result { + let Some(index) = position(columns, col) else { + return internal_err!("join enumeration lost column {col:?}"); + }; + Ok(Arc::new(Column::new(plan.schema().field(index).name(), index)) as _) +} + +/// Wraps `plan` in a projection keeping only `indices`, in that order. +fn narrow( + plan: Arc, + indices: &[usize], +) -> Result> { + let schema = plan.schema(); + let exprs: Vec = indices + .iter() + .map(|&index| { + let name = schema.field(index).name(); + ProjectionExpr { + expr: Arc::new(Column::new(name, index)), + alias: name.clone(), + } + }) + .collect(); + Ok(Arc::new(ProjectionExec::try_new(exprs, plan)?)) +} + +fn projection_for(required: &[ColRef], emitted: &[ColRef]) -> Result>> { + let mut projection = Vec::with_capacity(required.len()); + for col in required { + let Some(index) = position(emitted, *col) else { + return internal_err!("join enumeration lost column {col:?}"); + }; + projection.push(index); + } + let identity = projection.len() == emitted.len() + && projection.iter().enumerate().all(|(idx, col)| idx == *col); + Ok((!identity).then_some(projection)) +} + +/// One built input and the columns it emits. +struct Built { + plan: Arc, + columns: Vec, +} + +/// One join as the search decided it, including the mode it was costed with. +struct JoinSpec { + on: Vec<(PhysicalExprRef, PhysicalExprRef)>, + join_type: JoinType, + filter: Option, + mode: PartitionMode, +} + +/// Rebuilds the non-equi filters applied at one join as one conjunction. A [`JoinFilter`] +/// addresses its batch by index, so merging shifts all but the first filter's indices. +fn rebuild_filters( + filters: &[&Filter], + left_columns: &[ColRef], + right_columns: &[ColRef], +) -> Result> { + let Some((first, rest)) = filters.split_first() else { + return Ok(None); + }; + + let mut expression = Arc::clone(first.filter.expression()); + let mut fields: Vec = + first.filter.schema().fields().iter().cloned().collect(); + let mut columns: Vec = first.columns.clone(); + for filter in rest { + let shifted = + shift_columns(Arc::clone(filter.filter.expression()), fields.len())?; + expression = Arc::new(BinaryExpr::new(expression, Operator::And, shifted)) as _; + fields.extend(filter.filter.schema().fields().iter().cloned()); + columns.extend_from_slice(&filter.columns); + } + + let column_indices = columns + .iter() + .map(|col| { + if let Some(index) = position(left_columns, *col) { + Ok(ColumnIndex { + index, + side: JoinSide::Left, + }) + } else if let Some(index) = position(right_columns, *col) { + Ok(ColumnIndex { + index, + side: JoinSide::Right, + }) + } else { + internal_err!("join enumeration lost filter column {col:?}") + } + }) + .collect::>>()?; + + // Intermediate columns go left side first: a sort merge join rebuilds the batch as + // all left then all right columns, so interleaved sides would read wrong columns. + let mut order: Vec = (0..column_indices.len()).collect(); + order.sort_by_key(|i| column_indices[*i].side == JoinSide::Right); + let mut moved = vec![0; order.len()]; + for (to, from) in order.iter().enumerate() { + moved[*from] = to; + } + + Ok(Some(JoinFilter::new( + remap_columns(expression, &moved)?, + order.iter().map(|i| column_indices[*i].clone()).collect(), + Arc::new(Schema::new( + order + .iter() + .map(|i| Arc::clone(&fields[*i])) + .collect::>(), + )), + ))) +} + +fn shift_columns(expression: PhysicalExprRef, offset: usize) -> Result { + if offset == 0 { + return Ok(expression); + } + rewrite_columns(expression, &|index| index + offset) +} + +fn remap_columns( + expression: PhysicalExprRef, + moved: &[usize], +) -> Result { + if moved.iter().enumerate().all(|(from, to)| from == *to) { + return Ok(expression); + } + rewrite_columns(expression, &|index| moved[index]) +} + +fn rewrite_columns( + expression: PhysicalExprRef, + index: &dyn Fn(usize) -> usize, +) -> Result { + expression + .transform(|expr| { + Ok(match expr.downcast_ref::() { + Some(column) => Transformed::yes(Arc::new(Column::new( + column.name(), + index(column.index()), + )) as _), + None => Transformed::no(expr), + }) + }) + .data() +} + +type RewrittenRelation = (Arc, Arc); + +/// Substitutes rewritten relations into a subtree, leaving its shape untouched. The +/// relations are the exact `Arc`s taken from it, so pointer identity finds them. +fn replace_relations( + plan: &Arc, + rewritten: &[RewrittenRelation], +) -> Result>> { + if rewritten.is_empty() { + return Ok(None); + } + if let Some((_, new)) = rewritten + .iter() + .find(|(original, _)| Arc::ptr_eq(original, plan)) + { + return Ok(Some(Arc::clone(new))); + } + let mut changed = false; + let children = plan + .children() + .into_iter() + .map(|child| match replace_relations(child, rewritten)? { + Some(new_child) => { + changed = true; + Ok(new_child) + } + None => Ok(Arc::clone(child)), + }) + .collect::>>()?; + if changed { + Ok(Some(replace_children_if_necessary( + Arc::clone(plan), + children, + )?)) + } else { + Ok(None) + } +} + +pub(crate) fn enumerate_join_order( + plan: &Arc, + config: &ConfigOptions, + stats: &mut StatsFn, + cost_model: &dyn JoinCostModelFactory, +) -> Result>> { + if let Some(graph) = extract(plan, stats)? { + if let Some(reordered) = reorder(&graph, config, stats, cost_model)? { + return Ok(Some(reordered)); + } + // Rejected, so descend into the relations rather than the children: re-extracting + // here would re-search costed subsets and readmit what the margin turned down. + let mut rewritten: Vec = vec![]; + for relation in &graph.relations { + if let Some(new) = + enumerate_join_order(&relation.plan, config, stats, cost_model)? + { + rewritten.push((Arc::clone(&relation.plan), new)); + } + } + return replace_relations(plan, &rewritten); + } + + let mut changed = false; + let children = plan + .children() + .into_iter() + .map( + |child| match enumerate_join_order(child, config, stats, cost_model)? { + Some(new_child) => { + changed = true; + Ok(new_child) + } + None => Ok(Arc::clone(child)), + }, + ) + .collect::>>()?; + if changed { + Ok(Some(replace_children_if_necessary( + Arc::clone(plan), + children, + )?)) + } else { + Ok(None) + } +} + +fn reorder( + graph: &JoinGraph, + config: &ConfigOptions, + stats: &mut StatsFn, + cost_model: &dyn JoinCostModelFactory, +) -> Result>> { + let limit = config.optimizer.join_enumeration_limit.min(MAX_RELATIONS); + if graph.relations.len() > limit { + return Ok(None); + } + let model = cost_model.create(graph, config)?; + let Some(solution) = solve_dp(graph, model.as_ref()) else { + return Ok(None); + }; + + // Keep the planner's order unless the winner is clearly cheaper: where estimates + // cannot tell orders apart, the winner is arbitrary. + let margin = f64::from(config.optimizer.join_enumeration_min_improvement) / 100.0; + if solution.cost >= model.tree_cost(&graph.original_nodes) * (1.0 - margin) { + return Ok(None); + } + + let mut relations = Vec::with_capacity(graph.relations.len()); + for relation in &graph.relations { + relations.push( + enumerate_join_order(&relation.plan, config, stats, cost_model)? + .unwrap_or_else(|| Arc::clone(&relation.plan)), + ); + } + + let rebuilder = Rebuilder { + graph, + model: model.as_ref(), + solution: &solution, + relations: &relations, + }; + let (plan, columns) = rebuilder.node(graph.all(), &graph.output)?; + if columns == graph.output { + return Ok(Some(plan)); + } + // Only a sort merge subtree reaches here, when its root emits more than the output. + let Some(projection) = projection_for(&graph.output, &columns)? else { + return Ok(Some(plan)); + }; + Ok(Some(narrow(plan, &projection)?)) +} diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs index 4ce6329f327fb..c935b1a2763b3 100644 --- a/datafusion/physical-optimizer/src/join_selection.rs +++ b/datafusion/physical-optimizer/src/join_selection.rs @@ -21,7 +21,9 @@ //! tries to transform a non-runnable query (with the given infinite sources) //! into a runnable query by replacing pipeline-breaking join operations with //! pipeline-friendly ones. To achieve the second goal, it selects the proper -//! `PartitionMode` and the build side using the available statistics for hash joins. +//! `PartitionMode` and the build side using the available statistics for hash +//! joins. The shape of the join tree is chosen before this, by +//! [`JoinEnumeration`](crate::join_enumeration::JoinEnumeration). use crate::PhysicalOptimizerRule; use crate::optimizer::{ConfigOnlyContext, PhysicalOptimizerContext}; @@ -33,7 +35,12 @@ use datafusion_common::{JoinSide, JoinType, internal_err}; use datafusion_expr_common::sort_properties::SortProperties; use datafusion_physical_expr::LexOrdering; use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use datafusion_physical_plan::Distribution; use datafusion_physical_plan::execution_plan::EmissionType; +use datafusion_physical_plan::execution_plan::{ + ChildrenPropertiesMode, ReplaceChildrenOptions, +}; use datafusion_physical_plan::joins::utils::ColumnIndex; use datafusion_physical_plan::joins::{ CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, @@ -170,10 +177,13 @@ impl PhysicalOptimizerRule for JoinSelection { let new_plan = plan .transform_up(|p| apply_subrules(p, &subrules, config)) .data()?; - new_plan + let new_plan = new_plan .transform_up(|plan| { statistical_join_selection_subrule(plan, config, registry) }) + .data()?; + new_plan + .transform_down(|plan| keep_partitioning_needed_above(plan, registry)) .data() } @@ -186,6 +196,150 @@ impl PhysicalOptimizerRule for JoinSelection { } } +/// How far above a join a hash requirement still counts as that join's, enough to reach +/// through the partial aggregate a group-by is planned as. +const PARTITIONING_LOOKAHEAD: usize = 3; + +/// Restores a partitioned join whose partitioning an operator above needs. Collecting the +/// build side saves its exchanges but discards the partitioning, which is then rebuilt +/// above, at more than collecting saved. +fn keep_partitioning_needed_above( + plan: Arc, + registry: Option<&StatisticsRegistry>, +) -> Result>> { + let mut children: Vec<_> = plan.children().into_iter().map(Arc::clone).collect(); + let mut changed = false; + let requirements = plan.input_distribution_requirements(); + for (idx, required) in requirements.per_child_distributions().enumerate() { + #[expect( + deprecated, + reason = "HashPartitioned is still planned during the KeyPartitioned migration" + )] + let required_exprs = match required { + Distribution::KeyPartitioned(exprs) + | Distribution::HashPartitioned(exprs) => exprs, + _ => continue, + }; + if let Some(rewritten) = + repartition_collected_join(&children[idx], required_exprs, registry)? + { + children[idx] = rewritten; + changed = true; + } + } + if changed { + Ok(Transformed::yes(plan.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?)) + } else { + Ok(Transformed::no(plan)) + } +} + +/// Rebuilds the collected join under `plan` as partitioned when its keys are what +/// `required_exprs` asks for, looking through the operators that pass a partitioning up. +fn repartition_collected_join( + plan: &Arc, + required_exprs: &[Arc], + registry: Option<&StatisticsRegistry>, +) -> Result>> { + let Some(required) = column_names(required_exprs) else { + return Ok(None); + }; + let mut node = Arc::clone(plan); + for depth in 0..PARTITIONING_LOOKAHEAD { + if let Some(join) = node.downcast_ref::() { + if *join.partition_mode() != PartitionMode::CollectLeft + || join.on().is_empty() + { + return Ok(None); + } + // The join is not hash partitioned yet, its inputs still being single + // partitions, so compare against the keys it would be partitioned on. + let keys: Vec<_> = + join.on().iter().map(|(left, _)| Arc::clone(left)).collect(); + if column_names(&keys).is_none_or(|keys| keys != required) { + return Ok(None); + } + if !worth_partitioning(join, registry)? { + return Ok(None); + } + let partitioned: Arc = Arc::new( + join.builder() + .with_partition_mode(PartitionMode::Partitioned) + .build()?, + ); + return Ok(Some(rebuild_above(plan, &node, partitioned, depth)?)); + } + // Only a single-input operator with no requirement of its own passes one up. + let children = node.children(); + let [child] = children.as_slice() else { + return Ok(None); + }; + if !matches!( + node.input_distribution_requirements().child_distribution(0), + Some(Distribution::UnspecifiedDistribution) + ) { + return Ok(None); + } + let child = Arc::clone(child); + drop(children); + node = child; + } + Ok(None) +} + +/// Whether partitioning the join moves no more rows than collecting it does. Both move +/// the build side; partitioning then moves the probe side, collecting the output above. +fn worth_partitioning( + join: &HashJoinExec, + registry: Option<&StatisticsRegistry>, +) -> Result { + let rows = |plan: &dyn ExecutionPlan| -> Result> { + Ok(get_stats(plan, registry)?.num_rows.get_value().copied()) + }; + let (Some(probe), Some(output)) = (rows(join.right().as_ref())?, rows(join)?) else { + return Ok(false); + }; + Ok(output >= probe) +} + +/// The names the expressions partition on, or `None` if any is not a column. +fn column_names(exprs: &[Arc]) -> Option> { + exprs + .iter() + .map(|expr| expr.downcast_ref::().map(Column::name)) + .collect() +} + +/// Puts `replacement` back under the `depth` single-input operators above it. +fn rebuild_above( + top: &Arc, + old: &Arc, + replacement: Arc, + depth: usize, +) -> Result> { + if depth == 0 { + return Ok(replacement); + } + let mut rebuilt = replacement; + let mut chain = vec![]; + let mut node = Arc::clone(top); + while !Arc::ptr_eq(&node, old) { + chain.push(Arc::clone(&node)); + let child = Arc::clone(node.children()[0]); + node = child; + } + for parent in chain.into_iter().rev() { + rebuilt = parent.replace_children( + vec![rebuilt], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + } + Ok(rebuilt) +} + /// Determines whether it is possible to swap inputs of a hash join - for null-aware joins, we can only swap `LeftAnti` with no filters fn can_swap_hash_join(hash_join: &HashJoinExec) -> bool { hash_join.join_type().supports_swap() diff --git a/datafusion/physical-optimizer/src/lib.rs b/datafusion/physical-optimizer/src/lib.rs index b9eb248f6e843..4d45f573685bd 100644 --- a/datafusion/physical-optimizer/src/lib.rs +++ b/datafusion/physical-optimizer/src/lib.rs @@ -34,6 +34,7 @@ pub mod ensure_requirements; // modules keep their public paths. pub use ensure_requirements::{enforce_distribution, enforce_sorting}; pub mod filter_pushdown; +pub mod join_enumeration; pub mod join_selection; pub mod limit_pushdown; pub mod limit_pushdown_past_window; diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs index aed25546cd09b..21be83dc1e246 100644 --- a/datafusion/physical-optimizer/src/optimizer.rs +++ b/datafusion/physical-optimizer/src/optimizer.rs @@ -25,6 +25,7 @@ use crate::combine_partial_final_agg::CombinePartialFinalAggregate; use crate::ensure_coop::EnsureCooperative; use crate::ensure_requirements::EnsureRequirements; use crate::filter_pushdown::FilterPushdown; +use crate::join_enumeration::JoinEnumeration; use crate::join_selection::JoinSelection; use crate::limit_pushdown::LimitPushdown; use crate::limited_distinct_aggregation::LimitedDistinctAggregation; @@ -93,6 +94,9 @@ impl PhysicalOptimizer { // this information is not lost across different rules during optimization. Arc::new(OutputRequirements::new_add_mode()), Arc::new(AggregateStatistics::new()), + // Must run before JoinSelection, which decides each join's build side and + // partition mode. + Arc::new(JoinEnumeration::new()), // Statistics-based join selection will change the Auto mode to a real join implementation, // like collect left, or hash join, or future sort merge join, which will influence the // EnsureRequirements rule as it decides whether to add additional repartitioning and diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index 414e5a6d8586a..441237d11a91b 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -35,6 +35,7 @@ use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, PushedDown, }; +use crate::joins::utils::max_distinct_count; use crate::limit::LocalLimitExec; use crate::metrics::{MetricBuilder, MetricType}; use crate::projection::{ @@ -63,13 +64,13 @@ use datafusion_execution::TaskContext; use datafusion_expr::Operator; use datafusion_physical_expr::equivalence::ProjectionMapping; use datafusion_physical_expr::expressions::{ - BinaryExpr, Column, IsNotNullExpr, Literal, lit, + BinaryExpr, Column, InListExpr, IsNotNullExpr, Literal, lit, }; use datafusion_physical_expr::intervals::utils::check_support; use datafusion_physical_expr::utils::{collect_columns, reassign_expr_columns}; use datafusion_physical_expr::{ AcrossPartitions, AnalysisContext, ConstExpr, ExprBoundaries, PhysicalExpr, analyze, - conjunction, split_conjunction, + conjunction, conjunction_opt, split_conjunction, }; use datafusion_physical_expr_common::physical_expr::fmt_sql; @@ -359,16 +360,49 @@ impl FilterExec { } else { let null_rejecting_columns = collect_null_rejecting_columns(predicate); - if check_support(predicate, schema) { + // Estimate one top-level conjunct at a time: interval analysis rejects a whole + // predicate if any part is out of reach, and an expanded `IN` list is. + let (supported, rest): (Vec<_>, Vec<_>) = split_conjunction(predicate) + .into_iter() + .partition(|conjunct| check_support(conjunct, schema)); + let split_anything = !rest.is_empty(); + + // An unrecognized conjunct contributes the default once, as the whole predicate + // used to. + let mut unanalyzed_selectivity = 1.0; + let mut has_unknown_conjunct = false; + for conjunct in rest { + match in_list_selectivity( + conjunct, + &input_stats.column_statistics, + &input_num_rows, + ) { + Some(selectivity) => unanalyzed_selectivity *= selectivity, + None => has_unknown_conjunct = true, + } + } + if has_unknown_conjunct { + unanalyzed_selectivity *= default_selectivity as f64 / 100.0; + } + + // Rebuilding re-associates the conjunction and interval propagation is + // shape sensitive, so pass the predicate through untouched if nothing split. + let analyzable = if split_anything { + conjunction_opt(supported.into_iter().cloned()) + } else { + Some(Arc::clone(predicate)) + }; + if let Some(analyzable) = analyzable { let input_analysis_ctx = AnalysisContext::try_from_statistics( schema, &input_stats.column_statistics, )?; - let analysis_ctx = analyze(predicate, input_analysis_ctx, schema)?; - let selectivity = analysis_ctx.selectivity.unwrap_or(1.0); + let analysis_ctx = analyze(&analyzable, input_analysis_ctx, schema)?; + let selectivity = + analysis_ctx.selectivity.unwrap_or(1.0) * unanalyzed_selectivity; let filtered_num_rows = input_num_rows.with_estimated_selectivity(selectivity); - let cs = collect_new_statistics( + let mut cs = collect_new_statistics( schema, &input_stats.column_statistics, analysis_ctx.boundaries, @@ -376,12 +410,19 @@ impl FilterExec { &null_rejecting_columns, filtered_num_rows, ); + // An equality that was split off pins its column to one distinct value. + for idx in &eq_columns { + if let Some(col_stat) = cs.get_mut(*idx) + && col_stat.distinct_count != Precision::Exact(1) + { + col_stat.distinct_count = + distinct_count_for_singleton_domain(filtered_num_rows); + } + } (selectivity, filtered_num_rows, cs) } else { - // Without interval boundaries, use the default selectivity and - // apply the row-count constraints that still follow from the - // filter predicate. - let selectivity = default_selectivity as f64 / 100.0; + // No boundaries to derive, so keep the input's value statistics. + let selectivity = unanalyzed_selectivity; let filtered_num_rows = input_num_rows.with_estimated_selectivity(selectivity); let mut cs = input_stats.to_inexact().column_statistics; @@ -959,6 +1000,91 @@ impl EmbeddedProjection for FilterExec { /// /// Only AND conjunctions are traversed; OR is intentionally skipped /// since `a = 1 OR a = 2` does not pin NDV to 1. +/// Estimates `col IN (a, b, c)`, and the `OR` chain a short list expands into, as +/// `distinct literals / distinct values`. Interval arithmetic cannot narrow a column +/// from a disjunction, so such a conjunct would otherwise only take the default. +/// +/// `None` when the conjunct is not a list of literals over one column. +fn in_list_selectivity( + conjunct: &Arc, + column_statistics: &[ColumnStatistics], + num_rows: &Precision, +) -> Option { + let mut column = None; + let mut values: Vec = vec![]; + if let Some(in_list) = conjunct.downcast_ref::() { + if in_list.negated() { + // `NOT IN` selects most of the column; left to the default. + return None; + } + column = Some(in_list.expr().downcast_ref::()?); + for value in in_list.list() { + push_literal(&mut values, value)?; + } + } else { + // Disjunctions only, so a plain `col = literal` keeps the estimate it has. + let binary = conjunct.downcast_ref::()?; + if *binary.op() != Operator::Or { + return None; + } + collect_or_equalities(conjunct, &mut column, &mut values)?; + } + + let column = column?; + if values.len() < 2 { + return None; + } + let stats = column_statistics.get(column.index())?; + let distinct = *max_distinct_count(num_rows, stats).get_value()?; + if distinct == 0 { + return None; + } + Some((values.len() as f64 / distinct as f64).min(1.0)) +} + +fn push_literal( + values: &mut Vec, + expr: &Arc, +) -> Option<()> { + let value = expr.downcast_ref::()?.value(); + if !values.contains(value) { + values.push(value.clone()); + } + Some(()) +} + +/// Collects the literals of an `OR` chain of equalities over one column. `None` +/// for anything else, so `a = 1 OR b = 2` is not mistaken for a list. +fn collect_or_equalities<'a>( + expr: &'a Arc, + column: &mut Option<&'a Column>, + values: &mut Vec, +) -> Option<()> { + let binary = expr.downcast_ref::()?; + match binary.op() { + Operator::Or => { + collect_or_equalities(binary.left(), column, values)?; + collect_or_equalities(binary.right(), column, values) + } + Operator::Eq => { + let (found, literal) = match ( + binary.left().downcast_ref::(), + binary.right().downcast_ref::(), + ) { + (Some(col), None) => (col, binary.right()), + (None, Some(col)) => (col, binary.left()), + _ => return None, + }; + if column.is_some_and(|current| current != found) { + return None; + } + *column = Some(found); + push_literal(values, literal) + } + _ => None, + } +} + fn collect_equality_columns(predicate: &Arc) -> (HashSet, bool) { let mut eq_values: HashMap = HashMap::new(); let mut infeasible = false; @@ -2739,9 +2865,9 @@ mod tests { Arc::new(Literal::new(ScalarValue::Utf8(Some("b".to_string())))), )), )), - // Input NDV is 50, but the 20% default selectivity on 100 rows - // estimates 20 output rows, so NDV is capped at 20. - vec![Precision::Inexact(20)], + // The two listed values are 2 of the column's 50, so 4 of the 100 rows + // are expected and NDV is capped at 4, not collapsed to 1. + vec![Precision::Inexact(4)], ), ( "AND with mixed types (Utf8 + Int32)", diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 62771af4c413d..f01350aaa0b35 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -466,12 +466,39 @@ pub(crate) fn estimate_join_statistics( join_type: &JoinType, schema: &Schema, ) -> Result { + // Width of one output row, from the sides this join emits. Without it a join reports + // no size and `hash_join_single_partition_threshold` falls back to counting rows. + let width = |stats: &Statistics| match ( + stats.total_byte_size.get_value(), + stats.num_rows.get_value(), + ) { + (Some(bytes), Some(rows)) if *rows > 0 => Some(*bytes as f64 / *rows as f64), + _ => None, + }; + let output_width = match join_type { + JoinType::LeftSemi | JoinType::LeftAnti => width(&left_stats), + JoinType::RightSemi | JoinType::RightAnti => width(&right_stats), + JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => { + width(&left_stats) + .zip(width(&right_stats)) + .map(|(l, r)| l + r) + } + // A mark join emits one side plus a boolean. + JoinType::LeftMark => width(&left_stats).map(|w| w + 1.0), + JoinType::RightMark => width(&right_stats).map(|w| w + 1.0), + }; + let join_stats = estimate_join_cardinality(join_type, left_stats, right_stats, on, null_equality); let (num_rows, total_byte_size, column_statistics) = match join_stats { Some(stats) => ( Precision::Inexact(stats.num_rows), - stats.total_byte_size, + match output_width { + Some(width) => { + Precision::Inexact((stats.num_rows as f64 * width) as usize) + } + None => stats.total_byte_size, + }, stats.column_statistics, ), None => ( @@ -1023,7 +1050,7 @@ fn estimate_semi_join_cardinality( /// directly. Otherwise, if the column is numeric and has min/max values, it /// estimates the maximum distinct count from those. Otherwise, the num_rows /// is used. -fn max_distinct_count( +pub fn max_distinct_count( num_rows: &Precision, stats: &ColumnStatistics, ) -> Precision { diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index b6837002086ad..ad989647b6213 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -234,6 +234,7 @@ physical_plan after OutputRequirements 01)OutputRequirementExec: order_by=[], dist_by=Unspecified 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true physical_plan after aggregate_statistics SAME TEXT AS ABOVE +physical_plan after join_enumeration SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE physical_plan after FilterPushdown SAME TEXT AS ABOVE @@ -314,6 +315,7 @@ physical_plan after OutputRequirements 02)--GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]: ScanBytes=Exact(32)),(Col[1]: ScanBytes=Inexact(24)),(Col[2]: ScanBytes=Exact(32)),(Col[3]: ScanBytes=Exact(32)),(Col[4]: ScanBytes=Exact(32)),(Col[5]: ScanBytes=Exact(64)),(Col[6]: ScanBytes=Exact(32)),(Col[7]: ScanBytes=Exact(64)),(Col[8]: ScanBytes=Inexact(88)),(Col[9]: ScanBytes=Inexact(49)),(Col[10]: ScanBytes=Exact(64))]] 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]: ScanBytes=Exact(32)),(Col[1]: ScanBytes=Inexact(24)),(Col[2]: ScanBytes=Exact(32)),(Col[3]: ScanBytes=Exact(32)),(Col[4]: ScanBytes=Exact(32)),(Col[5]: ScanBytes=Exact(64)),(Col[6]: ScanBytes=Exact(32)),(Col[7]: ScanBytes=Exact(64)),(Col[8]: ScanBytes=Inexact(88)),(Col[9]: ScanBytes=Inexact(49)),(Col[10]: ScanBytes=Exact(64))]] physical_plan after aggregate_statistics SAME TEXT AS ABOVE +physical_plan after join_enumeration SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE physical_plan after FilterPushdown SAME TEXT AS ABOVE @@ -360,6 +362,7 @@ physical_plan after OutputRequirements 02)--GlobalLimitExec: skip=0, fetch=10 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet physical_plan after aggregate_statistics SAME TEXT AS ABOVE +physical_plan after join_enumeration SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE physical_plan after FilterPushdown SAME TEXT AS ABOVE @@ -613,6 +616,7 @@ physical_plan after OutputRequirements 01)OutputRequirementExec: order_by=[], dist_by=Unspecified 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true physical_plan after aggregate_statistics SAME TEXT AS ABOVE +physical_plan after join_enumeration SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE physical_plan after FilterPushdown SAME TEXT AS ABOVE diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 573fb04b3451b..68879996934b5 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -324,8 +324,11 @@ datafusion.optimizer.expand_views_at_output false datafusion.optimizer.filter_null_join_keys false datafusion.optimizer.hash_join_inlist_pushdown_max_distinct_values 150 datafusion.optimizer.hash_join_inlist_pushdown_max_size 131072 -datafusion.optimizer.hash_join_single_partition_threshold 1048576 +datafusion.optimizer.hash_join_single_partition_threshold 4194304 datafusion.optimizer.hash_join_single_partition_threshold_rows 131072 +datafusion.optimizer.join_enumeration true +datafusion.optimizer.join_enumeration_limit 12 +datafusion.optimizer.join_enumeration_min_improvement 10 datafusion.optimizer.join_reordering true datafusion.optimizer.max_passes 3 datafusion.optimizer.prefer_existing_sort false @@ -484,8 +487,11 @@ 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 1048576 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 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.join_enumeration true When set to true, the physical plan optimizer enumerates join orders for subtrees of joins and picks the cheapest from cardinality estimates, considering bushy shapes as well as left-deep ones. Subtrees whose inputs lack row count statistics are left untouched. +datafusion.optimizer.join_enumeration_limit 12 Maximum inputs in a join subtree for which `join_enumeration` searches, at a cost of `O(3^n)`. Larger subtrees keep the planner's order, as do subtrees of more than 16 inputs regardless of this setting. +datafusion.optimizer.join_enumeration_min_improvement 10 How much cheaper an enumerated join order must be, in percent, before it replaces the order the planner produced. Estimates are often too close to tell two orders apart, so a small gain is not worth acting on. 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 Number of times that the optimizer will attempt to optimize the plan 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/datafusion/sqllogictest/test_files/join.slt.part b/datafusion/sqllogictest/test_files/join.slt.part index 00bea008fc2fc..958b556e9ced2 100644 --- a/datafusion/sqllogictest/test_files/join.slt.part +++ b/datafusion/sqllogictest/test_files/join.slt.part @@ -934,10 +934,12 @@ logical_plan 05)--SubqueryAlias: d 06)----TableScan: department projection=[dept_name] physical_plan -01)CrossJoinExec -02)--FilterExec: name@1 = Alice OR name@1 = Bob +01)ProjectionExec: expr=[emp_id@1 as emp_id, name@2 as name, dept_name@0 as dept_name] +02)--CrossJoinExec 03)----DataSourceExec: partitions=1, partition_sizes=[1] -04)--DataSourceExec: partitions=1, partition_sizes=[1] +04)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +05)------FilterExec: name@1 = Alice OR name@1 = Bob +06)--------DataSourceExec: partitions=1, partition_sizes=[1] # expect no row for Carol query ITT @@ -947,10 +949,10 @@ JOIN department AS d ON (e.name = 'Alice' OR e.name = 'Bob'); ---- 1 Alice HR -1 Alice Engineering -1 Alice Sales 2 Bob HR +1 Alice Engineering 2 Bob Engineering +1 Alice Sales 2 Bob Sales # OR conditions on Filter (not join filter) diff --git a/datafusion/sqllogictest/test_files/join_limit_pushdown.slt b/datafusion/sqllogictest/test_files/join_limit_pushdown.slt index 933b03e7ebd93..d0afb510dbc21 100644 --- a/datafusion/sqllogictest/test_files/join_limit_pushdown.slt +++ b/datafusion/sqllogictest/test_files/join_limit_pushdown.slt @@ -220,9 +220,9 @@ logical_plan 05)------TableScan: t2 projection=[x] 06)----TableScan: t3 projection=[p] physical_plan -01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(p@0, x@1)], projection=[a@1, x@2, p@0], fetch=2 +01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(a@0, x@0)], projection=[a@0, x@1, p@2], fetch=2 02)--DataSourceExec: partitions=1, partition_sizes=[1] -03)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(a@0, x@0)] +03)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(p@0, x@0)], projection=[x@1, p@0] 04)----DataSourceExec: partitions=1, partition_sizes=[1] 05)----DataSourceExec: partitions=1, partition_sizes=[1] diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index 7a706836f44d6..703529ec2f7e9 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -1337,13 +1337,12 @@ logical_plan 03)----TableScan: join_t1 projection=[t1_id] 04)----TableScan: join_t2 projection=[t2_id] physical_plan -01)AggregateExec: mode=FinalPartitioned, gby=[t1_id@0 as t1_id], aggr=[] -02)--RepartitionExec: partitioning=Hash([t1_id@0], 2), input_partitions=2 -03)----AggregateExec: mode=Partial, gby=[t1_id@0 as t1_id], aggr=[] -04)------HashJoinExec: mode=CollectLeft, join_type=LeftSemi, on=[(t1_id@0, t2_id@0)] -05)--------DataSourceExec: partitions=1, partition_sizes=[1] -06)--------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1 -07)----------DataSourceExec: partitions=1, partition_sizes=[1] +01)AggregateExec: mode=SinglePartitioned, gby=[t1_id@0 as t1_id], aggr=[] +02)--HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(t1_id@0, t2_id@0)] +03)----RepartitionExec: partitioning=Hash([t1_id@0], 2), input_partitions=1 +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)----RepartitionExec: partitioning=Hash([t2_id@0], 2), input_partitions=1 +06)------DataSourceExec: partitions=1, partition_sizes=[1] statement ok set datafusion.explain.logical_plan_only = true; @@ -1431,13 +1430,12 @@ logical_plan 06)--------TableScan: join_t2 projection=[t2_id] physical_plan 01)ProjectionExec: expr=[count(Int64(1))@1 as count(*)] -02)--AggregateExec: mode=FinalPartitioned, gby=[t1_id@0 as t1_id], aggr=[count(Int64(1))] -03)----RepartitionExec: partitioning=Hash([t1_id@0], 2), input_partitions=2 -04)------AggregateExec: mode=Partial, gby=[t1_id@0 as t1_id], aggr=[count(Int64(1))] -05)--------HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(t1_id@0, t2_id@0)], projection=[t1_id@0] -06)----------DataSourceExec: partitions=1, partition_sizes=[1] -07)----------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1 -08)------------DataSourceExec: partitions=1, partition_sizes=[1] +02)--AggregateExec: mode=SinglePartitioned, gby=[t1_id@0 as t1_id], aggr=[count(Int64(1))] +03)----HashJoinExec: mode=Partitioned, join_type=Inner, on=[(t1_id@0, t2_id@0)], projection=[t1_id@0] +04)------RepartitionExec: partitioning=Hash([t1_id@0], 2), input_partitions=1 +05)--------DataSourceExec: partitions=1, partition_sizes=[1] +06)------RepartitionExec: partitioning=Hash([t2_id@0], 2), input_partitions=1 +07)--------DataSourceExec: partitions=1, partition_sizes=[1] query TT EXPLAIN diff --git a/datafusion/sqllogictest/test_files/statistics_registry.slt b/datafusion/sqllogictest/test_files/statistics_registry.slt index 89258bec299c1..afb5ae12e5ae7 100644 --- a/datafusion/sqllogictest/test_files/statistics_registry.slt +++ b/datafusion/sqllogictest/test_files/statistics_registry.slt @@ -38,6 +38,13 @@ set datafusion.optimizer.hash_join_single_partition_threshold = 1; statement ok set datafusion.optimizer.hash_join_single_partition_threshold_rows = 1; +# This file tests the build side decision, so cost-based join order enumeration +# is kept out of the way: with it on, both cases below settle on the same +# (cheaper) join order and the registry's effect on the build side is no longer +# visible in the plan. Enumeration is covered on its own at the end of the file. +statement ok +set datafusion.optimizer.join_enumeration = false; + # -- Create test data -------------------------------------------------------- query I @@ -153,6 +160,37 @@ JOIN dim_small d ON o.small_id = d.small_id; ---- 66 +# -- With join order enumeration --------------------------------------------- +# The `FROM` order joins customers with orders first, for a 66 row intermediate. +# Enumeration instead joins orders with dim_small first, which produces 10 rows +# because every order matches exactly one dim_small row. + +statement ok +set datafusion.optimizer.join_enumeration = true; + +query TT +EXPLAIN SELECT o.order_id, c.region_id, d.label +FROM customers c +JOIN orders o ON c.customer_id = o.customer_id +JOIN dim_small d ON o.small_id = d.small_id; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(customer_id@0, customer_id@0)], projection=[order_id@3, region_id@1, label@4] +02)--RepartitionExec: partitioning=Hash([customer_id@0], 4), input_partitions=1, maintains_sort_order=true +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/customers.parquet]]}, projection=[customer_id, region_id], output_ordering=[region_id@1 ASC NULLS LAST], file_type=parquet +04)--RepartitionExec: partitioning=Hash([customer_id@0], 4), input_partitions=1 +05)----HashJoinExec: mode=Partitioned, join_type=Inner, on=[(small_id@2, small_id@0)], projection=[customer_id@1, order_id@0, label@4] +06)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/orders.parquet]]}, projection=[order_id, customer_id, small_id], output_ordering=[order_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible +07)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/dim_small.parquet]]}, projection=[small_id, label], output_ordering=[small_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + +query III +SELECT count(*), min(o.order_id), max(d.label) +FROM customers c +JOIN orders o ON c.customer_id = o.customer_id +JOIN dim_small d ON o.small_id = d.small_id; +---- +66 1 10 + # -- Cleanup ----------------------------------------------------------------- statement ok @@ -162,7 +200,7 @@ statement ok set datafusion.optimizer.use_statistics_registry = false; statement ok -set datafusion.optimizer.hash_join_single_partition_threshold = 1048576; +set datafusion.optimizer.hash_join_single_partition_threshold = 4194304; statement ok set datafusion.optimizer.hash_join_single_partition_threshold_rows = 131072; diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index e02ada03fc413..d6b19a43878fb 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -170,10 +170,13 @@ The following configuration settings are available: | datafusion.optimizer.max_passes | 3 | Number of times that the optimizer will attempt to optimize the plan | | datafusion.optimizer.top_down_join_key_reordering | true | When set to true, the physical plan optimizer will run a top down process to reorder the join keys | | 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.join_enumeration | true | When set to true, the physical plan optimizer enumerates join orders for subtrees of joins and picks the cheapest from cardinality estimates, considering bushy shapes as well as left-deep ones. Subtrees whose inputs lack row count statistics are left untouched. | +| datafusion.optimizer.join_enumeration_min_improvement | 10 | How much cheaper an enumerated join order must be, in percent, before it replaces the order the planner produced. Estimates are often too close to tell two orders apart, so a small gain is not worth acting on. | +| datafusion.optimizer.join_enumeration_limit | 12 | Maximum inputs in a join subtree for which `join_enumeration` searches, at a cost of `O(3^n)`. Larger subtrees keep the planner's order, as do subtrees of more than 16 inputs regardless of this setting. | | datafusion.optimizer.use_statistics_registry | false | When set to true, the physical plan optimizer uses the pluggable `StatisticsRegistry` for statistics propagation across operators. This enables more accurate cardinality estimates compared to each operator's built-in `partition_statistics`. | | 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 | 1048576 | 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 | 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_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: |