Skip to main content

uqa_planner/statement_planning/
statistics.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Catalog statistics and retrieval access estimates with first-error preservation.
8
9use crate::{ColumnStats, LocalAccessEstimate};
10use std::collections::BTreeMap;
11use std::sync::Arc;
12use uqa_core::catalog_index::CatalogIndexRow;
13use uqa_sql::{
14    ast::OperatorJoinRelations,
15    semantics::volatility::{self, VolatilityCatalog},
16    SQLError, ScalarExpr,
17};
18
19/// Retain the loaded table generation across its dependent metadata reads.
20pub trait StatisticsTableState: Send + Sync {}
21
22/// Loaded metadata needed to estimate relational access without depending on a storage provider.
23pub trait PlannerStatisticsCatalog {
24    fn storage_table(&self, table: &str) -> Result<Option<Arc<dyn StatisticsTableState>>, String>;
25    fn hierarchy_scan_tables(&self, table: &str) -> Result<Vec<String>, SQLError>;
26    fn table_row_count(&self, table: &str) -> Result<u64, SQLError>;
27    fn column_statistics(&self, table: &str) -> Result<BTreeMap<String, ColumnStats>, String>;
28    fn resolved_table_name(&self, table: &str) -> Result<Option<String>, SQLError>;
29    fn catalog_indexes(&self) -> Result<Vec<CatalogIndexRow>, SQLError>;
30}
31/// Access estimates supplied by the retrieval planner selected by the composition boundary.
32pub trait RetrievalSourceCosting {
33    fn operator_join_access(
34        &self,
35        name: &str,
36        relations: Option<&OperatorJoinRelations>,
37        args: &[ScalarExpr],
38    ) -> Result<LocalAccessEstimate, SQLError>;
39    fn local_access(
40        &self,
41        table: &str,
42        predicate: &ScalarExpr,
43    ) -> Result<Option<LocalAccessEstimate>, SQLError>;
44}
45#[derive(Clone, Copy)]
46pub struct StatementStatisticsContext<'a> {
47    pub catalog: &'a dyn PlannerStatisticsCatalog,
48    pub volatility: &'a dyn VolatilityCatalog,
49    pub retrieval: &'a dyn RetrievalSourceCosting,
50}
51pub(super) struct CatalogSourceStatistics<'a> {
52    pub(super) context: StatementStatisticsContext<'a>,
53    pub(super) error: &'a std::cell::RefCell<Option<SQLError>>,
54}
55
56impl CatalogSourceStatistics<'_> {
57    fn record_error(&self, error: SQLError) {
58        if self.error.borrow().is_none() {
59            *self.error.borrow_mut() = Some(error);
60        }
61    }
62}
63
64impl crate::SourceStatistics for CatalogSourceStatistics<'_> {
65    fn relation_statistics(&self, table: &str) -> Option<crate::RelationStats> {
66        match self.context.catalog.storage_table(table) {
67            Ok(None) => None,
68            Ok(Some(_)) => match (
69                hierarchy_row_count(self.context.catalog, table),
70                self.context.catalog.column_statistics(table),
71            ) {
72                (Ok(row_count), Ok(columns)) => Some(crate::RelationStats { row_count, columns }),
73                (Err(error), _) => {
74                    self.record_error(error);
75                    None
76                }
77                (_, Err(error)) => {
78                    self.record_error(SQLError::Internal(format!(
79                        "read optimizer statistics for `{table}`: {error}"
80                    )));
81                    None
82                }
83            },
84            Err(error) => {
85                self.record_error(SQLError::Internal(format!(
86                    "resolve optimizer storage table `{table}`: {error}"
87                )));
88                None
89            }
90        }
91    }
92
93    fn source_access_estimate(
94        &self,
95        source: &crate::SourcePlan,
96    ) -> Option<crate::LocalAccessEstimate> {
97        let crate::SourcePlan::Function {
98            name,
99            relations,
100            args,
101            ..
102        } = source
103        else {
104            return None;
105        };
106        if args.iter().any(|argument| {
107            argument.contains_parameter()
108                || volatility::expr_contains_volatile_function(self.context.volatility, argument)
109        }) {
110            return None;
111        }
112        let identity = name.to_ascii_lowercase();
113        let lower = uqa_sql::semantics::builtin_function_dispatch_name(&identity);
114        if !uqa_sql::registry::is_operator_join_table_function(&lower) {
115            return None;
116        }
117        match self
118            .context
119            .retrieval
120            .operator_join_access(&lower, relations.as_ref(), args)
121        {
122            Ok(estimate) => Some(estimate),
123            Err(error) => {
124                self.record_error(error);
125                None
126            }
127        }
128    }
129
130    fn local_access_estimate(
131        &self,
132        table: &str,
133        predicate: &uqa_sql::ScalarExpr,
134    ) -> Option<crate::LocalAccessEstimate> {
135        if volatility::expr_contains_volatile_function(self.context.volatility, predicate) {
136            return None;
137        }
138        if predicate.contains_parameter() {
139            return match super::parameterized::parameterized_access(self, table, predicate) {
140                Ok(estimate) => estimate,
141                Err(error) => {
142                    self.record_error(error);
143                    None
144                }
145            };
146        }
147        match self.context.catalog.storage_table(table) {
148            Ok(Some(_)) => {}
149            Ok(None) => return None,
150            Err(error) => {
151                self.record_error(SQLError::Internal(format!(
152                    "resolve optimizer storage table `{table}`: {error}"
153                )));
154                return None;
155            }
156        }
157        match self.context.retrieval.local_access(table, predicate) {
158            Ok(estimate) => estimate,
159            Err(error) => {
160                self.record_error(error);
161                None
162            }
163        }
164    }
165}
166
167fn hierarchy_row_count(
168    catalog: &dyn PlannerStatisticsCatalog,
169    table: &str,
170) -> Result<u64, SQLError> {
171    let mut total = 0_u64;
172    for member in catalog.hierarchy_scan_tables(table)? {
173        total = total
174            .checked_add(catalog.table_row_count(&member)?)
175            .ok_or_else(|| {
176                SQLError::Internal(format!(
177                    "optimizer hierarchy row count overflow for `{table}`"
178                ))
179            })?;
180    }
181    Ok(total)
182}