basalt/
planner.rs

1//! Lightweight cost-based access-path selection.
2//!
3//! The planner deliberately keeps the physical plan small: a table scan is
4//! always available, and an equality predicate on an indexed column can use
5//! an index when its candidate set is cheaper than visiting every row. More
6//! join and aggregate planning remains in the executor's relational pipeline.
7
8use crate::db::Table;
9use crate::sql::ast::{BinOp, Expr};
10use crate::types::Value;
11
12type RangeCandidate = (usize, Option<Value>, Option<Value>, Vec<u64>, String);
13
14#[derive(Debug, Clone)]
15pub enum AccessPath {
16    TableScan,
17    IndexScan {
18        index_name: String,
19        column: usize,
20        key: Value,
21        row_ids: Vec<u64>,
22    },
23    IndexRange {
24        index_name: String,
25        column: usize,
26        low: Option<Value>,
27        high: Option<Value>,
28        row_ids: Vec<u64>,
29    },
30}
31
32#[derive(Debug, Clone)]
33pub struct Plan {
34    pub table: String,
35    pub access: AccessPath,
36    pub estimated_rows: usize,
37}
38
39/// Choose an access path for a single table's predicate. The returned path is
40/// guaranteed to be semantically equivalent to a full scan; it only narrows
41/// the rows that still go through normal predicate evaluation.
42pub fn choose(table: &Table, relation: Option<&str>, predicate: Option<&Expr>) -> Plan {
43    let full_scan_cost = table.row_count().max(1);
44    let mut best: Option<(usize, AccessPath)> = None;
45    let mut candidate = None;
46    if let Some(predicate) = predicate {
47        find_equality(table, relation, predicate, &mut candidate);
48    }
49    if let Some((column, key, row_ids, index_name)) = candidate {
50        // An index traversal plus fetching candidates is worthwhile only when
51        // the candidate set is smaller than a full table walk. Primary keys
52        // and UNIQUE indexes naturally satisfy this for all but tiny tables.
53        if row_ids.len().saturating_mul(2) <= full_scan_cost {
54            best = Some((
55                row_ids.len(),
56                AccessPath::IndexScan {
57                    index_name,
58                    column,
59                    key,
60                    row_ids,
61                },
62            ));
63        }
64    }
65    let mut range = None;
66    if let Some(predicate) = predicate {
67        find_range(table, relation, predicate, &mut range);
68    }
69    if let Some((column, low, high, row_ids, index_name)) = range
70        && row_ids.len().saturating_mul(2) <= full_scan_cost
71        && best
72            .as_ref()
73            .map(|(estimated, _)| row_ids.len() < *estimated)
74            .unwrap_or(true)
75    {
76        best = Some((
77            row_ids.len(),
78            AccessPath::IndexRange {
79                index_name,
80                column,
81                low,
82                high,
83                row_ids,
84            },
85        ));
86    }
87    if let Some((estimated_rows, access)) = best {
88        return Plan {
89            table: table.name.clone(),
90            estimated_rows,
91            access,
92        };
93    }
94    Plan {
95        table: table.name.clone(),
96        estimated_rows: table.row_count(),
97        access: AccessPath::TableScan,
98    }
99}
100
101fn find_equality(
102    table: &Table,
103    relation: Option<&str>,
104    expr: &Expr,
105    candidate: &mut Option<(usize, Value, Vec<u64>, String)>,
106) {
107    match expr {
108        Expr::Binary {
109            left,
110            op: BinOp::And,
111            right,
112        } => {
113            find_equality(table, relation, left, candidate);
114            find_equality(table, relation, right, candidate);
115        }
116        Expr::Binary {
117            left,
118            op: BinOp::Eq,
119            right,
120        } => {
121            if let Some((column_name, value)) = column_literal(left, right, relation)
122                && let Ok(column) = table.column_index(column_name)
123                && let Some((index_name, row_ids)) = table.lookup_eq_index(column, value)
124            {
125                let is_better = candidate
126                    .as_ref()
127                    .map(|existing| row_ids.len() < existing.2.len())
128                    .unwrap_or(true);
129                if is_better {
130                    *candidate = Some((column, value.clone(), row_ids, index_name));
131                }
132            }
133        }
134        _ => {}
135    }
136}
137
138fn column_literal<'a>(
139    left: &'a Expr,
140    right: &'a Expr,
141    relation: Option<&str>,
142) -> Option<(&'a str, &'a Value)> {
143    match (left, right) {
144        (Expr::Column(column), Expr::Literal(value)) => Some((column, value)),
145        (
146            Expr::ColumnRef {
147                relation: found,
148                column,
149            },
150            Expr::Literal(value),
151        ) if relation
152            .map(|wanted| found.eq_ignore_ascii_case(wanted))
153            .unwrap_or(true) =>
154        {
155            Some((column, value))
156        }
157        (Expr::Literal(value), Expr::Column(column)) => Some((column, value)),
158        (
159            Expr::Literal(value),
160            Expr::ColumnRef {
161                relation: found,
162                column,
163            },
164        ) if relation
165            .map(|wanted| found.eq_ignore_ascii_case(wanted))
166            .unwrap_or(true) =>
167        {
168            Some((column, value))
169        }
170        _ => None,
171    }
172}
173
174fn find_range(
175    table: &Table,
176    relation: Option<&str>,
177    expr: &Expr,
178    candidate: &mut Option<RangeCandidate>,
179) {
180    match expr {
181        Expr::Binary {
182            left,
183            op: BinOp::And,
184            right,
185        } => {
186            find_range(table, relation, left, candidate);
187            find_range(table, relation, right, candidate);
188        }
189        Expr::Binary { left, op, right }
190            if matches!(op, BinOp::Lt | BinOp::LtEq | BinOp::Gt | BinOp::GtEq) =>
191        {
192            let (column, value, reversed) =
193                if let Some((column, value)) = column_literal(left, right, relation) {
194                    (column, value, false)
195                } else if let Some((column, value)) = column_literal(right, left, relation) {
196                    (column, value, true)
197                } else {
198                    return;
199                };
200            let (low, high) = match (op, reversed) {
201                (BinOp::Lt | BinOp::LtEq, false) => (None, Some(value.clone())),
202                (BinOp::Gt | BinOp::GtEq, false) => (Some(value.clone()), None),
203                (BinOp::Lt | BinOp::LtEq, true) => (Some(value.clone()), None),
204                (BinOp::Gt | BinOp::GtEq, true) => (None, Some(value.clone())),
205                _ => unreachable!(),
206            };
207            if let Ok(column) = table.column_index(column)
208                && let Some((name, rows)) =
209                    table.lookup_range_index(column, low.as_ref(), high.as_ref())
210            {
211                let is_better = candidate
212                    .as_ref()
213                    .map(|existing| rows.len() < existing.3.len())
214                    .unwrap_or(true);
215                if is_better {
216                    *candidate = Some((column, low, high, rows, name));
217                }
218            }
219        }
220        _ => {}
221    }
222}
223
224#[cfg(test)]
225mod tests {
226    use super::*;
227    use crate::db::Column;
228    use crate::sql::ast::Statement;
229    use crate::sql::parser::parse;
230    use crate::types::ColumnType;
231
232    #[test]
233    fn picks_selective_index() {
234        let mut table = Table::new(
235            "t",
236            vec![Column {
237                name: "id".into(),
238                ty: ColumnType::Integer,
239                not_null: true,
240                unique: false,
241                primary_key: false,
242            }],
243        )
244        .unwrap();
245        for value in 0..100 {
246            table.insert_row(vec![Value::Integer(value)]).unwrap();
247        }
248        table.create_index("id_idx", 0, false).unwrap();
249        let Statement::Select { where_clause, .. } =
250            &parse("SELECT * FROM t WHERE id = 42").unwrap()[0]
251        else {
252            panic!()
253        };
254        let plan = choose(&table, None, where_clause.as_ref());
255        assert!(matches!(plan.access, AccessPath::IndexScan { .. }));
256    }
257
258    #[test]
259    fn picks_range_index() {
260        let mut table = Table::new(
261            "t",
262            vec![Column {
263                name: "id".into(),
264                ty: ColumnType::Integer,
265                not_null: true,
266                unique: false,
267                primary_key: false,
268            }],
269        )
270        .unwrap();
271        for value in 0..100 {
272            table.insert_row(vec![Value::Integer(value)]).unwrap();
273        }
274        table.create_index("id_idx", 0, false).unwrap();
275        let Statement::Select { where_clause, .. } =
276            &parse("SELECT * FROM t WHERE id >= 90").unwrap()[0]
277        else {
278            panic!()
279        };
280        assert!(matches!(
281            choose(&table, None, where_clause.as_ref()).access,
282            AccessPath::IndexRange { .. }
283        ));
284    }
285
286    #[test]
287    fn picks_the_most_selective_predicate() {
288        let mut table = Table::new(
289            "t",
290            vec![
291                Column {
292                    name: "bucket".into(),
293                    ty: ColumnType::Integer,
294                    not_null: true,
295                    unique: false,
296                    primary_key: false,
297                },
298                Column {
299                    name: "id".into(),
300                    ty: ColumnType::Integer,
301                    not_null: true,
302                    unique: false,
303                    primary_key: false,
304                },
305            ],
306        )
307        .unwrap();
308        for id in 0..100 {
309            table
310                .insert_row(vec![Value::Integer(id % 10), Value::Integer(id)])
311                .unwrap();
312        }
313        table.create_index("bucket_idx", 0, false).unwrap();
314        table.create_index("id_idx", 1, false).unwrap();
315        let Statement::Select { where_clause, .. } =
316            &parse("SELECT * FROM t WHERE bucket = 1 AND id = 42").unwrap()[0]
317        else {
318            panic!()
319        };
320        let plan = choose(&table, None, where_clause.as_ref());
321        assert!(matches!(
322            plan.access,
323            AccessPath::IndexScan { column: 1, .. }
324        ));
325    }
326}
basalt/planner.rs

basalt/
planner.rs
