Dandandan commented on code in PR #24456: URL: https://github.com/apache/datafusion/pull/24456#discussion_r3831868391
########## 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: need to have a look what makes sense here -- 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]
