uqa_planner/statement_planning/
statistics.rs1use 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
19pub trait StatisticsTableState: Send + Sync {}
21
22pub 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}
31pub 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}