1use crate::{JoinType, LogicalPlan};
20use std::collections::HashMap;
21
22#[derive(Debug, Clone, PartialEq)]
26pub struct CostEstimate {
27 pub rows: f64,
29 pub cpu_cost: f64,
31 pub io_cost: f64,
33 pub total_cost: f64,
35}
36
37impl CostEstimate {
38 pub fn zero() -> Self {
40 Self {
41 rows: 0.0,
42 cpu_cost: 0.0,
43 io_cost: 0.0,
44 total_cost: 0.0,
45 }
46 }
47
48 pub fn combine(a: &Self, b: &Self) -> Self {
52 Self {
53 rows: a.rows + b.rows,
54 cpu_cost: a.cpu_cost + b.cpu_cost,
55 io_cost: a.io_cost + b.io_cost,
56 total_cost: a.total_cost + b.total_cost,
57 }
58 }
59}
60
61#[derive(Debug, Clone)]
65pub struct TableStats {
66 pub row_count: u64,
68 pub avg_row_size_bytes: u64,
70 pub index_on: Vec<String>,
72 pub column_stats: HashMap<String, ColumnStats>,
74}
75
76impl Default for TableStats {
77 fn default() -> Self {
78 Self {
79 row_count: 10_000,
80 avg_row_size_bytes: 100,
81 index_on: vec![],
82 column_stats: HashMap::new(),
83 }
84 }
85}
86
87impl TableStats {
88 pub fn with_column(mut self, name: impl Into<String>, stats: ColumnStats) -> Self {
92 self.column_stats.insert(name.into(), stats);
93 self
94 }
95}
96
97#[derive(Debug, Clone, Default)]
101pub struct ColumnStats {
102 pub ndv: u64,
104 pub null_fraction: f64,
106 pub min: Option<f64>,
108 pub max: Option<f64>,
110}
111
112pub struct CostModel {
118 stats: HashMap<String, TableStats>,
120 cpu_weight: f64,
122 filter_selectivity: f64,
124}
125
126impl Default for CostModel {
127 fn default() -> Self {
128 Self::new()
129 }
130}
131
132impl CostModel {
133 pub fn new() -> Self {
135 Self {
136 stats: HashMap::new(),
137 cpu_weight: 0.01,
138 filter_selectivity: 0.1,
139 }
140 }
141
142 pub fn with_table_stats(mut self, table: &str, stats: TableStats) -> Self {
146 self.stats.insert(table.to_lowercase(), stats);
147 self
148 }
149
150 pub fn with_cpu_weight(mut self, w: f64) -> Self {
152 self.cpu_weight = w;
153 self
154 }
155
156 pub fn with_filter_selectivity(mut self, s: f64) -> Self {
158 self.filter_selectivity = s;
159 self
160 }
161
162 fn table_stats(&self, table: &str) -> TableStats {
164 self.stats
165 .get(&table.to_lowercase())
166 .cloned()
167 .unwrap_or_default()
168 }
169
170 fn total(&self, cpu_cost: f64, io_cost: f64) -> f64 {
172 io_cost + self.cpu_weight * cpu_cost
173 }
174
175 pub(crate) fn estimate_scan(&self, table: &str) -> CostEstimate {
181 let s = self.table_stats(table);
182 let rows = s.row_count as f64;
183 let io_cost = rows * s.avg_row_size_bytes as f64 / 8192.0;
184 let cpu_cost = rows * self.cpu_weight;
185 CostEstimate {
186 rows,
187 cpu_cost,
188 io_cost,
189 total_cost: self.total(cpu_cost, io_cost),
190 }
191 }
192
193 pub(crate) fn estimate_join_from(
199 &self,
200 l: &CostEstimate,
201 r: &CostEstimate,
202 selectivity: f64,
203 ) -> CostEstimate {
204 let rows = l.rows * r.rows * selectivity;
205 let cpu_cost = l.cpu_cost + r.cpu_cost + l.rows + r.rows + rows;
206 let io_cost = l.io_cost + r.io_cost;
207 CostEstimate {
208 rows,
209 cpu_cost,
210 io_cost,
211 total_cost: self.total(cpu_cost, io_cost),
212 }
213 }
214
215 pub fn estimate(&self, plan: &LogicalPlan) -> CostEstimate {
217 match plan {
218 LogicalPlan::Scan { table, .. } => self.estimate_scan(table),
220
221 LogicalPlan::Filter { input, .. } => {
223 let inp = self.estimate(input);
224 let rows = inp.rows * self.filter_selectivity;
225 let cpu_cost = inp.cpu_cost + inp.rows * 0.1;
226 let io_cost = inp.io_cost;
227 CostEstimate {
228 rows,
229 cpu_cost,
230 io_cost,
231 total_cost: self.total(cpu_cost, io_cost),
232 }
233 }
234
235 LogicalPlan::Project { input, .. } => {
237 let inp = self.estimate(input);
238 let rows = inp.rows;
239 let cpu_cost = inp.cpu_cost + rows * 0.01;
240 let io_cost = inp.io_cost;
241 CostEstimate {
242 rows,
243 cpu_cost,
244 io_cost,
245 total_cost: self.total(cpu_cost, io_cost),
246 }
247 }
248
249 LogicalPlan::Join {
251 left,
252 right,
253 join_type,
254 ..
255 } => {
256 let l = self.estimate(left);
257 let r = self.estimate(right);
258 match join_type {
259 JoinType::LeftSemi => {
260 let rows = l.rows.min(l.rows * self.filter_selectivity);
262 let cpu_cost = l.cpu_cost + r.cpu_cost + l.rows + r.rows + rows;
263 let io_cost = l.io_cost + r.io_cost;
264 CostEstimate {
265 rows,
266 cpu_cost,
267 io_cost,
268 total_cost: self.total(cpu_cost, io_cost),
269 }
270 }
271 JoinType::LeftAnti => {
272 let rows = (l.rows * (1.0 - self.filter_selectivity)).min(l.rows);
274 let cpu_cost = l.cpu_cost + r.cpu_cost + l.rows + r.rows + rows;
275 let io_cost = l.io_cost + r.io_cost;
276 CostEstimate {
277 rows,
278 cpu_cost,
279 io_cost,
280 total_cost: self.total(cpu_cost, io_cost),
281 }
282 }
283 _ => self.estimate_join_from(&l, &r, self.filter_selectivity),
284 }
285 }
286
287 LogicalPlan::Aggregate { input, .. } => {
289 let inp = self.estimate(input);
290 let rows = inp.rows * 0.1;
291 let cpu_cost = inp.cpu_cost + inp.rows * 0.5;
292 let io_cost = inp.io_cost;
293 CostEstimate {
294 rows,
295 cpu_cost,
296 io_cost,
297 total_cost: self.total(cpu_cost, io_cost),
298 }
299 }
300
301 LogicalPlan::Sort { input, .. } => {
303 let inp = self.estimate(input);
304 let rows = inp.rows;
305 let sort_cpu = rows * rows.ln().max(1.0) * 0.01;
306 let cpu_cost = inp.cpu_cost + sort_cpu;
307 let io_cost = inp.io_cost;
308 CostEstimate {
309 rows,
310 cpu_cost,
311 io_cost,
312 total_cost: self.total(cpu_cost, io_cost),
313 }
314 }
315
316 LogicalPlan::Limit { input, count, .. } => {
318 let inp = self.estimate(input);
319 let rows = match count {
320 Some(c) => inp.rows.min(*c as f64),
321 None => inp.rows,
322 };
323 let cpu_cost = inp.cpu_cost;
324 let io_cost = inp.io_cost;
325 CostEstimate {
326 rows,
327 cpu_cost,
328 io_cost,
329 total_cost: self.total(cpu_cost, io_cost),
330 }
331 }
332
333 LogicalPlan::Values { .. } | LogicalPlan::Empty => CostEstimate::zero(),
335
336 LogicalPlan::SetOp { left, right, .. } => {
338 let l = self.estimate(left);
339 let r = self.estimate(right);
340 CostEstimate::combine(&l, &r)
341 }
342
343 LogicalPlan::Cte { query, .. } => {
345 let inner = self.estimate(query);
346 let cpu_cost = inner.cpu_cost + inner.rows * 0.05;
347 let io_cost = inner.io_cost;
348 let rows = inner.rows;
349 CostEstimate {
350 rows,
351 cpu_cost,
352 io_cost,
353 total_cost: self.total(cpu_cost, io_cost),
354 }
355 }
356
357 LogicalPlan::CteRef { name } => {
359 let s = self.table_stats(name);
360 let rows = s.row_count as f64;
361 let io_cost = rows * s.avg_row_size_bytes as f64 / 8192.0;
362 let cpu_cost = rows * self.cpu_weight + rows * 0.05;
363 CostEstimate {
364 rows,
365 cpu_cost,
366 io_cost,
367 total_cost: self.total(cpu_cost, io_cost),
368 }
369 }
370
371 LogicalPlan::Window { input, .. } => {
373 let inp = self.estimate(input);
374 let rows = inp.rows;
375 let cpu_cost = inp.cpu_cost + rows * 0.1;
376 let io_cost = inp.io_cost;
377 CostEstimate {
378 rows,
379 cpu_cost,
380 io_cost,
381 total_cost: self.total(cpu_cost, io_cost),
382 }
383 }
384
385 LogicalPlan::Subquery { query, .. } => {
387 let inner = self.estimate(query);
388 let rows = inner.rows;
389 let cpu_cost = inner.cpu_cost + rows * 0.2;
390 let io_cost = inner.io_cost;
391 CostEstimate {
392 rows,
393 cpu_cost,
394 io_cost,
395 total_cost: self.total(cpu_cost, io_cost),
396 }
397 }
398
399 LogicalPlan::Exists { subquery, .. } => {
401 let inner = self.estimate(subquery);
402 let rows = 1.0_f64.min(inner.rows);
404 let cpu_cost = inner.cpu_cost + 0.1;
405 let io_cost = inner.io_cost;
406 CostEstimate {
407 rows,
408 cpu_cost,
409 io_cost,
410 total_cost: self.total(cpu_cost, io_cost),
411 }
412 }
413
414 LogicalPlan::InSubquery { subquery, .. } => {
416 let inner = self.estimate(subquery);
417 let rows = inner.rows * self.filter_selectivity;
418 let cpu_cost = inner.cpu_cost + inner.rows * 0.15;
419 let io_cost = inner.io_cost;
420 CostEstimate {
421 rows,
422 cpu_cost,
423 io_cost,
424 total_cost: self.total(cpu_cost, io_cost),
425 }
426 }
427
428 LogicalPlan::Compute { input, .. } => self.estimate(input),
433 }
434 }
435}
436
437#[derive(Debug, Clone)]
445pub struct NodeCost {
446 pub op: &'static str,
448 pub estimate: CostEstimate,
450 pub children: Vec<NodeCost>,
452}
453
454impl CostModel {
455 pub fn explain_costs(&self, plan: &LogicalPlan) -> NodeCost {
457 match plan {
458 LogicalPlan::Scan { .. } => NodeCost {
459 op: "Scan",
460 estimate: self.estimate(plan),
461 children: vec![],
462 },
463 LogicalPlan::Filter { input, .. } => NodeCost {
464 op: "Filter",
465 estimate: self.estimate(plan),
466 children: vec![self.explain_costs(input)],
467 },
468 LogicalPlan::Project { input, .. } => NodeCost {
469 op: "Project",
470 estimate: self.estimate(plan),
471 children: vec![self.explain_costs(input)],
472 },
473 LogicalPlan::Join { left, right, .. } => NodeCost {
474 op: "Join",
475 estimate: self.estimate(plan),
476 children: vec![self.explain_costs(left), self.explain_costs(right)],
477 },
478 LogicalPlan::Aggregate { input, .. } => NodeCost {
479 op: "Aggregate",
480 estimate: self.estimate(plan),
481 children: vec![self.explain_costs(input)],
482 },
483 LogicalPlan::Sort { input, .. } => NodeCost {
484 op: "Sort",
485 estimate: self.estimate(plan),
486 children: vec![self.explain_costs(input)],
487 },
488 LogicalPlan::Limit { input, .. } => NodeCost {
489 op: "Limit",
490 estimate: self.estimate(plan),
491 children: vec![self.explain_costs(input)],
492 },
493 LogicalPlan::Values { .. } => NodeCost {
494 op: "Values",
495 estimate: CostEstimate::zero(),
496 children: vec![],
497 },
498 LogicalPlan::Empty => NodeCost {
499 op: "Empty",
500 estimate: CostEstimate::zero(),
501 children: vec![],
502 },
503 LogicalPlan::SetOp { left, right, .. } => NodeCost {
504 op: "SetOp",
505 estimate: self.estimate(plan),
506 children: vec![self.explain_costs(left), self.explain_costs(right)],
507 },
508 LogicalPlan::Cte { query, .. } => NodeCost {
509 op: "Cte",
510 estimate: self.estimate(plan),
511 children: vec![self.explain_costs(query)],
512 },
513 LogicalPlan::CteRef { .. } => NodeCost {
514 op: "CteRef",
515 estimate: self.estimate(plan),
516 children: vec![],
517 },
518 LogicalPlan::Window { input, .. } => NodeCost {
519 op: "Window",
520 estimate: self.estimate(plan),
521 children: vec![self.explain_costs(input)],
522 },
523 LogicalPlan::Subquery { query, .. } => NodeCost {
524 op: "Subquery",
525 estimate: self.estimate(plan),
526 children: vec![self.explain_costs(query)],
527 },
528 LogicalPlan::Exists { subquery, .. } => NodeCost {
529 op: "Exists",
530 estimate: self.estimate(plan),
531 children: vec![self.explain_costs(subquery)],
532 },
533 LogicalPlan::InSubquery { subquery, .. } => NodeCost {
534 op: "InSubquery",
535 estimate: self.estimate(plan),
536 children: vec![self.explain_costs(subquery)],
537 },
538 LogicalPlan::Compute { input, .. } => NodeCost {
539 op: "Compute",
540 estimate: self.estimate(plan),
541 children: vec![self.explain_costs(input)],
542 },
543 }
544 }
545}
546
547#[cfg(test)]
550mod tests {
551 use super::*;
552 use crate::{parse_one, plan_statement};
553
554 #[test]
556 fn test_cost_scan_default() {
557 let stmt = parse_one("SELECT * FROM users").expect("parse");
558 let plan = plan_statement(&stmt).expect("plan");
559 let model = CostModel::new();
560 let est = model.estimate(&plan);
561 assert!(
563 (est.rows - 10_000.0).abs() < 1e-9,
564 "expected 10_000 rows from default stats, got {}",
565 est.rows
566 );
567 }
568
569 #[test]
571 fn test_cost_scan_custom() {
572 let stmt = parse_one("SELECT * FROM orders").expect("parse");
573 let plan = plan_statement(&stmt).expect("plan");
574 let model = CostModel::new().with_table_stats(
575 "orders",
576 TableStats {
577 row_count: 50_000,
578 avg_row_size_bytes: 128,
579 index_on: vec![],
580 ..Default::default()
581 },
582 );
583 let est = model.estimate(&plan);
584 assert!(
585 (est.rows - 50_000.0).abs() < 1e-9,
586 "expected 50_000 rows, got {}",
587 est.rows
588 );
589 }
590
591 #[test]
593 fn test_cost_filter_reduces_rows() {
594 let stmt = parse_one("SELECT * FROM users WHERE active = true").expect("parse");
595 let plan = plan_statement(&stmt).expect("plan");
596 let model = CostModel::new();
597 let est = model.estimate(&plan);
598 assert!(
600 est.rows < 10_000.0,
601 "filter should reduce rows below scan rows, got {}",
602 est.rows
603 );
604 }
605
606 #[test]
608 fn test_cost_join_two_scans() {
609 let stmt = parse_one("SELECT a.id, b.name FROM a JOIN b ON a.id = b.a_id").expect("parse");
610 let plan = plan_statement(&stmt).expect("plan");
611 let model = CostModel::new();
612 let est = model.estimate(&plan);
613 let expected = 10_000.0_f64 * 10_000.0 * 0.1;
614 assert!(
615 (est.rows - expected).abs() < 1.0,
616 "join rows expected ~{expected}, got {}",
617 est.rows
618 );
619 }
620
621 #[test]
623 fn test_cost_limit_caps_rows() {
624 let stmt = parse_one("SELECT * FROM users LIMIT 100").expect("parse");
625 let plan = plan_statement(&stmt).expect("plan");
626 let model = CostModel::new();
627 let est = model.estimate(&plan);
628 assert!(
629 (est.rows - 100.0).abs() < 1e-9,
630 "limit should cap rows to 100, got {}",
631 est.rows
632 );
633 }
634
635 #[test]
637 fn test_cost_aggregate_reduces_rows() {
638 let stmt = parse_one("SELECT dept, COUNT(*) FROM employees GROUP BY dept").expect("parse");
639 let plan = plan_statement(&stmt).expect("plan");
640 let model = CostModel::new();
641 let est = model.estimate(&plan);
642 assert!(
643 est.rows < 10_000.0,
644 "aggregate should reduce rows, got {}",
645 est.rows
646 );
647 }
648
649 #[test]
651 fn test_cost_sort_preserves_rows() {
652 let stmt = parse_one("SELECT * FROM products ORDER BY price DESC").expect("parse");
653 let plan = plan_statement(&stmt).expect("plan");
654 let model = CostModel::new();
655 let full_est = model.estimate(&plan);
656
657 let stmt_no_sort = parse_one("SELECT * FROM products").expect("parse_no_sort");
659 let plan_no_sort = plan_statement(&stmt_no_sort).expect("plan_no_sort");
660 let base_est = model.estimate(&plan_no_sort);
661
662 assert!(
664 (full_est.rows - base_est.rows).abs() < 1e-9,
665 "sort should preserve row count: {} vs {}",
666 full_est.rows,
667 base_est.rows
668 );
669 assert!(
671 full_est.cpu_cost > base_est.cpu_cost,
672 "sort should increase cpu_cost: {} vs {}",
673 full_est.cpu_cost,
674 base_est.cpu_cost
675 );
676 }
677
678 #[test]
680 fn test_cost_combines_correctly() {
681 let a = CostEstimate {
682 rows: 100.0,
683 cpu_cost: 5.0,
684 io_cost: 10.0,
685 total_cost: 10.05,
686 };
687 let b = CostEstimate {
688 rows: 200.0,
689 cpu_cost: 3.0,
690 io_cost: 7.0,
691 total_cost: 7.03,
692 };
693 let c = CostEstimate::combine(&a, &b);
694 assert!((c.rows - 300.0).abs() < 1e-9, "rows: {}", c.rows);
695 assert!((c.cpu_cost - 8.0).abs() < 1e-9, "cpu_cost: {}", c.cpu_cost);
696 assert!((c.io_cost - 17.0).abs() < 1e-9, "io_cost: {}", c.io_cost);
697 assert!(
698 (c.total_cost - 17.08).abs() < 1e-9,
699 "total_cost: {}",
700 c.total_cost
701 );
702 }
703}