alamb commented on code in PR #24456: URL: https://github.com/apache/datafusion/pull/24456#discussion_r3824173460
########## datafusion/physical-optimizer/src/join_enumeration.rs: ########## @@ -0,0 +1,1651 @@ +// 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. +//! +//! 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. + +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, NullEquality, 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::execution_plan::replace_children_if_necessary; +use datafusion_physical_plan::joins::utils::{ + ColumnIndex, JoinFilter, max_distinct_count, +}; +use datafusion_physical_plan::joins::{ + CrossJoinExec, HashJoinExec, HashJoinExecBuilder, NestedLoopJoinExec, PartitionMode, + SortMergeJoinExec, +}; +use datafusion_physical_plan::operator_statistics::StatisticsRegistry; +use datafusion_physical_plan::projection::{ProjectionExec, all_alias_free_columns}; +use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; + +/// Chooses the shape of the join tree, before [`JoinSelection`] decides how each +/// join runs. +/// +/// [`JoinSelection`]: crate::join_selection::JoinSelection +#[derive(Default, Debug)] +pub struct JoinEnumeration {} + +impl JoinEnumeration { + #[expect(missing_docs)] + pub fn new() -> Self { + Self {} + } +} + +impl PhysicalOptimizerRule for JoinEnumeration { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.optimize_with_context(plan, &ConfigOnlyContext::new(config)) + } + + fn optimize_with_context( + &self, + plan: Arc<dyn ExecutionPlan>, + context: &dyn PhysicalOptimizerContext, + ) -> Result<Arc<dyn ExecutionPlan>> { + 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::<Statistics>::clone(s.base_arc())) + } else { + StatisticsContext::new().compute(plan, &StatisticsArgs::new()) + } + }; + Ok(enumerate_join_order(&plan, config, &mut stats)?.unwrap_or(plan)) + } + + fn name(&self) -> &str { + "join_enumeration" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// 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. +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<Arc<Statistics>> + 'a; + +/// A bitmask over relation indices. +type RelSet = u64; + +fn bit(rel: usize) -> RelSet { + 1u64 << rel +} + +fn iter_rels(mask: RelSet) -> impl Iterator<Item = usize> { + std::iter::successors(Some(mask), |m| Some(m & m.wrapping_sub(1))) + .take_while(|m| *m != 0) + .map(|m| m.trailing_zeros() as usize) +} + +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)] +struct ColRef { + rel: usize, + col: usize, +} + +/// What a relation contributes to the join. +#[derive(Debug)] +enum Role { + Output, + /// The quantified side of a semi or anti join, which filters instead of + /// contributing columns. + Reducer(Reducer), +} + +#[derive(Debug)] +struct Reducer { + /// `true` for an anti join, which keeps the rows that do *not* match. + anti: bool, + /// Keys, as `(column of the filtered side, column index here)`. + keys: Vec<(ColRef, usize)>, + /// Relations the keys reference; this reducer applies only to a set covering them. + required: RelSet, +} + +/// One leaf of the join graph: a subplan the enumerator does not look inside. +#[derive(Debug)] +struct Relation { + plan: Arc<dyn ExecutionPlan>, + /// Estimated row count, clamped to at least 1. + rows: f64, + /// Estimated bytes per row, when the input reports a size. + width: Option<f64>, + /// Per-column distinct value estimate, clamped to `[1, rows]`. + ndv: Vec<f64>, + role: Role, +} + +/// An equi-join predicate `left = right` between two distinct relations. +#[derive(Clone, Copy, Debug)] +struct Edge { + left: ColRef, + right: ColRef, +} + +/// A non-equi join predicate, moved along with its column references rewritten. +#[derive(Debug)] +struct Filter { + filter: JoinFilter, + /// The column each entry of the filter's intermediate schema comes from. + columns: Vec<ColRef>, + /// The relations those columns belong to. + required: RelSet, +} + +/// A connected set of joins as relations plus the predicates between them. +#[derive(Debug)] +struct JoinGraph { Review Comment: we finally have a JoinGraph in DataFusion! (this will probably get crazy, so I recommend it put in its own module to start) ########## datafusion/physical-optimizer/src/join_enumeration.rs: ########## @@ -0,0 +1,1651 @@ +// 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. +//! +//! 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. + +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, NullEquality, 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::execution_plan::replace_children_if_necessary; +use datafusion_physical_plan::joins::utils::{ + ColumnIndex, JoinFilter, max_distinct_count, +}; +use datafusion_physical_plan::joins::{ + CrossJoinExec, HashJoinExec, HashJoinExecBuilder, NestedLoopJoinExec, PartitionMode, + SortMergeJoinExec, +}; +use datafusion_physical_plan::operator_statistics::StatisticsRegistry; +use datafusion_physical_plan::projection::{ProjectionExec, all_alias_free_columns}; +use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; + +/// Chooses the shape of the join tree, before [`JoinSelection`] decides how each +/// join runs. +/// +/// [`JoinSelection`]: crate::join_selection::JoinSelection +#[derive(Default, Debug)] +pub struct JoinEnumeration {} + +impl JoinEnumeration { + #[expect(missing_docs)] + pub fn new() -> Self { + Self {} + } +} + +impl PhysicalOptimizerRule for JoinEnumeration { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.optimize_with_context(plan, &ConfigOnlyContext::new(config)) + } + + fn optimize_with_context( + &self, + plan: Arc<dyn ExecutionPlan>, + context: &dyn PhysicalOptimizerContext, + ) -> Result<Arc<dyn ExecutionPlan>> { + 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::<Statistics>::clone(s.base_arc())) + } else { + StatisticsContext::new().compute(plan, &StatisticsArgs::new()) + } + }; + Ok(enumerate_join_order(&plan, config, &mut stats)?.unwrap_or(plan)) + } + + fn name(&self) -> &str { + "join_enumeration" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// 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. +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<Arc<Statistics>> + 'a; + +/// A bitmask over relation indices. +type RelSet = u64; + +fn bit(rel: usize) -> RelSet { + 1u64 << rel +} + +fn iter_rels(mask: RelSet) -> impl Iterator<Item = usize> { + std::iter::successors(Some(mask), |m| Some(m & m.wrapping_sub(1))) + .take_while(|m| *m != 0) + .map(|m| m.trailing_zeros() as usize) +} + +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)] +struct ColRef { + rel: usize, Review Comment: would help to document what this is an index into ########## datafusion/physical-optimizer/src/join_enumeration.rs: ########## @@ -0,0 +1,1651 @@ +// 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. +//! +//! 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. + +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, NullEquality, 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::execution_plan::replace_children_if_necessary; +use datafusion_physical_plan::joins::utils::{ + ColumnIndex, JoinFilter, max_distinct_count, +}; +use datafusion_physical_plan::joins::{ + CrossJoinExec, HashJoinExec, HashJoinExecBuilder, NestedLoopJoinExec, PartitionMode, + SortMergeJoinExec, +}; +use datafusion_physical_plan::operator_statistics::StatisticsRegistry; +use datafusion_physical_plan::projection::{ProjectionExec, all_alias_free_columns}; +use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; + +/// Chooses the shape of the join tree, before [`JoinSelection`] decides how each +/// join runs. +/// +/// [`JoinSelection`]: crate::join_selection::JoinSelection +#[derive(Default, Debug)] +pub struct JoinEnumeration {} + +impl JoinEnumeration { + #[expect(missing_docs)] + pub fn new() -> Self { + Self {} + } +} + +impl PhysicalOptimizerRule for JoinEnumeration { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.optimize_with_context(plan, &ConfigOnlyContext::new(config)) + } + + fn optimize_with_context( + &self, + plan: Arc<dyn ExecutionPlan>, + context: &dyn PhysicalOptimizerContext, + ) -> Result<Arc<dyn ExecutionPlan>> { + 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::<Statistics>::clone(s.base_arc())) + } else { + StatisticsContext::new().compute(plan, &StatisticsArgs::new()) + } + }; + Ok(enumerate_join_order(&plan, config, &mut stats)?.unwrap_or(plan)) + } + + fn name(&self) -> &str { + "join_enumeration" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// 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. +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<Arc<Statistics>> + 'a; + +/// A bitmask over relation indices. +type RelSet = u64; + +fn bit(rel: usize) -> RelSet { + 1u64 << rel +} + +fn iter_rels(mask: RelSet) -> impl Iterator<Item = usize> { + std::iter::successors(Some(mask), |m| Some(m & m.wrapping_sub(1))) + .take_while(|m| *m != 0) + .map(|m| m.trailing_zeros() as usize) +} + +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)] +struct ColRef { + rel: usize, + col: usize, +} + +/// What a relation contributes to the join. +#[derive(Debug)] +enum Role { + Output, + /// The quantified side of a semi or anti join, which filters instead of + /// contributing columns. + Reducer(Reducer), +} + +#[derive(Debug)] +struct Reducer { + /// `true` for an anti join, which keeps the rows that do *not* match. + anti: bool, + /// Keys, as `(column of the filtered side, column index here)`. + keys: Vec<(ColRef, usize)>, + /// Relations the keys reference; this reducer applies only to a set covering them. + required: RelSet, +} + +/// One leaf of the join graph: a subplan the enumerator does not look inside. +#[derive(Debug)] +struct Relation { + plan: Arc<dyn ExecutionPlan>, + /// Estimated row count, clamped to at least 1. + rows: f64, + /// Estimated bytes per row, when the input reports a size. + width: Option<f64>, + /// Per-column distinct value estimate, clamped to `[1, rows]`. + ndv: Vec<f64>, + role: Role, +} + +/// An equi-join predicate `left = right` between two distinct relations. +#[derive(Clone, Copy, Debug)] +struct Edge { + left: ColRef, + right: ColRef, +} + +/// A non-equi join predicate, moved along with its column references rewritten. +#[derive(Debug)] +struct Filter { + filter: JoinFilter, + /// The column each entry of the filter's intermediate schema comes from. + columns: Vec<ColRef>, + /// The relations those columns belong to. + required: RelSet, +} + +/// A connected set of joins as relations plus the predicates between them. Review Comment: I think it might make sense to point out that this is an alternative representation of at a sub clause in a query -- basically a `SELECT .... ` but with nodes representing tables and edges representing joins between them It looks like this representation takes some sub tree of joins and then replaces the JoinExec (maybe?) It wold help to describe with an example how an input ExecutionPlan goes into a JoinGraph ########## datafusion/physical-optimizer/src/join_enumeration.rs: ########## @@ -0,0 +1,1651 @@ +// 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. +//! +//! 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. + +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, NullEquality, 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::execution_plan::replace_children_if_necessary; +use datafusion_physical_plan::joins::utils::{ + ColumnIndex, JoinFilter, max_distinct_count, +}; +use datafusion_physical_plan::joins::{ + CrossJoinExec, HashJoinExec, HashJoinExecBuilder, NestedLoopJoinExec, PartitionMode, + SortMergeJoinExec, +}; +use datafusion_physical_plan::operator_statistics::StatisticsRegistry; +use datafusion_physical_plan::projection::{ProjectionExec, all_alias_free_columns}; +use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; + +/// Chooses the shape of the join tree, before [`JoinSelection`] decides how each +/// join runs. +/// +/// [`JoinSelection`]: crate::join_selection::JoinSelection +#[derive(Default, Debug)] +pub struct JoinEnumeration {} + +impl JoinEnumeration { + #[expect(missing_docs)] + pub fn new() -> Self { + Self {} + } +} + +impl PhysicalOptimizerRule for JoinEnumeration { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.optimize_with_context(plan, &ConfigOnlyContext::new(config)) + } + + fn optimize_with_context( + &self, + plan: Arc<dyn ExecutionPlan>, + context: &dyn PhysicalOptimizerContext, + ) -> Result<Arc<dyn ExecutionPlan>> { + 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::<Statistics>::clone(s.base_arc())) + } else { + StatisticsContext::new().compute(plan, &StatisticsArgs::new()) + } + }; + Ok(enumerate_join_order(&plan, config, &mut stats)?.unwrap_or(plan)) + } + + fn name(&self) -> &str { + "join_enumeration" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// 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. +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<Arc<Statistics>> + 'a; + +/// A bitmask over relation indices. +type RelSet = u64; + +fn bit(rel: usize) -> RelSet { + 1u64 << rel +} + +fn iter_rels(mask: RelSet) -> impl Iterator<Item = usize> { + std::iter::successors(Some(mask), |m| Some(m & m.wrapping_sub(1))) + .take_while(|m| *m != 0) + .map(|m| m.trailing_zeros() as usize) +} + +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)] +struct ColRef { + rel: usize, + col: usize, +} + +/// What a relation contributes to the join. +#[derive(Debug)] +enum Role { + Output, + /// The quantified side of a semi or anti join, which filters instead of + /// contributing columns. + Reducer(Reducer), +} + +#[derive(Debug)] +struct Reducer { + /// `true` for an anti join, which keeps the rows that do *not* match. + anti: bool, + /// Keys, as `(column of the filtered side, column index here)`. + keys: Vec<(ColRef, usize)>, + /// Relations the keys reference; this reducer applies only to a set covering them. + required: RelSet, +} + +/// One leaf of the join graph: a subplan the enumerator does not look inside. +#[derive(Debug)] +struct Relation { + plan: Arc<dyn ExecutionPlan>, + /// Estimated row count, clamped to at least 1. + rows: f64, + /// Estimated bytes per row, when the input reports a size. + width: Option<f64>, + /// Per-column distinct value estimate, clamped to `[1, rows]`. + ndv: Vec<f64>, + role: Role, +} + +/// An equi-join predicate `left = right` between two distinct relations. +#[derive(Clone, Copy, Debug)] +struct Edge { + left: ColRef, + right: ColRef, +} + +/// A non-equi join predicate, moved along with its column references rewritten. +#[derive(Debug)] +struct Filter { + filter: JoinFilter, + /// The column each entry of the filter's intermediate schema comes from. + columns: Vec<ColRef>, + /// The relations those columns belong to. + required: RelSet, +} + +/// A connected set of joins as relations plus the predicates between them. +#[derive(Debug)] +struct JoinGraph { + relations: Vec<Relation>, + edges: Vec<Edge>, + filters: Vec<Filter>, + /// Columns the original subtree emitted; the rebuilt one reproduces this exactly. + output: Vec<ColRef>, + /// Null handling shared by the subtree's joins. A join that differs becomes a + /// relation instead. + null_equality: Option<NullEquality>, + /// 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. + original_nodes: Vec<(RelSet, RelSet)>, + /// The relations that are reducers rather than ordinary inputs. + reducers: RelSet, + /// Which join operator the subtree used, and which the rebuild emits. + kind: Option<JoinKind>, +} + +impl JoinGraph { + fn ndv(&self, col: ColRef) -> f64 { + self.relations[col.rel].ndv[col.col] + } + + fn all(&self) -> RelSet { + (0..self.relations.len()).fold(0, |mask, rel| mask | bit(rel)) + } + + fn reducer(&self, rel: usize) -> Option<&Reducer> { + match &self.relations[rel].role { + Role::Reducer(reducer) => Some(reducer), + Role::Output => None, + } + } + + fn kind(&self) -> JoinKind { + self.kind.unwrap_or(JoinKind::Hash) + } + + fn null_equality(&self) -> NullEquality { + self.null_equality + .unwrap_or(NullEquality::NullEqualsNothing) + } +} + +/// A valid way of combining two relation sets. +#[derive(Clone, Copy, Debug)] +enum Combine { + Inner, + /// A semi or anti join applying `reducer` to the opposite set. + Reducer { + reducer: usize, + }, +} + +/// Cardinality and cost estimates over the subsets of a [`JoinGraph`]. +struct CostModel<'a> { Review Comment: Can we please make this a trait (so that other people can plug in their own cost models, based on whatever statistics they may have, or other knowledge they have about their data and the joins? Something like this ```rust trait JoinCostModel { /// Estimated rows from joining every relation in `mask`. fn cardinality(&self, mask: RelSet) -> f64; /// Whether `left` and `right` can be combined, and how (inner join, or one /// side applying as a reducer). fn combine(&self, left: RelSet, right: RelSet) -> Option<Combine>; /// Cost of joining `left` (partitioned as `left_part`) with `right` /// (`right_part`), for each way of exchanging their inputs. fn exchanges( &self, left: RelSet, right: RelSet, left_part: PartSet, right_part: PartSet, collect_only: Option<RelSet>, ) -> Vec<(f64, PartSet, RelSet, PartitionMode)>; } ``` ########## 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 Review Comment: this set of settings is exactly why I am hesitant to put something too complicated into the datafusion core -- the join ordering algorithms can get so out of hand complicated I think we need to find a way to keep the core (relatively) simple and leave API hooks to support the more complicated cases I have more suggestions on how to do this below ########## datafusion/physical-optimizer/src/join_enumeration.rs: ########## @@ -0,0 +1,1651 @@ +// 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. +//! +//! 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. + +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, NullEquality, 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::execution_plan::replace_children_if_necessary; +use datafusion_physical_plan::joins::utils::{ + ColumnIndex, JoinFilter, max_distinct_count, +}; +use datafusion_physical_plan::joins::{ + CrossJoinExec, HashJoinExec, HashJoinExecBuilder, NestedLoopJoinExec, PartitionMode, + SortMergeJoinExec, +}; +use datafusion_physical_plan::operator_statistics::StatisticsRegistry; +use datafusion_physical_plan::projection::{ProjectionExec, all_alias_free_columns}; +use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; + +/// Chooses the shape of the join tree, before [`JoinSelection`] decides how each +/// join runs. +/// +/// [`JoinSelection`]: crate::join_selection::JoinSelection +#[derive(Default, Debug)] +pub struct JoinEnumeration {} + +impl JoinEnumeration { + #[expect(missing_docs)] + pub fn new() -> Self { + Self {} + } +} + +impl PhysicalOptimizerRule for JoinEnumeration { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.optimize_with_context(plan, &ConfigOnlyContext::new(config)) + } + + fn optimize_with_context( + &self, + plan: Arc<dyn ExecutionPlan>, + context: &dyn PhysicalOptimizerContext, + ) -> Result<Arc<dyn ExecutionPlan>> { + 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::<Statistics>::clone(s.base_arc())) + } else { + StatisticsContext::new().compute(plan, &StatisticsArgs::new()) + } + }; + Ok(enumerate_join_order(&plan, config, &mut stats)?.unwrap_or(plan)) + } + + fn name(&self) -> &str { + "join_enumeration" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// 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. +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<Arc<Statistics>> + 'a; + +/// A bitmask over relation indices. +type RelSet = u64; + +fn bit(rel: usize) -> RelSet { + 1u64 << rel +} + +fn iter_rels(mask: RelSet) -> impl Iterator<Item = usize> { + std::iter::successors(Some(mask), |m| Some(m & m.wrapping_sub(1))) + .take_while(|m| *m != 0) + .map(|m| m.trailing_zeros() as usize) +} + +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)] +struct ColRef { + rel: usize, + col: usize, +} + +/// What a relation contributes to the join. +#[derive(Debug)] +enum Role { + Output, + /// The quantified side of a semi or anti join, which filters instead of + /// contributing columns. + Reducer(Reducer), +} + +#[derive(Debug)] +struct Reducer { + /// `true` for an anti join, which keeps the rows that do *not* match. + anti: bool, + /// Keys, as `(column of the filtered side, column index here)`. + keys: Vec<(ColRef, usize)>, + /// Relations the keys reference; this reducer applies only to a set covering them. + required: RelSet, +} + +/// One leaf of the join graph: a subplan the enumerator does not look inside. +#[derive(Debug)] +struct Relation { + plan: Arc<dyn ExecutionPlan>, + /// Estimated row count, clamped to at least 1. + rows: f64, + /// Estimated bytes per row, when the input reports a size. + width: Option<f64>, + /// Per-column distinct value estimate, clamped to `[1, rows]`. + ndv: Vec<f64>, + role: Role, +} + +/// An equi-join predicate `left = right` between two distinct relations. +#[derive(Clone, Copy, Debug)] +struct Edge { + left: ColRef, + right: ColRef, +} + +/// A non-equi join predicate, moved along with its column references rewritten. +#[derive(Debug)] +struct Filter { + filter: JoinFilter, + /// The column each entry of the filter's intermediate schema comes from. + columns: Vec<ColRef>, + /// The relations those columns belong to. + required: RelSet, +} + +/// A connected set of joins as relations plus the predicates between them. Review Comment: Here is an example that Claude came up with example four tables, B goes through GROUP BY before the join: ``` HashJoin(a.x = b.x) ├── HashJoin(a.y = c.y) │ ├── Scan A │ └── Scan C └── HashJoin(b.z = d.z) ├── AggregateExec(group by: b.x, b.z; cnt = count(*)) │ └── Scan B └── Scan D ``` And the join representation ``` JoinGraph { kind: Hash, reducers: {} (none), filters: [] edge: A.y = C.y ┌───────────┐ ───────────────────────────── ┌───────────┐ │ rel 0: A │ │ rel 1: C │ │ (Scan A) │ │ (Scan C) │ └─────┬─────┘ └───────────┘ │ │ edge: A.x = Agg(B).col0 │ ┌─────┴───────────┐ │ rel 2: Agg(B) │ │ [OPAQUE] │ └────────┬─────────┘ │ │ edge: Agg(B).col1 = D.z │ ┌─────┴─────┐ │ rel 3: D │ │ (Scan D) │ └───────────┘ output: columns of A ++ C ++ Agg(B) ++ D (flattened left-to-right, depth-first) } ``` ########## datafusion/physical-optimizer/src/join_enumeration.rs: ########## @@ -0,0 +1,1651 @@ +// 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. +//! +//! 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. + +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, NullEquality, 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::execution_plan::replace_children_if_necessary; +use datafusion_physical_plan::joins::utils::{ + ColumnIndex, JoinFilter, max_distinct_count, +}; +use datafusion_physical_plan::joins::{ + CrossJoinExec, HashJoinExec, HashJoinExecBuilder, NestedLoopJoinExec, PartitionMode, + SortMergeJoinExec, +}; +use datafusion_physical_plan::operator_statistics::StatisticsRegistry; +use datafusion_physical_plan::projection::{ProjectionExec, all_alias_free_columns}; +use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; + +/// Chooses the shape of the join tree, before [`JoinSelection`] decides how each +/// join runs. +/// +/// [`JoinSelection`]: crate::join_selection::JoinSelection +#[derive(Default, Debug)] +pub struct JoinEnumeration {} + +impl JoinEnumeration { + #[expect(missing_docs)] + pub fn new() -> Self { + Self {} + } +} + +impl PhysicalOptimizerRule for JoinEnumeration { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.optimize_with_context(plan, &ConfigOnlyContext::new(config)) + } + + fn optimize_with_context( + &self, + plan: Arc<dyn ExecutionPlan>, + context: &dyn PhysicalOptimizerContext, + ) -> Result<Arc<dyn ExecutionPlan>> { + 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::<Statistics>::clone(s.base_arc())) + } else { + StatisticsContext::new().compute(plan, &StatisticsArgs::new()) + } + }; + Ok(enumerate_join_order(&plan, config, &mut stats)?.unwrap_or(plan)) + } + + fn name(&self) -> &str { + "join_enumeration" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// 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. +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<Arc<Statistics>> + 'a; + +/// A bitmask over relation indices. +type RelSet = u64; + +fn bit(rel: usize) -> RelSet { + 1u64 << rel +} + +fn iter_rels(mask: RelSet) -> impl Iterator<Item = usize> { + std::iter::successors(Some(mask), |m| Some(m & m.wrapping_sub(1))) + .take_while(|m| *m != 0) + .map(|m| m.trailing_zeros() as usize) +} + +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)] +struct ColRef { + rel: usize, + col: usize, +} + +/// What a relation contributes to the join. +#[derive(Debug)] +enum Role { + Output, + /// The quantified side of a semi or anti join, which filters instead of + /// contributing columns. + Reducer(Reducer), +} + +#[derive(Debug)] +struct Reducer { + /// `true` for an anti join, which keeps the rows that do *not* match. + anti: bool, + /// Keys, as `(column of the filtered side, column index here)`. + keys: Vec<(ColRef, usize)>, + /// Relations the keys reference; this reducer applies only to a set covering them. + required: RelSet, +} + +/// One leaf of the join graph: a subplan the enumerator does not look inside. +#[derive(Debug)] +struct Relation { + plan: Arc<dyn ExecutionPlan>, + /// Estimated row count, clamped to at least 1. + rows: f64, + /// Estimated bytes per row, when the input reports a size. + width: Option<f64>, + /// Per-column distinct value estimate, clamped to `[1, rows]`. + ndv: Vec<f64>, + role: Role, +} + +/// An equi-join predicate `left = right` between two distinct relations. +#[derive(Clone, Copy, Debug)] +struct Edge { + left: ColRef, + right: ColRef, +} + +/// A non-equi join predicate, moved along with its column references rewritten. +#[derive(Debug)] +struct Filter { + filter: JoinFilter, + /// The column each entry of the filter's intermediate schema comes from. + columns: Vec<ColRef>, + /// The relations those columns belong to. + required: RelSet, +} + +/// A connected set of joins as relations plus the predicates between them. +#[derive(Debug)] +struct JoinGraph { + relations: Vec<Relation>, + edges: Vec<Edge>, + filters: Vec<Filter>, Review Comment: I normally think of a join graph where the edges themselves have filters. But maybe that is implicit in the fact that the edges and filters are parallell lists -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
