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