icydb_core/db/query/
planner.rs

1use crate::{
2    IndexSpec, Key, Value,
3    db::primitives::filter::{Cmp, FilterExpr},
4    obs::metrics,
5    traits::EntityKind,
6};
7use std::fmt::{self, Display};
8
9///
10/// QueryPlan
11///
12
13#[derive(Debug)]
14pub enum QueryPlan {
15    FullScan,
16    Index(IndexPlan),
17    Keys(Vec<Key>),
18    Range(Key, Key),
19}
20
21impl fmt::Display for QueryPlan {
22    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
23        match self {
24            Self::Index(plan) => write!(f, "Index({plan})"),
25
26            Self::Keys(keys) => {
27                // Show up to 5 keys, then ellipsize
28                let preview: Vec<String> = keys.iter().take(5).map(|k| format!("{k:?}")).collect();
29
30                if keys.len() > 5 {
31                    write!(f, "Keys[{}… total {}]", preview.join(", "), keys.len())
32                } else {
33                    write!(f, "Keys[{}]", preview.join(", "))
34                }
35            }
36
37            Self::Range(start, end) => {
38                write!(f, "Range({start:?} → {end:?})")
39            }
40
41            Self::FullScan => write!(f, "FullScan"),
42        }
43    }
44}
45
46///
47/// IndexPlan
48///
49
50#[derive(Debug)]
51pub struct IndexPlan {
52    pub index: &'static IndexSpec,
53    pub values: Vec<Value>,
54}
55
56impl Display for IndexPlan {
57    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
58        let values: Vec<String> = self.values.iter().map(|v| format!("{v:?}")).collect();
59        write!(f, "index={} values=[{}]", self.index, values.join(", "))
60    }
61}
62
63///
64/// QueryPlanner
65///
66
67#[derive(Debug)]
68pub struct QueryPlanner {
69    pub filter: Option<FilterExpr>,
70}
71
72impl QueryPlanner {
73    #[must_use]
74    pub fn new(filter: Option<&FilterExpr>) -> Self {
75        Self {
76            filter: filter.cloned(),
77        }
78    }
79
80    #[must_use]
81    pub fn plan<E: EntityKind>(&self) -> QueryPlan {
82        // If filter is a primary key match
83        // this would handle One and Many queries
84        if let Some(plan) = self.extract_from_filter::<E>() {
85            metrics::with_state_mut(|m| match plan {
86                QueryPlan::Keys(_) => m.ops.plan_keys += 1,
87                QueryPlan::Index(_) => m.ops.plan_index += 1,
88                QueryPlan::Range(_, _) | QueryPlan::FullScan => m.ops.plan_range += 1,
89            });
90            return plan;
91        }
92
93        // check for index matches
94        // THIS WILL DO THE INDEX LOOKUPS
95        if !E::INDEXES.is_empty()
96            && let Some(plan) = self.extract_from_index::<E>()
97        {
98            metrics::with_state_mut(|m| m.ops.plan_index += 1);
99            return plan;
100        }
101
102        // Fallback: do a full scan
103        metrics::with_state_mut(|m| m.ops.plan_range += 1);
104
105        QueryPlan::FullScan
106    }
107
108    // extract_from_filter
109    fn extract_from_filter<E: EntityKind>(&self) -> Option<QueryPlan> {
110        let Some(filter) = &self.filter else {
111            return None;
112        };
113
114        match filter {
115            FilterExpr::Clause(clause) if clause.field == E::PRIMARY_KEY => match clause.cmp {
116                Cmp::Eq => clause.value.as_key().map(|key| QueryPlan::Keys(vec![key])),
117
118                Cmp::In => {
119                    if let Value::List(values) = &clause.value {
120                        let keys = values.iter().filter_map(Value::as_key).collect::<Vec<_>>();
121
122                        if keys.is_empty() {
123                            None
124                        } else {
125                            Some(QueryPlan::Keys(keys))
126                        }
127                    } else {
128                        None
129                    }
130                }
131
132                _ => None,
133            },
134
135            _ => None,
136        }
137    }
138
139    // extract_from_index: build a leftmost equality prefix in terms of Value
140    fn extract_from_index<E: EntityKind>(&self) -> Option<QueryPlan> {
141        let Some(filter) = &self.filter else {
142            return None;
143        };
144
145        let mut best: Option<(usize, IndexPlan)> = None;
146
147        for index in E::INDEXES {
148            // Build leftmost equality prefix (only == supported for hashed indexes)
149            let mut values: Vec<Value> = Vec::with_capacity(index.fields.len());
150
151            for field in index.fields {
152                if let Some(v) = Self::find_eq_value(filter, field) {
153                    values.push(v);
154                } else {
155                    break; // stop at first non-match
156                }
157            }
158
159            // Skip indexes that produced no equality prefix
160            if values.is_empty() {
161                continue;
162            }
163
164            let score = values.len();
165            let cand = (score, IndexPlan { index, values });
166
167            match &best {
168                Some((best_score, _)) if *best_score >= score => { /* keep current best */ }
169                _ => best = Some(cand),
170            }
171        }
172
173        best.map(|(_, plan)| QueryPlan::Index(plan))
174    }
175
176    /// Find an equality clause (`field == ?`) anywhere in the filter tree and return the Value.
177    fn find_eq_value(filter: &FilterExpr, field: &str) -> Option<Value> {
178        match filter {
179            FilterExpr::Clause(c) if c.field == field && matches!(c.cmp, Cmp::Eq) => {
180                Some(c.value.clone())
181            }
182            // Walk conjunctive subtrees
183            FilterExpr::And(list) => list.iter().find_map(|f| Self::find_eq_value(f, field)),
184            _ => None,
185        }
186    }
187}