Skip to main content

uqa_planner/optimizer/
api.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Optimizer configuration, statistics seam, and public entry points.
8
9use super::{
10    optimize_unified_plan, reorder_unified_plan_joins, AggregateClassifier, JoinGraphError,
11    RelationStats, ScalarExpr, SourcePlan, UnifiedPlan,
12};
13use crate::LocalAccessEstimate;
14
15/// Errors raised while simplifying expressions or choosing a physical join order.
16#[derive(Debug, thiserror::Error)]
17pub enum OptimizerError {
18    #[error(transparent)]
19    Expression(#[from] uqa_sql::SQLError),
20    #[error(transparent)]
21    JoinGraph(#[from] JoinGraphError),
22}
23
24pub type OptimizerResult<T> = Result<T, OptimizerError>;
25
26/// Evaluate a planner-proven constant using the runtime selected by the caller.
27pub type ConstantEvaluator = fn(&ScalarExpr) -> Result<uqa_core::Value, uqa_sql::SQLError>;
28
29#[derive(Debug, Clone)]
30pub struct OptimizerConfig {
31    pub enable_filter_pushdown: bool,
32    pub enable_boolean_simplify: bool,
33    pub enable_vector_threshold_merge: bool,
34    pub enable_join_reordering: bool,
35    /// Shared scalar execution supplied by the engine composition boundary.
36    pub constant_evaluator: ConstantEvaluator,
37}
38
39impl OptimizerConfig {
40    #[must_use]
41    pub const fn new(constant_evaluator: ConstantEvaluator) -> Self {
42        Self {
43            enable_filter_pushdown: true,
44            enable_boolean_simplify: true,
45            enable_vector_threshold_merge: true,
46            enable_join_reordering: true,
47            constant_evaluator,
48        }
49    }
50}
51
52/// Cardinality and column statistics used to cost base relations during
53/// join enumeration. Engines implement this against their live catalogue;
54/// callers without a catalogue still get deterministic DPccp enumeration
55/// from the optimizer's fallback cardinality.
56pub trait SourceStatistics {
57    fn relation_statistics(&self, table: &str) -> Option<RelationStats>;
58
59    /// Estimate a non-table source that can participate as one atom in an
60    /// inner-join region. Returning `None` keeps that region in SQL order.
61    fn source_access_estimate(&self, _source: &SourcePlan) -> Option<LocalAccessEstimate> {
62        None
63    }
64
65    /// Estimate a predicate that references only `table`. Returning `None`
66    /// delegates to the planner's scalar selectivity model.
67    fn local_access_estimate(
68        &self,
69        _table: &str,
70        _predicate: &ScalarExpr,
71    ) -> Option<LocalAccessEstimate> {
72        None
73    }
74}
75
76impl<F> SourceStatistics for F
77where
78    F: Fn(&str) -> Option<RelationStats>,
79{
80    fn relation_statistics(&self, table: &str) -> Option<RelationStats> {
81        self(table)
82    }
83}
84
85struct NoSourceStatistics;
86
87impl SourceStatistics for NoSourceStatistics {
88    fn relation_statistics(&self, _table: &str) -> Option<RelationStats> {
89        None
90    }
91}
92
93struct NoRegisteredAggregates;
94
95impl AggregateClassifier for NoRegisteredAggregates {
96    fn is_registered_aggregate(&self, _name: &str) -> bool {
97        false
98    }
99}
100
101/// Optimize a fully lowered plan using the built-in aggregate catalogue.
102pub fn optimize(plan: UnifiedPlan, config: &OptimizerConfig) -> OptimizerResult<UnifiedPlan> {
103    optimize_with_aggregates_and_statistics(
104        plan,
105        config,
106        &NoRegisteredAggregates,
107        &NoSourceStatistics,
108    )
109}
110
111/// Optimize a fully lowered plan while classifying engine-local aggregates.
112pub fn optimize_with_aggregates(
113    plan: UnifiedPlan,
114    config: &OptimizerConfig,
115    aggregates: &dyn AggregateClassifier,
116) -> OptimizerResult<UnifiedPlan> {
117    optimize_with_aggregates_and_statistics(plan, config, aggregates, &NoSourceStatistics)
118}
119
120/// Optimize a fully lowered plan with caller-provided relation statistics.
121pub fn optimize_with_statistics(
122    plan: UnifiedPlan,
123    config: &OptimizerConfig,
124    statistics: &dyn SourceStatistics,
125) -> OptimizerResult<UnifiedPlan> {
126    optimize_with_aggregates_and_statistics(plan, config, &NoRegisteredAggregates, statistics)
127}
128
129/// Optimize a fully lowered plan while classifying engine-local aggregates
130/// and costing join orders from the engine's relation statistics.
131pub fn optimize_with_aggregates_and_statistics(
132    mut plan: UnifiedPlan,
133    config: &OptimizerConfig,
134    aggregates: &dyn AggregateClassifier,
135    statistics: &dyn SourceStatistics,
136) -> OptimizerResult<UnifiedPlan> {
137    optimize_unified_plan(&mut plan, config, aggregates)?;
138    if config.enable_join_reordering {
139        reorder_unified_plan_joins(&mut plan, statistics)?;
140    }
141    Ok(plan)
142}