add statistics

Author: JanKaul

datafusion/catalog/src/default_table_source.rs

@@ -23,7 +23,7 @@ use std::sync::Arc;
 use crate::TableProvider;
 
 use arrow::datatypes::SchemaRef;
-use datafusion_common::{Constraints, internal_err};
+use datafusion_common::{Constraints, Statistics, internal_err};
 use datafusion_expr::{Expr, TableProviderFilterPushDown, TableSource, TableType};
 
 /// Implements [`TableSource`] for a [`TableProvider`]
@@ -77,6 +77,10 @@ impl TableSource for DefaultTableSource {
     fn get_column_default(&self, column: &str) -> Option<&Expr> {
         self.table_provider.get_column_default(column)
     }
+
+    fn statistics(&self) -> Option<Statistics> {
+        self.table_provider.statistics()
+    }
 }
 
 /// Wrap TableProvider in TableSource

datafusion/expr/src/table_source.rs

@@ -20,7 +20,7 @@
 use crate::{Expr, LogicalPlan};
 
 use arrow::datatypes::SchemaRef;
-use datafusion_common::{Constraints, Result};
+use datafusion_common::{Constraints, Result, Statistics};
 
 use std::{any::Any, borrow::Cow};
 
@@ -127,6 +127,14 @@ pub trait TableSource: Any + Sync + Send {
     fn get_column_default(&self, _column: &str) -> Option<&Expr> {
         None
     }
+
+    /// Get statistics for this table, if available
+    /// Although not presently used in mainline DataFusion, this allows implementation specific
+    /// behavior for downstream repositories, in conjunction with specialized optimizer rules to
+    /// perform operations such as re-ordering of joins.
+    fn statistics(&self) -> Option<Statistics> {
+        None
+    }
 }
 
 impl dyn TableSource {

datafusion/optimizer/src/reorder_join/cost.rs

@@ -15,20 +15,45 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use datafusion_common::{Result, plan_err};
-use datafusion_expr::{JoinType, LogicalPlan};
+use datafusion_common::{Column, Result, plan_err, stats::Precision};
+use datafusion_expr::{Expr, JoinType, LogicalPlan};
 
 use super::join_graph::Edge;
 
 pub trait JoinCostEstimator: std::fmt::Debug {
-    fn cardinality(&self, plan: &LogicalPlan) -> Option<f64> {
-        estimate_cardinality(plan).ok()
+    /// Cardinality of `plan`.
+    ///
+    /// - `column = None`: number of output rows of `plan`.
+    /// - `column = Some(c)`: number of distinct values of column `c`
+    ///   in `plan`'s output (NDV).
+    fn cardinality(&self, plan: &LogicalPlan, column: Option<&Column>) -> Option<f64> {
+        estimate_cardinality(plan, column).ok()
     }
 
-    fn selectivity(&self, edge: &Edge) -> f64 {
-        match edge.join_type {
+    /// Estimated selectivity of joining `left` with `right` via `edge`.
+    ///
+    /// Default: `1 / max(NDV(left.key), NDV(right.key))` for inner equi-joins
+    /// when both NDVs are available; otherwise a per-join-type constant.
+    fn selectivity(&self, edge: &Edge, left: &LogicalPlan, right: &LogicalPlan) -> f64 {
+        let fallback = match edge.join_type {
             JoinType::Inner => 0.1,
             _ => 1.0,
+        };
+        if edge.join_type != JoinType::Inner || edge.on.is_empty() {
+            return fallback;
+        }
+        // Use only the first equi-pair. Compounding pairwise selectivities
+        // under independence assumptions overestimates selectivity when
+        // composite-key columns are correlated, which is the common case.
+        let (a, b) = &edge.on[0];
+        let (Some(col_a), Some(col_b)) = (key_column(a), key_column(b)) else {
+            return fallback;
+        };
+        let ndv_a = ndv_for(self, col_a, left, right);
+        let ndv_b = ndv_for(self, col_b, left, right);
+        match (ndv_a, ndv_b) {
+            (Some(a), Some(b)) if a.max(b) > 0.0 => 1.0 / a.max(b),
+            _ => fallback,
         }
     }
 
@@ -43,29 +68,114 @@ pub struct DefaultCostEstimator;
 
 impl JoinCostEstimator for DefaultCostEstimator {}
 
-fn estimate_cardinality(plan: &LogicalPlan) -> Result<f64> {
+fn key_column(expr: &Expr) -> Option<&Column> {
+    match expr {
+        Expr::Column(c) => Some(c),
+        _ => None,
+    }
+}
+
+/// Look up NDV of `column` on whichever side (left or right) owns it.
+fn ndv_for<E: JoinCostEstimator + ?Sized>(
+    estimator: &E,
+    column: &Column,
+    left: &LogicalPlan,
+    right: &LogicalPlan,
+) -> Option<f64> {
+    if left.schema().has_column(column) {
+        estimator.cardinality(left, Some(column))
+    } else if right.schema().has_column(column) {
+        estimator.cardinality(right, Some(column))
+    } else {
+        None
+    }
+}
+
+fn estimate_cardinality(plan: &LogicalPlan, column: Option<&Column>) -> Result<f64> {
     match plan {
-        LogicalPlan::Filter(filter) => {
-            let input_cardinality = estimate_cardinality(&filter.input)?;
-            Ok(0.1 * input_cardinality)
-        }
-        LogicalPlan::Aggregate(agg) => {
-            let input_cardinality = estimate_cardinality(&agg.input)?;
-            Ok(0.1 * input_cardinality)
-        }
-        LogicalPlan::TableScan(_) => {
-            // The logical-plan-level `TableSource` trait doesn't expose row
-            // statistics. To use cardinalities for base relations, override
-            // `JoinCostEstimator::cardinality`.
-            plan_err!(
-                "Default JoinCostEstimator cannot size TableScan; \
-                 override `cardinality` to provide statistics"
-            )
+        LogicalPlan::Filter(filter) => match column {
+            None => {
+                let input = estimate_cardinality(&filter.input, None)?;
+                Ok(0.1 * input)
+            }
+            Some(c) => {
+                // NDV is bounded above by the input's NDV and by the
+                // surviving row count.
+                let ndv_in = estimate_cardinality(&filter.input, Some(c))?;
+                let rows_out = estimate_cardinality(plan, None).unwrap_or(ndv_in);
+                Ok(ndv_in.min(rows_out))
+            }
+        },
+        LogicalPlan::Aggregate(agg) => match column {
+            None => {
+                let input = estimate_cardinality(&agg.input, None)?;
+                Ok(0.1 * input)
+            }
+            Some(c) => {
+                // Group-by keys are unique in the aggregate's output, so
+                // NDV(group_key) equals the post-aggregate row count.
+                let is_group_key = agg.group_expr.iter().any(|e| match e {
+                    Expr::Column(g) => g.name == c.name && g.relation == c.relation,
+                    _ => false,
+                });
+                if is_group_key {
+                    estimate_cardinality(plan, None)
+                } else {
+                    plan_err!(
+                        "Cannot estimate NDV of non-group-by column \
+                         `{}` through Aggregate",
+                        c.name
+                    )
+                }
+            }
+        },
+        LogicalPlan::TableScan(scan) => {
+            let stats = scan.source.statistics().ok_or_else(|| {
+                datafusion_common::DataFusionError::Plan(format!(
+                    "TableSource for `{}` does not expose statistics",
+                    scan.table_name
+                ))
+            })?;
+            match column {
+                None => match stats.num_rows {
+                    Precision::Exact(n) | Precision::Inexact(n) => Ok(n as f64),
+                    Precision::Absent => plan_err!(
+                        "TableSource for `{}` does not provide a row count",
+                        scan.table_name
+                    ),
+                },
+                Some(c) => {
+                    // `column_statistics` is indexed by the source schema
+                    // (pre-projection), so resolve the column there.
+                    let idx = scan.source.schema().index_of(&c.name).map_err(|_| {
+                        datafusion_common::DataFusionError::Plan(format!(
+                            "Column `{}` not found in source schema of `{}`",
+                            c.name, scan.table_name
+                        ))
+                    })?;
+                    let col_stats =
+                        stats.column_statistics.get(idx).ok_or_else(|| {
+                            datafusion_common::DataFusionError::Plan(format!(
+                                "Column statistics missing for index {idx} \
+                                 on `{}`",
+                                scan.table_name
+                            ))
+                        })?;
+                    match col_stats.distinct_count {
+                        Precision::Exact(n) | Precision::Inexact(n) => Ok(n as f64),
+                        Precision::Absent => plan_err!(
+                            "Column `{}` on `{}` has no distinct-count statistic",
+                            c.name,
+                            scan.table_name
+                        ),
+                    }
+                }
+            }
         }
         x => {
             let inputs = x.inputs();
             if inputs.len() == 1 {
-                estimate_cardinality(inputs[0])
+                estimate_cardinality(inputs[0], column)
             } else {
                 plan_err!("Cannot estimate cardinality for plan with multiple inputs")
             }

datafusion/optimizer/src/reorder_join/left_deep_join_plan.rs

@@ -207,7 +207,8 @@ impl<'graph> PrecedenceTreeNode<'graph> {
         let node = query_graph
             .get_node(node_id)
             .ok_or_else(|| plan_datafusion_err!("Root node not found"))?;
-        let input_cardinality = cost_estimator.cardinality(&node.plan).unwrap_or(1.0);
+        let input_cardinality =
+            cost_estimator.cardinality(&node.plan, None).unwrap_or(1.0);
 
         let children = node
             .connections()
@@ -220,7 +221,9 @@ impl<'graph> PrecedenceTreeNode<'graph> {
                     .find(|x| *x != node_id && remaining.contains(x))?;
 
                 remaining.remove(&other);
-                let child_selectivity = cost_estimator.selectivity(edge);
+                let other_plan = &query_graph.get_node(other)?.plan;
+                let child_selectivity =
+                    cost_estimator.selectivity(edge, &node.plan, other_plan);
                 Some(PrecedenceTreeNode::from_query_node(
                     other,
                     child_selectivity,