mentedb_query/
planner.rs

1//! Query planner: converts a parsed `Statement` into a `QueryPlan`.
2
3use crate::ast::*;
4use mentedb_core::edge::EdgeType;
5use mentedb_core::error::{MenteError, MenteResult};
6use mentedb_core::types::{MemoryId, Timestamp};
7use serde::{Deserialize, Serialize};
8
9/// A physical query plan describing how to execute a query.
10#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
11pub enum QueryPlan {
12    VectorSearch {
13        query: Vec<f32>,
14        k: usize,
15        filters: Vec<Filter>,
16    },
17    TagScan {
18        tags: Vec<String>,
19        filters: Vec<Filter>,
20        limit: Option<usize>,
21    },
22    TemporalScan {
23        start: Timestamp,
24        end: Timestamp,
25        filters: Vec<Filter>,
26    },
27    GraphTraversal {
28        start: MemoryId,
29        depth: usize,
30        edge_types: Vec<EdgeType>,
31    },
32    PointLookup {
33        id: MemoryId,
34    },
35    EdgeInsert {
36        source: MemoryId,
37        target: MemoryId,
38        edge_type: EdgeType,
39        weight: f32,
40    },
41    Delete {
42        id: MemoryId,
43    },
44    Consolidate {
45        filters: Vec<Filter>,
46    },
47}
48
49const DEFAULT_LIMIT: usize = 20;
50
51/// Produce an execution plan from a parsed statement.
52pub fn plan(statement: &Statement) -> MenteResult<QueryPlan> {
53    match statement {
54        Statement::Recall(recall) => plan_recall(recall),
55        Statement::Relate(relate) => plan_relate(relate),
56        Statement::Forget(forget) => Ok(QueryPlan::Delete { id: forget.target }),
57        Statement::Consolidate(cons) => Ok(QueryPlan::Consolidate {
58            filters: cons.filters.clone(),
59        }),
60        Statement::Traverse(trav) => plan_traverse(trav),
61    }
62}
63
64fn plan_recall(recall: &RecallStatement) -> MenteResult<QueryPlan> {
65    let limit = recall.limit.unwrap_or(DEFAULT_LIMIT);
66
67    // If there's a NEAR clause or a SimilarTo filter, use vector search
68    if let Some(ref vec) = recall.near {
69        return Ok(QueryPlan::VectorSearch {
70            query: vec.clone(),
71            k: limit,
72            filters: recall.filters.clone(),
73        });
74    }
75
76    // Check for SimilarTo operator in filters — implies vector search via text embedding
77    if let Some(sim_filter) = recall.filters.iter().find(|f| f.op == Operator::SimilarTo) {
78        if let Value::Text(ref _text) = sim_filter.value {
79            // The text is embedded at execution time (no embedder here), so keep
80            // ALL filters, including the SimilarTo one: the executor reads its
81            // text to build the query vector and applies the rest (type, tag).
82            return Ok(QueryPlan::VectorSearch {
83                query: Vec::new(), // placeholder — executor embeds the SimilarTo text
84                k: limit,
85                filters: recall.filters.clone(),
86            });
87        }
88        // SimilarTo with non-text value doesn't make sense
89        return Err(MenteError::Query(
90            "~> operator requires a text value on the right-hand side".into(),
91        ));
92    }
93
94    // If only tag filters, use TagScan
95    let tag_filters: Vec<&Filter> = recall
96        .filters
97        .iter()
98        .filter(|f| f.field == Field::Tag)
99        .collect();
100    if !tag_filters.is_empty() && recall.filters.iter().all(|f| f.field == Field::Tag) {
101        let tags: Vec<String> = tag_filters
102            .iter()
103            .filter_map(|f| match &f.value {
104                Value::Text(t) => Some(t.clone()),
105                _ => None,
106            })
107            .collect();
108        return Ok(QueryPlan::TagScan {
109            tags,
110            filters: Vec::new(),
111            limit: Some(limit),
112        });
113    }
114
115    // If time-range filters exist (created or accessed with range ops), use TemporalScan
116    let time_filters: Vec<&Filter> = recall
117        .filters
118        .iter()
119        .filter(|f| {
120            (f.field == Field::Created || f.field == Field::Accessed)
121                && matches!(
122                    f.op,
123                    Operator::Gt | Operator::Lt | Operator::Gte | Operator::Lte
124                )
125        })
126        .collect();
127
128    if !time_filters.is_empty() {
129        let remaining: Vec<Filter> = recall
130            .filters
131            .iter()
132            .filter(|f| {
133                !((f.field == Field::Created || f.field == Field::Accessed)
134                    && matches!(
135                        f.op,
136                        Operator::Gt | Operator::Lt | Operator::Gte | Operator::Lte
137                    ))
138            })
139            .cloned()
140            .collect();
141
142        let mut start: Timestamp = 0;
143        let mut end: Timestamp = u64::MAX;
144        for f in &time_filters {
145            if let Value::Text(ref s) = f.value {
146                // Simple heuristic: treat text timestamps as orderable strings for now.
147                // A real implementation would parse dates. We use 0/MAX as placeholders.
148                let _ = s; // acknowledged
149            }
150            if let Value::Integer(ts) = f.value {
151                let ts = ts as u64;
152                match f.op {
153                    Operator::Gt | Operator::Gte => start = ts,
154                    Operator::Lt | Operator::Lte => end = ts,
155                    _ => {}
156                }
157            }
158        }
159
160        return Ok(QueryPlan::TemporalScan {
161            start,
162            end,
163            filters: remaining,
164        });
165    }
166
167    // Fallback: tag scan with no tags (full scan with filters)
168    Ok(QueryPlan::TagScan {
169        tags: Vec::new(),
170        filters: recall.filters.clone(),
171        limit: Some(limit),
172    })
173}
174
175fn plan_relate(relate: &RelateStatement) -> MenteResult<QueryPlan> {
176    Ok(QueryPlan::EdgeInsert {
177        source: relate.source,
178        target: relate.target,
179        edge_type: relate.edge_type,
180        weight: relate.weight.unwrap_or(1.0),
181    })
182}
183
184fn plan_traverse(trav: &TraverseStatement) -> MenteResult<QueryPlan> {
185    Ok(QueryPlan::GraphTraversal {
186        start: trav.start,
187        depth: trav.depth,
188        edge_types: trav.edge_filter.clone().unwrap_or_default(),
189    })
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195    use crate::lexer::tokenize;
196    use crate::parser::Parser;
197
198    fn plan_mql(input: &str) -> QueryPlan {
199        let tokens = tokenize(input).unwrap();
200        let stmt = Parser::parse(&tokens).unwrap();
201        plan(&stmt).unwrap()
202    }
203
204    #[test]
205    fn test_near_produces_vector_search() {
206        let qp = plan_mql("RECALL memories NEAR [0.1, 0.2, 0.3] LIMIT 5");
207        match qp {
208            QueryPlan::VectorSearch { query, k, .. } => {
209                assert_eq!(query, vec![0.1, 0.2, 0.3]);
210                assert_eq!(k, 5);
211            }
212            _ => panic!("expected VectorSearch, got {:?}", qp),
213        }
214    }
215
216    #[test]
217    fn test_similar_to_produces_vector_search() {
218        let qp = plan_mql(r#"RECALL memories WHERE content ~> "database migration" LIMIT 10"#);
219        match qp {
220            QueryPlan::VectorSearch { k, .. } => {
221                assert_eq!(k, 10);
222            }
223            _ => panic!("expected VectorSearch, got {:?}", qp),
224        }
225    }
226
227    #[test]
228    fn test_tag_filter_produces_tag_scan() {
229        let qp = plan_mql(r#"RECALL memories WHERE tag = "backend" LIMIT 5"#);
230        match qp {
231            QueryPlan::TagScan { tags, limit, .. } => {
232                assert_eq!(tags, vec!["backend".to_string()]);
233                assert_eq!(limit, Some(5));
234            }
235            _ => panic!("expected TagScan, got {:?}", qp),
236        }
237    }
238
239    #[test]
240    fn test_forget_produces_delete() {
241        let qp = plan_mql("FORGET 550e8400-e29b-41d4-a716-446655440000");
242        match qp {
243            QueryPlan::Delete { id } => {
244                assert_eq!(
245                    id,
246                    "550e8400-e29b-41d4-a716-446655440000"
247                        .parse::<MemoryId>()
248                        .unwrap()
249                );
250            }
251            _ => panic!("expected Delete, got {:?}", qp),
252        }
253    }
254
255    #[test]
256    fn test_traverse_produces_graph_traversal() {
257        let qp = plan_mql(
258            "TRAVERSE 550e8400-e29b-41d4-a716-446655440000 DEPTH 3 WHERE edge_type = caused",
259        );
260        match qp {
261            QueryPlan::GraphTraversal {
262                depth, edge_types, ..
263            } => {
264                assert_eq!(depth, 3);
265                assert_eq!(edge_types, vec![EdgeType::Caused]);
266            }
267            _ => panic!("expected GraphTraversal, got {:?}", qp),
268        }
269    }
270
271    #[test]
272    fn test_relate_produces_edge_insert() {
273        let qp = plan_mql(
274            "RELATE 550e8400-e29b-41d4-a716-446655440000 -> 660e8400-e29b-41d4-a716-446655440000 AS caused WITH weight = 0.8",
275        );
276        match qp {
277            QueryPlan::EdgeInsert {
278                edge_type, weight, ..
279            } => {
280                assert_eq!(edge_type, EdgeType::Caused);
281                assert!((weight - 0.8).abs() < f32::EPSILON);
282            }
283            _ => panic!("expected EdgeInsert, got {:?}", qp),
284        }
285    }
286}
mentedb_query/planner.rs

mentedb_query/
planner.rs
