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    /// Create a planner from an optional filter expression.
75    pub fn new(filter: Option<&FilterExpr>) -> Self {
76        Self {
77            filter: filter.cloned(),
78        }
79    }
80
81    #[must_use]
82    /// Generate a query plan for the given entity type.
83    pub fn plan<E: EntityKind>(&self) -> QueryPlan {
84        // If filter is a primary key match
85        // this would handle One and Many queries
86        if let Some(plan) = self.extract_from_filter::<E>() {
87            metrics::with_state_mut(|m| match plan {
88                QueryPlan::Keys(_) => m.ops.plan_keys += 1,
89                QueryPlan::Index(_) => m.ops.plan_index += 1,
90                QueryPlan::Range(_, _) | QueryPlan::FullScan => m.ops.plan_range += 1,
91            });
92            return plan;
93        }
94
95        // check for index matches
96        // THIS WILL DO THE INDEX LOOKUPS
97        if !E::INDEXES.is_empty()
98            && let Some(plan) = self.extract_from_index::<E>()
99        {
100            metrics::with_state_mut(|m| m.ops.plan_index += 1);
101            return plan;
102        }
103
104        // Fallback: do a full scan
105        metrics::with_state_mut(|m| m.ops.plan_range += 1);
106
107        QueryPlan::FullScan
108    }
109
110    // extract_from_filter
111    fn extract_from_filter<E: EntityKind>(&self) -> Option<QueryPlan> {
112        let Some(filter) = &self.filter else {
113            return None;
114        };
115
116        match filter {
117            FilterExpr::Clause(clause) if clause.field == E::PRIMARY_KEY => match clause.cmp {
118                Cmp::Eq => clause.value.as_key().map(|key| QueryPlan::Keys(vec![key])),
119
120                Cmp::In => {
121                    if let Value::List(values) = &clause.value {
122                        let keys = values.iter().filter_map(Value::as_key).collect::<Vec<_>>();
123
124                        if keys.is_empty() {
125                            None
126                        } else {
127                            Some(QueryPlan::Keys(keys))
128                        }
129                    } else {
130                        None
131                    }
132                }
133
134                _ => None,
135            },
136
137            _ => None,
138        }
139    }
140
141    // extract_from_index: build a leftmost equality prefix in terms of Value
142    fn extract_from_index<E: EntityKind>(&self) -> Option<QueryPlan> {
143        let Some(filter) = &self.filter else {
144            return None;
145        };
146
147        let mut best: Option<(usize, IndexPlan)> = None;
148
149        for index in E::INDEXES {
150            // Build leftmost equality prefix (only == supported for hashed indexes)
151            let mut values: Vec<Value> = Vec::with_capacity(index.fields.len());
152
153            for field in index.fields {
154                if let Some(v) = Self::find_eq_value(filter, field) {
155                    values.push(v);
156                } else {
157                    break; // stop at first non-match
158                }
159            }
160
161            // Skip indexes that produced no equality prefix
162            if values.is_empty() {
163                continue;
164            }
165
166            let score = values.len();
167            let cand = (score, IndexPlan { index, values });
168
169            match &best {
170                Some((best_score, _)) if *best_score >= score => { /* keep current best */ }
171                _ => best = Some(cand),
172            }
173        }
174
175        best.map(|(_, plan)| QueryPlan::Index(plan))
176    }
177
178    /// Find an equality clause (`field == ?`) anywhere in the filter tree and return the Value.
179    fn find_eq_value(filter: &FilterExpr, field: &str) -> Option<Value> {
180        match filter {
181            FilterExpr::Clause(c) if c.field == field && matches!(c.cmp, Cmp::Eq) => {
182                Some(c.value.clone())
183            }
184            // Walk conjunctive subtrees
185            FilterExpr::And(list) => list.iter().find_map(|f| Self::find_eq_value(f, field)),
186            _ => None,
187        }
188    }
189}