Skip to main content

ankurah_core/indexing/
key_spec.rs

1use crate::value::ValueType;
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
5pub struct KeySpec {
6    pub keyparts: Vec<IndexKeyPart>,
7}
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
10pub enum IndexDirection {
11    Asc,
12    Desc,
13}
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
16pub enum NullsOrder {
17    First,
18    Last,
19}
20
21#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
22pub struct IndexKeyPart {
23    pub column: String,
24    /// Optional path within property value (for JSON, future Ref, etc.)
25    pub sub_path: Option<Vec<String>>,
26    pub direction: IndexDirection, // ASC/DESC
27    pub value_type: ValueType,     // Expected type for this key component
28    pub nulls: Option<NullsOrder>, // PG-style NULLS FIRST/LAST (optional)
29    pub collation: Option<String>, // collation name/id if relevant (optional)
30}
31
32impl IndexKeyPart {
33    pub fn asc<S: Into<String>>(col: S, value_type: ValueType) -> Self {
34        Self { column: col.into(), sub_path: None, direction: IndexDirection::Asc, value_type, nulls: None, collation: None }
35    }
36    pub fn desc<S: Into<String>>(col: S, value_type: ValueType) -> Self {
37        Self { column: col.into(), sub_path: None, direction: IndexDirection::Desc, value_type, nulls: None, collation: None }
38    }
39
40    /// Create from a PathExpr (handles multi-step paths)
41    pub fn from_path(path: &ankql::ast::PathExpr, direction: IndexDirection, value_type: ValueType) -> Self {
42        let (column, sub_path) = if path.steps.len() == 1 {
43            (path.steps[0].clone(), None)
44        } else {
45            let column = path.steps[0].clone();
46            let sub_path = path.steps[1..].to_vec();
47            (column, Some(sub_path))
48        };
49        Self { column, sub_path, direction, value_type, nulls: None, collation: None }
50    }
51
52    /// Full path as a flat string (e.g., "context.session_id")
53    pub fn full_path(&self) -> String {
54        match &self.sub_path {
55            None => self.column.clone(),
56            Some(sub) => {
57                let mut parts = vec![self.column.clone()];
58                parts.extend(sub.clone());
59                parts.join(".")
60            }
61        }
62    }
63
64    /// Create from a flat path string (e.g., "context.session_id")
65    pub fn from_flat_path(path: &str, direction: IndexDirection, value_type: ValueType) -> Self {
66        let parts: Vec<&str> = path.split('.').collect();
67        let (column, sub_path) = if parts.len() == 1 {
68            (parts[0].to_string(), None)
69        } else {
70            let column = parts[0].to_string();
71            let sub_path: Vec<String> = parts[1..].iter().map(|s| s.to_string()).collect();
72            (column, Some(sub_path))
73        };
74        Self { column, sub_path, direction, value_type, nulls: None, collation: None }
75    }
76
77    /// Create ascending keypart from flat path
78    pub fn asc_path(path: &str, value_type: ValueType) -> Self { Self::from_flat_path(path, IndexDirection::Asc, value_type) }
79
80    /// Create descending keypart from flat path
81    pub fn desc_path(path: &str, value_type: ValueType) -> Self { Self::from_flat_path(path, IndexDirection::Desc, value_type) }
82}
83
84impl IndexDirection {
85    pub fn is_desc(&self) -> bool { matches!(self, IndexDirection::Desc) }
86}
87
88impl KeySpec {
89    pub fn new(keyparts: Vec<IndexKeyPart>) -> Self { Self { keyparts } }
90
91    /// Simple name generator similar to your existing helper.
92    pub fn name_with(&self, prefix: &str, delim: &str) -> String {
93        let fields: Vec<String> = self
94            .keyparts
95            .iter()
96            .map(|k| {
97                let dir = match k.direction {
98                    IndexDirection::Asc => "asc",
99                    IndexDirection::Desc => "desc",
100                };
101                let col_name = k.full_path();
102                if k.collation.is_some() || k.nulls.is_some() {
103                    // include extras only if present
104                    let mut extras = Vec::new();
105                    if let Some(c) = &k.collation {
106                        extras.push(format!("collate={}", c));
107                    }
108                    if let Some(n) = &k.nulls {
109                        extras.push(format!("nulls={:?}", n).to_lowercase());
110                    }
111                    format!("{} {}({})", col_name, dir, extras.join(","))
112                } else {
113                    format!("{} {}", col_name, dir)
114                }
115            })
116            .collect();
117
118        if prefix.is_empty() {
119            fields.join(delim)
120        } else {
121            format!("{}{}{}", prefix, delim, fields.join(delim))
122        }
123    }
124
125    /// Checks if this IndexSpec can be satisfied by another IndexSpec
126    /// Returns Yes if this is a prefix subset of other
127    /// Returns Inverse if this is a prefix subset of other with all directions flipped
128    /// Returns No if neither condition is met
129    pub fn matches(&self, other: &KeySpec) -> Option<IndexSpecMatch> {
130        if self.keyparts.len() > other.keyparts.len() {
131            return None;
132        }
133
134        let mut direct_match = true;
135        let mut inverse_match = true;
136
137        for (self_keypart, other_keypart) in self.keyparts.iter().zip(other.keyparts.iter()) {
138            // Both column and sub_path must match
139            if self_keypart.column != other_keypart.column || self_keypart.sub_path != other_keypart.sub_path {
140                return None;
141            }
142
143            if self_keypart.direction != other_keypart.direction {
144                direct_match = false;
145            }
146
147            if self_keypart.direction == other_keypart.direction {
148                inverse_match = false;
149            }
150        }
151
152        if direct_match {
153            Some(IndexSpecMatch::Match)
154        } else if inverse_match {
155            Some(IndexSpecMatch::Inverse)
156        } else {
157            None
158        }
159    }
160}
161
162#[derive(Debug, Clone, PartialEq, Eq, Hash)]
163pub enum IndexSpecMatch {
164    /// The index specs match
165    Match,
166    /// The index specs match, but scan direction must be inverted
167    Inverse,
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    #[test]
175    fn test_exact_match() {
176        let spec1 = KeySpec { keyparts: vec![IndexKeyPart::asc("a", ValueType::String), IndexKeyPart::desc("b", ValueType::String)] };
177        let spec2 = KeySpec { keyparts: vec![IndexKeyPart::asc("a", ValueType::String), IndexKeyPart::desc("b", ValueType::String)] };
178
179        assert_eq!(spec1.matches(&spec2), Some(IndexSpecMatch::Match));
180    }
181
182    #[test]
183    fn test_prefix_match() {
184        // +a, -b matches +a, -b, +c
185        let query_spec = KeySpec { keyparts: vec![IndexKeyPart::asc("a", ValueType::String), IndexKeyPart::desc("b", ValueType::String)] };
186        let index_spec = KeySpec {
187            keyparts: vec![
188                IndexKeyPart::asc("a", ValueType::String),
189                IndexKeyPart::desc("b", ValueType::String),
190                IndexKeyPart::asc("c", ValueType::String),
191            ],
192        };
193
194        assert_eq!(query_spec.matches(&index_spec), Some(IndexSpecMatch::Match));
195    }
196
197    #[test]
198    fn test_inverse_exact_match() {
199        // +a, -b matches -a, +b (inverse)
200        let query_spec = KeySpec { keyparts: vec![IndexKeyPart::asc("a", ValueType::String), IndexKeyPart::desc("b", ValueType::String)] };
201        let index_spec = KeySpec { keyparts: vec![IndexKeyPart::desc("a", ValueType::String), IndexKeyPart::asc("b", ValueType::String)] };
202
203        assert_eq!(query_spec.matches(&index_spec), Some(IndexSpecMatch::Inverse));
204    }
205
206    #[test]
207    fn test_inverse_prefix_match() {
208        // +a, -b matches -a, +b, +c (inverse prefix)
209        let query_spec = KeySpec { keyparts: vec![IndexKeyPart::asc("a", ValueType::String), IndexKeyPart::desc("b", ValueType::String)] };
210        let index_spec = KeySpec {
211            keyparts: vec![
212                IndexKeyPart::desc("a", ValueType::String),
213                IndexKeyPart::asc("b", ValueType::String),
214                IndexKeyPart::asc("c", ValueType::String),
215            ],
216        };
217
218        assert_eq!(query_spec.matches(&index_spec), Some(IndexSpecMatch::Inverse));
219    }
220
221    #[test]
222    fn test_user_example() {
223        // "+a, -b matches +a, -b, any c AND -a, +b, any c"
224        let query_spec = KeySpec { keyparts: vec![IndexKeyPart::asc("a", ValueType::String), IndexKeyPart::desc("b", ValueType::String)] };
225
226        // Test direct match: +a, -b, +c
227        let index_spec1 = KeySpec {
228            keyparts: vec![
229                IndexKeyPart::asc("a", ValueType::String),
230                IndexKeyPart::desc("b", ValueType::String),
231                IndexKeyPart::asc("c", ValueType::String),
232            ],
233        };
234        assert_eq!(query_spec.matches(&index_spec1), Some(IndexSpecMatch::Match));
235
236        // Test inverse match: -a, +b, -c
237        let index_spec2 = KeySpec {
238            keyparts: vec![
239                IndexKeyPart::desc("a", ValueType::String),
240                IndexKeyPart::asc("b", ValueType::String),
241                IndexKeyPart::desc("c", ValueType::String),
242            ],
243        };
244        assert_eq!(query_spec.matches(&index_spec2), Some(IndexSpecMatch::Inverse));
245    }
246
247    #[test]
248    fn test_no_match_different_fields() {
249        let query_spec = KeySpec { keyparts: vec![IndexKeyPart::asc("a", ValueType::String), IndexKeyPart::desc("b", ValueType::String)] };
250        let index_spec = KeySpec { keyparts: vec![IndexKeyPart::asc("x", ValueType::String), IndexKeyPart::desc("y", ValueType::String)] };
251
252        assert_eq!(query_spec.matches(&index_spec), None);
253    }
254
255    #[test]
256    fn test_no_match_partial_field_overlap() {
257        // +a, -b does not match +a, +b (different direction on second field)
258        let query_spec = KeySpec { keyparts: vec![IndexKeyPart::asc("a", ValueType::String), IndexKeyPart::desc("b", ValueType::String)] };
259        let index_spec = KeySpec { keyparts: vec![IndexKeyPart::asc("a", ValueType::String), IndexKeyPart::asc("b", ValueType::String)] };
260
261        assert_eq!(query_spec.matches(&index_spec), None);
262    }
263
264    #[test]
265    fn test_no_match_query_longer_than_index() {
266        // +a, -b, +c cannot match +a (query is longer than index)
267        let query_spec = KeySpec {
268            keyparts: vec![
269                IndexKeyPart::asc("a", ValueType::String),
270                IndexKeyPart::desc("b", ValueType::String),
271                IndexKeyPart::asc("c", ValueType::String),
272            ],
273        };
274        let index_spec = KeySpec { keyparts: vec![IndexKeyPart::asc("a", ValueType::String)] };
275
276        assert_eq!(query_spec.matches(&index_spec), None);
277    }
278
279    #[test]
280    fn test_empty_specs() {
281        let empty_spec = KeySpec { keyparts: vec![] };
282        let non_empty_spec = KeySpec { keyparts: vec![IndexKeyPart::asc("a", ValueType::String)] };
283
284        // Empty spec matches any spec (empty prefix)
285        assert_eq!(empty_spec.matches(&non_empty_spec), Some(IndexSpecMatch::Match));
286        assert_eq!(empty_spec.matches(&empty_spec), Some(IndexSpecMatch::Match));
287
288        // Non-empty spec does not match empty spec
289        assert_eq!(non_empty_spec.matches(&empty_spec), None);
290    }
291
292    #[test]
293    fn test_single_field_cases() {
294        let asc_spec = KeySpec { keyparts: vec![IndexKeyPart::asc("a", ValueType::String)] };
295        let desc_spec = KeySpec { keyparts: vec![IndexKeyPart::desc("a", ValueType::String)] };
296
297        // Direct match
298        assert_eq!(asc_spec.matches(&asc_spec), Some(IndexSpecMatch::Match));
299
300        // Inverse match
301        assert_eq!(asc_spec.matches(&desc_spec), Some(IndexSpecMatch::Inverse));
302        assert_eq!(desc_spec.matches(&asc_spec), Some(IndexSpecMatch::Inverse));
303    }
304
305    #[test]
306    fn test_complex_multi_field_scenarios() {
307        // Test various combinations with 3+ fields
308        let query_spec = KeySpec {
309            keyparts: vec![
310                IndexKeyPart::asc("a", ValueType::String),
311                IndexKeyPart::desc("b", ValueType::String),
312                IndexKeyPart::asc("c", ValueType::String),
313            ],
314        };
315
316        // Exact match with additional fields
317        let index_spec1 = KeySpec {
318            keyparts: vec![
319                IndexKeyPart::asc("a", ValueType::String),
320                IndexKeyPart::desc("b", ValueType::String),
321                IndexKeyPart::asc("c", ValueType::String),
322                IndexKeyPart::desc("d", ValueType::String),
323            ],
324        };
325        assert_eq!(query_spec.matches(&index_spec1), Some(IndexSpecMatch::Match));
326
327        // Inverse match with additional fields
328        let index_spec2 = KeySpec {
329            keyparts: vec![
330                IndexKeyPart::desc("a", ValueType::String),
331                IndexKeyPart::asc("b", ValueType::String),
332                IndexKeyPart::desc("c", ValueType::String),
333                IndexKeyPart::asc("d", ValueType::String),
334            ],
335        };
336        assert_eq!(query_spec.matches(&index_spec2), Some(IndexSpecMatch::Inverse));
337
338        // No match - mixed directions that don't form inverse
339        let index_spec3 = KeySpec {
340            keyparts: vec![
341                IndexKeyPart::asc("a", ValueType::String),
342                IndexKeyPart::asc("b", ValueType::String),
343                IndexKeyPart::desc("c", ValueType::String),
344            ],
345        };
346        assert_eq!(query_spec.matches(&index_spec3), None);
347    }
348
349    #[test]
350    fn test_helper_methods() {
351        // Test IndexKeyPart helper methods
352        let asc_keypart = IndexKeyPart::asc("test", ValueType::String);
353        assert_eq!(asc_keypart.column, "test");
354        assert_eq!(asc_keypart.direction, IndexDirection::Asc);
355        assert_eq!(asc_keypart.nulls, None);
356        assert_eq!(asc_keypart.collation, None);
357
358        let desc_keypart = IndexKeyPart::desc("test", ValueType::String);
359        assert_eq!(desc_keypart.column, "test");
360        assert_eq!(desc_keypart.direction, IndexDirection::Desc);
361        assert_eq!(desc_keypart.nulls, None);
362        assert_eq!(desc_keypart.collation, None);
363    }
364
365    #[test]
366    fn test_edge_case_behaviors() {
367        // Test that matches works correctly with various edge cases
368        let spec = KeySpec {
369            keyparts: vec![
370                IndexKeyPart::asc("a", ValueType::String),
371                IndexKeyPart::desc("b", ValueType::String),
372                IndexKeyPart::asc("c", ValueType::String),
373            ],
374        };
375
376        // Self-match should always be Yes
377        assert_eq!(spec.matches(&spec), Some(IndexSpecMatch::Match));
378
379        // Empty spec matches any spec
380        let empty = KeySpec { keyparts: vec![] };
381        assert_eq!(empty.matches(&spec), Some(IndexSpecMatch::Match));
382    }
383}