chryso_optimizer/
stats_collect.rs

1use chryso_core::ast::{Expr, OrderByExpr};
2use chryso_planner::LogicalPlan;
3use std::collections::HashSet;
4
5#[derive(Default)]
6pub struct StatsRequirements {
7    pub tables: HashSet<String>,
8    pub columns: HashSet<(String, String)>,
9}
10
11pub fn collect_requirements(plan: &LogicalPlan) -> StatsRequirements {
12    let mut requirements = StatsRequirements::default();
13    collect_tables(plan, &mut requirements.tables);
14    collect_columns(plan, &requirements.tables, &mut requirements.columns);
15    requirements
16}
17
18fn collect_tables(plan: &LogicalPlan, tables: &mut HashSet<String>) {
19    match plan {
20        LogicalPlan::Scan { table } | LogicalPlan::IndexScan { table, .. } => {
21            tables.insert(table.clone());
22        }
23        LogicalPlan::Dml { .. } => {}
24        LogicalPlan::Derived { input, .. } => collect_tables(input, tables),
25        LogicalPlan::Filter { input, .. }
26        | LogicalPlan::Projection { input, .. }
27        | LogicalPlan::Aggregate { input, .. }
28        | LogicalPlan::Distinct { input }
29        | LogicalPlan::TopN { input, .. }
30        | LogicalPlan::Sort { input, .. }
31        | LogicalPlan::Limit { input, .. } => collect_tables(input, tables),
32        LogicalPlan::Join { left, right, .. } => {
33            collect_tables(left, tables);
34            collect_tables(right, tables);
35        }
36    }
37}
38
39fn collect_columns(plan: &LogicalPlan, tables: &HashSet<String>, columns: &mut HashSet<(String, String)>) {
40    match plan {
41        LogicalPlan::Scan { .. } | LogicalPlan::IndexScan { .. } | LogicalPlan::Dml { .. } => {}
42        LogicalPlan::Derived { input, .. } => collect_columns(input, tables, columns),
43        LogicalPlan::Filter { predicate, input } => {
44            collect_expr_columns(predicate, tables, columns);
45            collect_columns(input, tables, columns);
46        }
47        LogicalPlan::Projection { exprs, input } => {
48            collect_exprs_columns(exprs, tables, columns);
49            collect_columns(input, tables, columns);
50        }
51        LogicalPlan::Join { on, left, right, .. } => {
52            collect_expr_columns(on, tables, columns);
53            collect_columns(left, tables, columns);
54            collect_columns(right, tables, columns);
55        }
56        LogicalPlan::Aggregate {
57            group_exprs,
58            aggr_exprs,
59            input,
60        } => {
61            collect_exprs_columns(group_exprs, tables, columns);
62            collect_exprs_columns(aggr_exprs, tables, columns);
63            collect_columns(input, tables, columns);
64        }
65        LogicalPlan::Distinct { input } => collect_columns(input, tables, columns),
66        LogicalPlan::TopN { order_by, input, .. } => {
67            collect_order_by_columns(order_by, tables, columns);
68            collect_columns(input, tables, columns);
69        }
70        LogicalPlan::Sort { order_by, input } => {
71            collect_order_by_columns(order_by, tables, columns);
72            collect_columns(input, tables, columns);
73        }
74        LogicalPlan::Limit { input, .. } => collect_columns(input, tables, columns),
75    }
76}
77
78fn collect_exprs_columns(exprs: &[Expr], tables: &HashSet<String>, columns: &mut HashSet<(String, String)>) {
79    for expr in exprs {
80        collect_expr_columns(expr, tables, columns);
81    }
82}
83
84fn collect_order_by_columns(order_by: &[OrderByExpr], tables: &HashSet<String>, columns: &mut HashSet<(String, String)>) {
85    for item in order_by {
86        collect_expr_columns(&item.expr, tables, columns);
87    }
88}
89
90fn collect_expr_columns(expr: &Expr, tables: &HashSet<String>, columns: &mut HashSet<(String, String)>) {
91    match expr {
92        Expr::Identifier(name) => {
93            if let Some((table, column)) = name.split_once('.') {
94                if tables.contains(table) {
95                    columns.insert((table.to_string(), column.to_string()));
96                }
97            } else if tables.len() == 1 {
98                if let Some(table) = tables.iter().next() {
99                    columns.insert((table.clone(), name.to_string()));
100                }
101            }
102        }
103        Expr::BinaryOp { left, right, .. } => {
104            collect_expr_columns(left, tables, columns);
105            collect_expr_columns(right, tables, columns);
106        }
107        Expr::IsNull { expr, .. } => {
108            collect_expr_columns(expr, tables, columns);
109        }
110        Expr::UnaryOp { expr, .. } => collect_expr_columns(expr, tables, columns),
111        Expr::FunctionCall { args, .. } => {
112            for arg in args {
113                collect_expr_columns(arg, tables, columns);
114            }
115        }
116        Expr::WindowFunction { function, spec } => {
117            collect_expr_columns(function, tables, columns);
118            collect_exprs_columns(&spec.partition_by, tables, columns);
119            collect_order_by_columns(&spec.order_by, tables, columns);
120        }
121        Expr::Subquery(select) | Expr::Exists(select) => {
122            for item in &select.projection {
123                collect_expr_columns(&item.expr, tables, columns);
124            }
125        }
126        Expr::InSubquery { expr, subquery } => {
127            collect_expr_columns(expr, tables, columns);
128            for item in &subquery.projection {
129                collect_expr_columns(&item.expr, tables, columns);
130            }
131        }
132        Expr::Case {
133            operand,
134            when_then,
135            else_expr,
136        } => {
137            if let Some(expr) = operand {
138                collect_expr_columns(expr, tables, columns);
139            }
140            for (when_expr, then_expr) in when_then {
141                collect_expr_columns(when_expr, tables, columns);
142                collect_expr_columns(then_expr, tables, columns);
143            }
144            if let Some(expr) = else_expr {
145                collect_expr_columns(expr, tables, columns);
146            }
147        }
148        Expr::Literal(_) | Expr::Wildcard => {}
149    }
150}
chryso_optimizer/stats_collect.rs

chryso_optimizer/
stats_collect.rs
