krishiv_plan/optimizer/
stats.rs

1//! Cost-based optimization (CBO) infrastructure: stats registry and
2//! NDV-aware cost model.
3//!
4//! See [`TableStatsRegistry`], [`TableCboStats`], and
5//! [`CboCostModel`] for the user-facing types. The cost model is
6//! plug-compatible with the static one in [`super::StaticCostModel`].
7
8use std::collections::HashMap;
9use std::sync::{Arc, RwLock};
10
11use crate::{LogicalPlan, NodeOp};
12
13use super::CostModel;
14
15/// Per-column statistics collected by `ANALYZE TABLE … FOR COLUMNS`.
16#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
17pub struct ColumnCboStats {
18    pub name: String,
19    /// Approximate number of distinct values (HLL-based).
20    pub ndv: Option<u64>,
21    /// Rendered minimum value (display form; type-erased on purpose).
22    pub min: Option<String>,
23    /// Rendered maximum value.
24    pub max: Option<String>,
25    pub null_count: Option<u64>,
26}
27
28/// Catalog statistics the cost model needs for one table.
29#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
30pub struct TableCboStats {
31    pub table: String,
32    pub row_count: Option<u64>,
33    pub ndv: Option<u64>,
34    pub avg_row_bytes: Option<u64>,
35    /// Column-level stats (empty unless ANALYZE ran with FOR COLUMNS).
36    #[serde(default)]
37    pub columns: Vec<ColumnCboStats>,
38}
39
40impl TableCboStats {
41    pub fn new(table: impl Into<String>) -> Self {
42        Self {
43            table: table.into(),
44            ..Default::default()
45        }
46    }
47
48    #[must_use]
49    pub fn with_row_count(mut self, n: u64) -> Self {
50        self.row_count = Some(n);
51        self
52    }
53
54    #[must_use]
55    pub fn with_ndv(mut self, n: u64) -> Self {
56        self.ndv = Some(n);
57        self
58    }
59
60    #[must_use]
61    pub fn with_avg_row_bytes(mut self, n: u64) -> Self {
62        self.avg_row_bytes = Some(n);
63        self
64    }
65}
66
67#[derive(Debug, Default, Clone)]
68pub struct TableStatsRegistry {
69    inner: Arc<RwLock<HashMap<String, TableCboStats>>>,
70}
71
72/// Process-global table statistics registry (Phase 54).
73///
74/// Written by the SQL layer (`ANALYZE TABLE`, Iceberg CTAS/DML auto-stats)
75/// and read by cost-based consumers — [`CboCostModel`] behind
76/// `default_aqe_optimizer_with_stats` on the coordinator, and any
77/// planning-time rule that wants row counts beyond the per-engine
78/// row-count registry. Clones share the same underlying map.
79pub fn global_table_stats() -> &'static TableStatsRegistry {
80    static GLOBAL: std::sync::OnceLock<TableStatsRegistry> = std::sync::OnceLock::new();
81    GLOBAL.get_or_init(TableStatsRegistry::new)
82}
83
84impl TableStatsRegistry {
85    pub fn new() -> Self {
86        Self::default()
87    }
88
89    pub fn put(&self, stats: TableCboStats) -> Option<TableCboStats> {
90        match self.inner.write() {
91            Ok(mut g) => g.insert(stats.table.clone(), stats),
92            Err(p) => p.into_inner().insert(stats.table.clone(), stats),
93        }
94    }
95
96    pub fn get(&self, table: &str) -> Option<TableCboStats> {
97        match self.inner.read() {
98            Ok(g) => g.get(table).cloned(),
99            Err(p) => p.into_inner().get(table).cloned(),
100        }
101    }
102
103    pub fn remove(&self, table: &str) -> Option<TableCboStats> {
104        match self.inner.write() {
105            Ok(mut g) => g.remove(table),
106            Err(p) => p.into_inner().remove(table),
107        }
108    }
109
110    pub fn len(&self) -> usize {
111        match self.inner.read() {
112            Ok(g) => g.len(),
113            Err(p) => p.into_inner().len(),
114        }
115    }
116
117    pub fn is_empty(&self) -> bool {
118        self.len() == 0
119    }
120
121    pub fn entries(&self) -> Vec<(String, TableCboStats)> {
122        match self.inner.read() {
123            Ok(g) => g.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),
124            Err(p) => p
125                .into_inner()
126                .iter()
127                .map(|(k, v)| (k.clone(), v.clone()))
128                .collect(),
129        }
130    }
131}
132
133/// NDV-aware cost model.
134pub struct CboCostModel {
135    pub registry: TableStatsRegistry,
136}
137
138impl std::fmt::Debug for CboCostModel {
139    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
140        f.debug_struct("CboCostModel")
141            .field("registry_entries", &self.registry.len())
142            .finish()
143    }
144}
145
146impl Default for CboCostModel {
147    fn default() -> Self {
148        Self {
149            registry: TableStatsRegistry::new(),
150        }
151    }
152}
153
154impl CostModel for CboCostModel {
155    fn estimate(&self, plan: &LogicalPlan) -> super::Cost {
156        const DEFAULT_ROWS: u64 = 10_000;
157        let mut cpu_nanos: u64 = 0;
158        let mut memory_bytes: u64 = 0;
159        let mut network_bytes: u64 = 0;
160
161        for node in plan.nodes() {
162            let (rows, row_bytes) = match node.op() {
163                Some(NodeOp::Scan { table, .. }) => match self.registry.get(table) {
164                    Some(stats) => (
165                        stats.row_count.unwrap_or(DEFAULT_ROWS),
166                        stats.avg_row_bytes.unwrap_or(64),
167                    ),
168                    None => (node.estimated_rows().unwrap_or(DEFAULT_ROWS), 64),
169                },
170                _ => (node.estimated_rows().unwrap_or(DEFAULT_ROWS), 64),
171            };
172
173            match node.op() {
174                Some(NodeOp::Scan { .. }) => {
175                    cpu_nanos = cpu_nanos.saturating_add(rows.saturating_mul(10));
176                    memory_bytes = memory_bytes.saturating_add(rows.saturating_mul(row_bytes));
177                }
178                Some(NodeOp::Filter { .. }) => {
179                    cpu_nanos = cpu_nanos.saturating_add(rows.saturating_mul(5));
180                }
181                Some(NodeOp::Project { .. }) => {
182                    cpu_nanos = cpu_nanos.saturating_add(rows.saturating_mul(2));
183                }
184                Some(NodeOp::Aggregate { group_keys, .. }) => {
185                    let ndv_proxy = group_keys.len() as u64;
186                    cpu_nanos = cpu_nanos
187                        .saturating_add(rows.saturating_mul(50))
188                        .saturating_add(ndv_proxy.saturating_mul(20));
189                    memory_bytes = memory_bytes.saturating_add(rows.saturating_mul(200));
190                }
191                Some(NodeOp::Join { .. }) | Some(NodeOp::SortMergeJoin { .. }) => {
192                    let ndv_cost = rows.saturating_mul(100);
193                    cpu_nanos = cpu_nanos.saturating_add(ndv_cost);
194                    memory_bytes = memory_bytes.saturating_add(rows.saturating_mul(100));
195                }
196                Some(NodeOp::Exchange { .. }) => {
197                    cpu_nanos = cpu_nanos.saturating_add(rows.saturating_mul(20));
198                    network_bytes = network_bytes.saturating_add(rows.saturating_mul(200));
199                }
200                _ => {
201                    cpu_nanos = cpu_nanos.saturating_add(rows.saturating_mul(15));
202                    memory_bytes = memory_bytes.saturating_add(rows.saturating_mul(64));
203                }
204            }
205        }
206
207        super::Cost {
208            cpu_nanos,
209            memory_bytes,
210            network_bytes,
211        }
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218    use crate::{ExecutionKind, PlanNode};
219
220    #[test]
221    fn empty_registry_estimate_is_zero() {
222        let model = CboCostModel::default();
223        let plan = LogicalPlan::new("p", ExecutionKind::Batch);
224        let cost = model.estimate(&plan);
225        assert_eq!(cost.cpu_nanos, 0);
226        assert_eq!(cost.memory_bytes, 0);
227        assert_eq!(cost.network_bytes, 0);
228    }
229
230    #[test]
231    fn registry_round_trip() {
232        let reg = TableStatsRegistry::new();
233        reg.put(
234            TableCboStats::new("orders")
235                .with_row_count(1_000_000)
236                .with_ndv(1_000_000)
237                .with_avg_row_bytes(128),
238        );
239        let got = reg.get("orders").expect("present");
240        assert_eq!(got.row_count, Some(1_000_000));
241        assert_eq!(got.ndv, Some(1_000_000));
242        assert_eq!(got.avg_row_bytes, Some(128));
243        assert_eq!(reg.get("unknown"), None);
244        assert_eq!(reg.len(), 1);
245    }
246
247    #[test]
248    fn registry_remove_returns_previous_value() {
249        let reg = TableStatsRegistry::new();
250        reg.put(TableCboStats::new("t1").with_row_count(100));
251        let prev = reg.remove("t1").expect("present");
252        assert_eq!(prev.row_count, Some(100));
253        assert!(reg.is_empty());
254    }
255
256    #[test]
257    fn cbo_cost_model_scan_charges_avg_row_bytes() {
258        let model = CboCostModel {
259            registry: {
260                let reg = TableStatsRegistry::new();
261                reg.put(
262                    TableCboStats::new("lineitem")
263                        .with_row_count(6_000_000)
264                        .with_avg_row_bytes(256),
265                );
266                reg
267            },
268        };
269        let mut plan = LogicalPlan::new("q", ExecutionKind::Batch);
270        plan.add_node(
271            PlanNode::new("scan", "scan lineitem", ExecutionKind::Batch).with_op(NodeOp::Scan {
272                table: "lineitem".to_string(),
273                filters: vec![],
274            }),
275        );
276        let cost = model.estimate(&plan);
277        // Scan: 6_000_000 * 10 = 60M CPU ns; 6_000_000 * 256 = 1.5 GB.
278        assert_eq!(cost.cpu_nanos, 60_000_000);
279        assert_eq!(cost.memory_bytes, 1_536_000_000);
280    }
281
282    #[test]
283    fn cbo_cost_model_aggregate_charges_group_key_proxy() {
284        let model = CboCostModel::default();
285        let mut plan = LogicalPlan::new("q", ExecutionKind::Batch);
286        plan.add_node(
287            PlanNode::new("agg", "aggregate", ExecutionKind::Batch).with_op(NodeOp::Aggregate {
288                group_keys: vec!["region".to_string(), "year".to_string()],
289            }),
290        );
291        let cost = model.estimate(&plan);
292        // rows = 10_000 default; CPU = 10_000 * 50 + 2 * 20 = 500_040
293        assert_eq!(cost.cpu_nanos, 500_040);
294    }
295}
krishiv_plan/optimizer/stats.rs

krishiv_plan/optimizer/
stats.rs
