aegis_query/
index.rs

1//! Query Engine Indexes
2//!
3//! B-tree and hash index implementations for query optimization.
4//!
5//! @version 0.1.0
6//! @author AutomataNexus Development Team
7
8use std::collections::{BTreeMap, HashMap, HashSet};
9use std::sync::RwLock;
10
11use aegis_common::Value;
12
13// =============================================================================
14// Index Types
15// =============================================================================
16
17/// The type of index to use.
18#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
19#[derive(Default)]
20pub enum IndexType {
21    /// B-tree index for range queries and ordered scans.
22    #[default]
23    BTree,
24    /// Hash index for equality lookups (faster for exact matches).
25    Hash,
26}
27
28
29// =============================================================================
30// Index Key
31// =============================================================================
32
33/// A key in an index, supporting composite keys.
34#[derive(Debug, Clone, PartialEq, Eq, Hash)]
35pub struct IndexKey {
36    pub values: Vec<IndexValue>,
37}
38
39/// A single value in an index key.
40#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
41pub enum IndexValue {
42    Null,
43    Bool(bool),
44    Int(i64),
45    String(String),
46    /// For floats, we use ordered representation for BTree compatibility
47    Float(OrderedFloat),
48}
49
50/// Wrapper for f64 that implements Ord for BTree storage.
51#[derive(Debug, Clone, Copy)]
52pub struct OrderedFloat(pub f64);
53
54impl PartialEq for OrderedFloat {
55    fn eq(&self, other: &Self) -> bool {
56        self.0.to_bits() == other.0.to_bits()
57    }
58}
59
60impl Eq for OrderedFloat {}
61
62impl std::hash::Hash for OrderedFloat {
63    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
64        self.0.to_bits().hash(state);
65    }
66}
67
68impl PartialOrd for OrderedFloat {
69    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
70        Some(self.cmp(other))
71    }
72}
73
74impl Ord for OrderedFloat {
75    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
76        self.0.partial_cmp(&other.0).unwrap_or(std::cmp::Ordering::Equal)
77    }
78}
79
80impl IndexKey {
81    pub fn new(values: Vec<IndexValue>) -> Self {
82        Self { values }
83    }
84
85    pub fn single(value: IndexValue) -> Self {
86        Self { values: vec![value] }
87    }
88
89    /// Convert a Value to an IndexValue.
90    pub fn from_value(value: &Value) -> IndexValue {
91        match value {
92            Value::Null => IndexValue::Null,
93            Value::Boolean(b) => IndexValue::Bool(*b),
94            Value::Integer(i) => IndexValue::Int(*i),
95            Value::Float(f) => IndexValue::Float(OrderedFloat(*f)),
96            Value::String(s) => IndexValue::String(s.clone()),
97            Value::Timestamp(t) => IndexValue::Int(t.timestamp()),
98            Value::Bytes(b) => IndexValue::String(format!("{:?}", b)),
99            Value::Array(arr) => IndexValue::String(format!("{:?}", arr)),
100            Value::Object(obj) => IndexValue::String(format!("{:?}", obj)),
101        }
102    }
103
104    /// Create an IndexKey from multiple Values.
105    pub fn from_values(values: &[Value]) -> Self {
106        Self {
107            values: values.iter().map(Self::from_value).collect(),
108        }
109    }
110}
111
112impl PartialOrd for IndexKey {
113    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
114        Some(self.cmp(other))
115    }
116}
117
118impl Ord for IndexKey {
119    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
120        self.values.cmp(&other.values)
121    }
122}
123
124// =============================================================================
125// Row ID
126// =============================================================================
127
128/// A row identifier in an index.
129pub type RowId = usize;
130
131// =============================================================================
132// B-Tree Index
133// =============================================================================
134
135/// A B-tree index for range queries and ordered scans.
136pub struct BTreeIndex {
137    /// Index name.
138    pub name: String,
139    /// Table this index belongs to.
140    pub table: String,
141    /// Column names in the index (supports composite keys).
142    pub columns: Vec<String>,
143    /// Whether this is a unique index.
144    pub unique: bool,
145    /// The actual B-tree storing key -> row IDs.
146    tree: RwLock<BTreeMap<IndexKey, HashSet<RowId>>>,
147    /// Number of entries in the index.
148    entry_count: RwLock<usize>,
149}
150
151impl BTreeIndex {
152    /// Create a new B-tree index.
153    pub fn new(name: String, table: String, columns: Vec<String>, unique: bool) -> Self {
154        Self {
155            name,
156            table,
157            columns,
158            unique,
159            tree: RwLock::new(BTreeMap::new()),
160            entry_count: RwLock::new(0),
161        }
162    }
163
164    /// Insert a key-rowid pair into the index.
165    pub fn insert(&self, key: IndexKey, row_id: RowId) -> Result<(), IndexError> {
166        let mut tree = self.tree.write().expect("BTreeIndex tree lock poisoned");
167
168        if self.unique {
169            if let Some(existing) = tree.get(&key) {
170                if !existing.is_empty() {
171                    return Err(IndexError::DuplicateKey(format!(
172                        "Duplicate key in unique index '{}'",
173                        self.name
174                    )));
175                }
176            }
177        }
178
179        let entry = tree.entry(key).or_default();
180        if entry.insert(row_id) {
181            *self.entry_count.write().expect("BTreeIndex entry_count lock poisoned") += 1;
182        }
183        Ok(())
184    }
185
186    /// Remove a key-rowid pair from the index.
187    pub fn remove(&self, key: &IndexKey, row_id: RowId) -> bool {
188        let mut tree = self.tree.write().expect("BTreeIndex tree lock poisoned");
189        if let Some(row_ids) = tree.get_mut(key) {
190            if row_ids.remove(&row_id) {
191                *self.entry_count.write().expect("BTreeIndex entry_count lock poisoned") -= 1;
192                if row_ids.is_empty() {
193                    tree.remove(key);
194                }
195                return true;
196            }
197        }
198        false
199    }
200
201    /// Look up an exact key.
202    pub fn get(&self, key: &IndexKey) -> Vec<RowId> {
203        let tree = self.tree.read().expect("BTreeIndex tree lock poisoned");
204        tree.get(key)
205            .map(|ids| ids.iter().copied().collect())
206            .unwrap_or_default()
207    }
208
209    /// Range scan from start to end (inclusive/exclusive based on flags).
210    pub fn range(
211        &self,
212        start: Option<&IndexKey>,
213        end: Option<&IndexKey>,
214        start_inclusive: bool,
215        end_inclusive: bool,
216    ) -> Vec<RowId> {
217        let tree = self.tree.read().expect("BTreeIndex tree lock poisoned");
218        let mut result = Vec::new();
219
220        use std::ops::Bound;
221
222        let start_bound = match start {
223            Some(k) if start_inclusive => Bound::Included(k.clone()),
224            Some(k) => Bound::Excluded(k.clone()),
225            None => Bound::Unbounded,
226        };
227
228        let end_bound = match end {
229            Some(k) if end_inclusive => Bound::Included(k.clone()),
230            Some(k) => Bound::Excluded(k.clone()),
231            None => Bound::Unbounded,
232        };
233
234        for (_, row_ids) in tree.range((start_bound, end_bound)) {
235            result.extend(row_ids.iter().copied());
236        }
237
238        result
239    }
240
241    /// Get all row IDs in index order.
242    pub fn scan_all(&self) -> Vec<RowId> {
243        let tree = self.tree.read().expect("BTreeIndex tree lock poisoned");
244        tree.values().flat_map(|ids| ids.iter().copied()).collect()
245    }
246
247    /// Get the number of entries in the index.
248    pub fn len(&self) -> usize {
249        *self.entry_count.read().expect("BTreeIndex entry_count lock poisoned")
250    }
251
252    /// Check if the index is empty.
253    pub fn is_empty(&self) -> bool {
254        self.len() == 0
255    }
256
257    /// Clear all entries from the index.
258    pub fn clear(&self) {
259        let mut tree = self.tree.write().expect("BTreeIndex tree lock poisoned");
260        tree.clear();
261        *self.entry_count.write().expect("BTreeIndex entry_count lock poisoned") = 0;
262    }
263}
264
265// =============================================================================
266// Hash Index
267// =============================================================================
268
269/// A hash index for fast equality lookups.
270pub struct HashIndex {
271    /// Index name.
272    pub name: String,
273    /// Table this index belongs to.
274    pub table: String,
275    /// Column names in the index.
276    pub columns: Vec<String>,
277    /// Whether this is a unique index.
278    pub unique: bool,
279    /// The hash map storing key -> row IDs.
280    map: RwLock<HashMap<IndexKey, HashSet<RowId>>>,
281    /// Number of entries in the index.
282    entry_count: RwLock<usize>,
283}
284
285impl HashIndex {
286    /// Create a new hash index.
287    pub fn new(name: String, table: String, columns: Vec<String>, unique: bool) -> Self {
288        Self {
289            name,
290            table,
291            columns,
292            unique,
293            map: RwLock::new(HashMap::new()),
294            entry_count: RwLock::new(0),
295        }
296    }
297
298    /// Insert a key-rowid pair into the index.
299    pub fn insert(&self, key: IndexKey, row_id: RowId) -> Result<(), IndexError> {
300        let mut map = self.map.write().expect("HashIndex map lock poisoned");
301
302        if self.unique {
303            if let Some(existing) = map.get(&key) {
304                if !existing.is_empty() {
305                    return Err(IndexError::DuplicateKey(format!(
306                        "Duplicate key in unique index '{}'",
307                        self.name
308                    )));
309                }
310            }
311        }
312
313        let entry = map.entry(key).or_default();
314        if entry.insert(row_id) {
315            *self.entry_count.write().expect("HashIndex entry_count lock poisoned") += 1;
316        }
317        Ok(())
318    }
319
320    /// Remove a key-rowid pair from the index.
321    pub fn remove(&self, key: &IndexKey, row_id: RowId) -> bool {
322        let mut map = self.map.write().expect("HashIndex map lock poisoned");
323        if let Some(row_ids) = map.get_mut(key) {
324            if row_ids.remove(&row_id) {
325                *self.entry_count.write().expect("HashIndex entry_count lock poisoned") -= 1;
326                if row_ids.is_empty() {
327                    map.remove(key);
328                }
329                return true;
330            }
331        }
332        false
333    }
334
335    /// Look up an exact key.
336    pub fn get(&self, key: &IndexKey) -> Vec<RowId> {
337        let map = self.map.read().expect("HashIndex map lock poisoned");
338        map.get(key)
339            .map(|ids| ids.iter().copied().collect())
340            .unwrap_or_default()
341    }
342
343    /// Get the number of entries in the index.
344    pub fn len(&self) -> usize {
345        *self.entry_count.read().expect("HashIndex entry_count lock poisoned")
346    }
347
348    /// Check if the index is empty.
349    pub fn is_empty(&self) -> bool {
350        self.len() == 0
351    }
352
353    /// Clear all entries from the index.
354    pub fn clear(&self) {
355        let mut map = self.map.write().expect("HashIndex map lock poisoned");
356        map.clear();
357        *self.entry_count.write().expect("HashIndex entry_count lock poisoned") = 0;
358    }
359}
360
361// =============================================================================
362// Unified Index Trait
363// =============================================================================
364
365/// A trait for index operations.
366pub trait Index: Send + Sync {
367    /// Get the index name.
368    fn name(&self) -> &str;
369
370    /// Get the table name.
371    fn table(&self) -> &str;
372
373    /// Get the column names.
374    fn columns(&self) -> &[String];
375
376    /// Check if this is a unique index.
377    fn is_unique(&self) -> bool;
378
379    /// Get the index type.
380    fn index_type(&self) -> IndexType;
381
382    /// Insert a key-rowid pair.
383    fn insert(&self, key: IndexKey, row_id: RowId) -> Result<(), IndexError>;
384
385    /// Remove a key-rowid pair.
386    fn remove(&self, key: &IndexKey, row_id: RowId) -> bool;
387
388    /// Look up an exact key.
389    fn get(&self, key: &IndexKey) -> Vec<RowId>;
390
391    /// Get the number of entries.
392    fn len(&self) -> usize;
393
394    /// Check if empty.
395    fn is_empty(&self) -> bool;
396
397    /// Clear the index.
398    fn clear(&self);
399}
400
401impl Index for BTreeIndex {
402    fn name(&self) -> &str { &self.name }
403    fn table(&self) -> &str { &self.table }
404    fn columns(&self) -> &[String] { &self.columns }
405    fn is_unique(&self) -> bool { self.unique }
406    fn index_type(&self) -> IndexType { IndexType::BTree }
407    fn insert(&self, key: IndexKey, row_id: RowId) -> Result<(), IndexError> { self.insert(key, row_id) }
408    fn remove(&self, key: &IndexKey, row_id: RowId) -> bool { self.remove(key, row_id) }
409    fn get(&self, key: &IndexKey) -> Vec<RowId> { self.get(key) }
410    fn len(&self) -> usize { self.len() }
411    fn is_empty(&self) -> bool { self.is_empty() }
412    fn clear(&self) { self.clear() }
413}
414
415impl Index for HashIndex {
416    fn name(&self) -> &str { &self.name }
417    fn table(&self) -> &str { &self.table }
418    fn columns(&self) -> &[String] { &self.columns }
419    fn is_unique(&self) -> bool { self.unique }
420    fn index_type(&self) -> IndexType { IndexType::Hash }
421    fn insert(&self, key: IndexKey, row_id: RowId) -> Result<(), IndexError> { self.insert(key, row_id) }
422    fn remove(&self, key: &IndexKey, row_id: RowId) -> bool { self.remove(key, row_id) }
423    fn get(&self, key: &IndexKey) -> Vec<RowId> { self.get(key) }
424    fn len(&self) -> usize { self.len() }
425    fn is_empty(&self) -> bool { self.is_empty() }
426    fn clear(&self) { self.clear() }
427}
428
429// =============================================================================
430// Index Error
431// =============================================================================
432
433/// Errors that can occur during index operations.
434#[derive(Debug, Clone)]
435pub enum IndexError {
436    DuplicateKey(String),
437    IndexNotFound(String),
438    InvalidKey(String),
439}
440
441impl std::fmt::Display for IndexError {
442    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
443        match self {
444            IndexError::DuplicateKey(msg) => write!(f, "Duplicate key: {}", msg),
445            IndexError::IndexNotFound(msg) => write!(f, "Index not found: {}", msg),
446            IndexError::InvalidKey(msg) => write!(f, "Invalid key: {}", msg),
447        }
448    }
449}
450
451impl std::error::Error for IndexError {}
452
453// =============================================================================
454// Index Manager
455// =============================================================================
456
457/// Manages all indexes for a table.
458pub struct TableIndexManager {
459    /// Table name.
460    table: String,
461    /// B-tree indexes by name.
462    btree_indexes: RwLock<HashMap<String, BTreeIndex>>,
463    /// Hash indexes by name.
464    hash_indexes: RwLock<HashMap<String, HashIndex>>,
465}
466
467impl TableIndexManager {
468    /// Create a new index manager for a table.
469    pub fn new(table: String) -> Self {
470        Self {
471            table,
472            btree_indexes: RwLock::new(HashMap::new()),
473            hash_indexes: RwLock::new(HashMap::new()),
474        }
475    }
476
477    /// Create a new index.
478    pub fn create_index(
479        &self,
480        name: String,
481        columns: Vec<String>,
482        unique: bool,
483        index_type: IndexType,
484    ) -> Result<(), IndexError> {
485        match index_type {
486            IndexType::BTree => {
487                let mut indexes = self.btree_indexes.write().expect("TableIndexManager btree_indexes lock poisoned");
488                if indexes.contains_key(&name) {
489                    return Err(IndexError::DuplicateKey(format!("Index '{}' already exists", name)));
490                }
491                indexes.insert(
492                    name.clone(),
493                    BTreeIndex::new(name, self.table.clone(), columns, unique),
494                );
495            }
496            IndexType::Hash => {
497                let mut indexes = self.hash_indexes.write().expect("TableIndexManager hash_indexes lock poisoned");
498                if indexes.contains_key(&name) {
499                    return Err(IndexError::DuplicateKey(format!("Index '{}' already exists", name)));
500                }
501                indexes.insert(
502                    name.clone(),
503                    HashIndex::new(name, self.table.clone(), columns, unique),
504                );
505            }
506        }
507        Ok(())
508    }
509
510    /// Drop an index.
511    pub fn drop_index(&self, name: &str) -> Result<(), IndexError> {
512        let mut btree = self.btree_indexes.write().expect("TableIndexManager btree_indexes lock poisoned");
513        if btree.remove(name).is_some() {
514            return Ok(());
515        }
516
517        let mut hash = self.hash_indexes.write().expect("TableIndexManager hash_indexes lock poisoned");
518        if hash.remove(name).is_some() {
519            return Ok(());
520        }
521
522        Err(IndexError::IndexNotFound(name.to_string()))
523    }
524
525    /// Insert a row into all indexes.
526    pub fn insert_row(
527        &self,
528        row_id: RowId,
529        column_values: &HashMap<String, Value>,
530    ) -> Result<(), IndexError> {
531        // Insert into B-tree indexes
532        let btree = self.btree_indexes.read().expect("TableIndexManager btree_indexes lock poisoned");
533        for index in btree.values() {
534            let key = self.build_key(&index.columns, column_values);
535            index.insert(key, row_id)?;
536        }
537
538        // Insert into hash indexes
539        let hash = self.hash_indexes.read().expect("TableIndexManager hash_indexes lock poisoned");
540        for index in hash.values() {
541            let key = self.build_key(&index.columns, column_values);
542            index.insert(key, row_id)?;
543        }
544
545        Ok(())
546    }
547
548    /// Remove a row from all indexes.
549    pub fn remove_row(&self, row_id: RowId, column_values: &HashMap<String, Value>) {
550        // Remove from B-tree indexes
551        let btree = self.btree_indexes.read().expect("TableIndexManager btree_indexes lock poisoned");
552        for index in btree.values() {
553            let key = self.build_key(&index.columns, column_values);
554            index.remove(&key, row_id);
555        }
556
557        // Remove from hash indexes
558        let hash = self.hash_indexes.read().expect("TableIndexManager hash_indexes lock poisoned");
559        for index in hash.values() {
560            let key = self.build_key(&index.columns, column_values);
561            index.remove(&key, row_id);
562        }
563    }
564
565    /// Update a row in all indexes (remove old, insert new).
566    pub fn update_row(
567        &self,
568        row_id: RowId,
569        old_values: &HashMap<String, Value>,
570        new_values: &HashMap<String, Value>,
571    ) -> Result<(), IndexError> {
572        self.remove_row(row_id, old_values);
573        self.insert_row(row_id, new_values)
574    }
575
576    /// Look up rows by index for equality.
577    pub fn lookup_eq(&self, index_name: &str, key: &IndexKey) -> Option<Vec<RowId>> {
578        // Check B-tree indexes
579        let btree = self.btree_indexes.read().expect("TableIndexManager btree_indexes lock poisoned");
580        if let Some(index) = btree.get(index_name) {
581            return Some(index.get(key));
582        }
583
584        // Check hash indexes
585        let hash = self.hash_indexes.read().expect("TableIndexManager hash_indexes lock poisoned");
586        if let Some(index) = hash.get(index_name) {
587            return Some(index.get(key));
588        }
589
590        None
591    }
592
593    /// Range lookup (only works for B-tree indexes).
594    pub fn lookup_range(
595        &self,
596        index_name: &str,
597        start: Option<&IndexKey>,
598        end: Option<&IndexKey>,
599        start_inclusive: bool,
600        end_inclusive: bool,
601    ) -> Option<Vec<RowId>> {
602        let btree = self.btree_indexes.read().expect("TableIndexManager btree_indexes lock poisoned");
603        btree.get(index_name).map(|index| {
604            index.range(start, end, start_inclusive, end_inclusive)
605        })
606    }
607
608    /// Find the best index for a set of columns.
609    pub fn find_index_for_columns(&self, columns: &[String]) -> Option<(String, IndexType)> {
610        // First try to find an exact match in hash indexes (faster for equality)
611        let hash = self.hash_indexes.read().expect("TableIndexManager hash_indexes lock poisoned");
612        for (name, index) in hash.iter() {
613            if index.columns == columns {
614                return Some((name.clone(), IndexType::Hash));
615            }
616        }
617
618        // Then try B-tree indexes (can use prefix of columns)
619        let btree = self.btree_indexes.read().expect("TableIndexManager btree_indexes lock poisoned");
620        for (name, index) in btree.iter() {
621            if columns.len() <= index.columns.len()
622                && index.columns[..columns.len()] == columns[..]
623            {
624                return Some((name.clone(), IndexType::BTree));
625            }
626        }
627
628        None
629    }
630
631    /// Get all index names.
632    pub fn index_names(&self) -> Vec<String> {
633        let btree = self.btree_indexes.read().expect("TableIndexManager btree_indexes lock poisoned");
634        let hash = self.hash_indexes.read().expect("TableIndexManager hash_indexes lock poisoned");
635        btree.keys().chain(hash.keys()).cloned().collect()
636    }
637
638    /// Get index info by name.
639    pub fn get_index_info(&self, name: &str) -> Option<(Vec<String>, bool, IndexType)> {
640        let btree = self.btree_indexes.read().expect("TableIndexManager btree_indexes lock poisoned");
641        if let Some(index) = btree.get(name) {
642            return Some((index.columns.clone(), index.unique, IndexType::BTree));
643        }
644
645        let hash = self.hash_indexes.read().expect("TableIndexManager hash_indexes lock poisoned");
646        if let Some(index) = hash.get(name) {
647            return Some((index.columns.clone(), index.unique, IndexType::Hash));
648        }
649
650        None
651    }
652
653    /// Insert a row into a specific index by name.
654    pub fn insert_into_index(
655        &self,
656        index_name: &str,
657        key: IndexKey,
658        row_id: RowId,
659    ) -> Result<(), IndexError> {
660        // Check B-tree indexes
661        let btree = self.btree_indexes.read().expect("TableIndexManager btree_indexes lock poisoned");
662        if let Some(index) = btree.get(index_name) {
663            return index.insert(key, row_id);
664        }
665        drop(btree);
666
667        // Check hash indexes
668        let hash = self.hash_indexes.read().expect("TableIndexManager hash_indexes lock poisoned");
669        if let Some(index) = hash.get(index_name) {
670            return index.insert(key, row_id);
671        }
672
673        Err(IndexError::IndexNotFound(index_name.to_string()))
674    }
675
676    /// Remove a row from a specific index by name.
677    pub fn remove_from_index(
678        &self,
679        index_name: &str,
680        key: &IndexKey,
681        row_id: RowId,
682    ) -> Result<bool, IndexError> {
683        // Check B-tree indexes
684        let btree = self.btree_indexes.read().expect("TableIndexManager btree_indexes lock poisoned");
685        if let Some(index) = btree.get(index_name) {
686            return Ok(index.remove(key, row_id));
687        }
688        drop(btree);
689
690        // Check hash indexes
691        let hash = self.hash_indexes.read().expect("TableIndexManager hash_indexes lock poisoned");
692        if let Some(index) = hash.get(index_name) {
693            return Ok(index.remove(key, row_id));
694        }
695
696        Err(IndexError::IndexNotFound(index_name.to_string()))
697    }
698
699    fn build_key(&self, columns: &[String], values: &HashMap<String, Value>) -> IndexKey {
700        let key_values: Vec<IndexValue> = columns
701            .iter()
702            .map(|col| {
703                values
704                    .get(col)
705                    .map(IndexKey::from_value)
706                    .unwrap_or(IndexValue::Null)
707            })
708            .collect();
709        IndexKey::new(key_values)
710    }
711}
712
713// =============================================================================
714// Tests
715// =============================================================================
716
717#[cfg(test)]
718mod tests {
719    use super::*;
720
721    #[test]
722    fn test_btree_index_basic() {
723        let index = BTreeIndex::new(
724            "idx_test".to_string(),
725            "users".to_string(),
726            vec!["id".to_string()],
727            false,
728        );
729
730        // Insert some keys
731        index.insert(IndexKey::single(IndexValue::Int(1)), 0).unwrap();
732        index.insert(IndexKey::single(IndexValue::Int(2)), 1).unwrap();
733        index.insert(IndexKey::single(IndexValue::Int(3)), 2).unwrap();
734
735        // Lookup
736        assert_eq!(index.get(&IndexKey::single(IndexValue::Int(2))), vec![1]);
737        assert_eq!(index.get(&IndexKey::single(IndexValue::Int(4))), Vec::<RowId>::new());
738    }
739
740    #[test]
741    fn test_btree_index_range() {
742        let index = BTreeIndex::new(
743            "idx_test".to_string(),
744            "users".to_string(),
745            vec!["age".to_string()],
746            false,
747        );
748
749        for i in 0..10 {
750            index.insert(IndexKey::single(IndexValue::Int(i * 10)), i as RowId).unwrap();
751        }
752
753        // Range [20, 50]
754        let result = index.range(
755            Some(&IndexKey::single(IndexValue::Int(20))),
756            Some(&IndexKey::single(IndexValue::Int(50))),
757            true,
758            true,
759        );
760        assert_eq!(result.len(), 4); // 20, 30, 40, 50
761
762        // Range (20, 50)
763        let result = index.range(
764            Some(&IndexKey::single(IndexValue::Int(20))),
765            Some(&IndexKey::single(IndexValue::Int(50))),
766            false,
767            false,
768        );
769        assert_eq!(result.len(), 2); // 30, 40
770    }
771
772    #[test]
773    fn test_btree_unique_constraint() {
774        let index = BTreeIndex::new(
775            "idx_unique".to_string(),
776            "users".to_string(),
777            vec!["email".to_string()],
778            true,
779        );
780
781        index.insert(IndexKey::single(IndexValue::String("a@b.com".to_string())), 0).unwrap();
782
783        let result = index.insert(IndexKey::single(IndexValue::String("a@b.com".to_string())), 1);
784        assert!(matches!(result, Err(IndexError::DuplicateKey(_))));
785    }
786
787    #[test]
788    fn test_hash_index_basic() {
789        let index = HashIndex::new(
790            "idx_hash".to_string(),
791            "users".to_string(),
792            vec!["status".to_string()],
793            false,
794        );
795
796        index.insert(IndexKey::single(IndexValue::String("active".to_string())), 0).unwrap();
797        index.insert(IndexKey::single(IndexValue::String("active".to_string())), 1).unwrap();
798        index.insert(IndexKey::single(IndexValue::String("inactive".to_string())), 2).unwrap();
799
800        let active = index.get(&IndexKey::single(IndexValue::String("active".to_string())));
801        assert_eq!(active.len(), 2);
802
803        let inactive = index.get(&IndexKey::single(IndexValue::String("inactive".to_string())));
804        assert_eq!(inactive.len(), 1);
805    }
806
807    #[test]
808    fn test_composite_key() {
809        let index = BTreeIndex::new(
810            "idx_composite".to_string(),
811            "orders".to_string(),
812            vec!["user_id".to_string(), "order_date".to_string()],
813            false,
814        );
815
816        let key1 = IndexKey::new(vec![
817            IndexValue::Int(1),
818            IndexValue::String("2024-01-01".to_string()),
819        ]);
820        let key2 = IndexKey::new(vec![
821            IndexValue::Int(1),
822            IndexValue::String("2024-01-02".to_string()),
823        ]);
824        let key3 = IndexKey::new(vec![
825            IndexValue::Int(2),
826            IndexValue::String("2024-01-01".to_string()),
827        ]);
828
829        index.insert(key1.clone(), 0).unwrap();
830        index.insert(key2.clone(), 1).unwrap();
831        index.insert(key3.clone(), 2).unwrap();
832
833        assert_eq!(index.get(&key1), vec![0]);
834        assert_eq!(index.get(&key2), vec![1]);
835        assert_eq!(index.len(), 3);
836    }
837
838    #[test]
839    fn test_table_index_manager() {
840        let manager = TableIndexManager::new("users".to_string());
841
842        // Create indexes
843        manager.create_index(
844            "idx_id".to_string(),
845            vec!["id".to_string()],
846            true,
847            IndexType::BTree,
848        ).unwrap();
849
850        manager.create_index(
851            "idx_status".to_string(),
852            vec!["status".to_string()],
853            false,
854            IndexType::Hash,
855        ).unwrap();
856
857        // Insert a row
858        let mut values = HashMap::new();
859        values.insert("id".to_string(), Value::Integer(1));
860        values.insert("status".to_string(), Value::String("active".to_string()));
861
862        manager.insert_row(0, &values).unwrap();
863
864        // Lookup
865        let key_id = IndexKey::single(IndexValue::Int(1));
866        let key_status = IndexKey::single(IndexValue::String("active".to_string()));
867
868        assert_eq!(manager.lookup_eq("idx_id", &key_id), Some(vec![0]));
869        assert_eq!(manager.lookup_eq("idx_status", &key_status), Some(vec![0]));
870    }
871
872    #[test]
873    fn test_find_index_for_columns() {
874        let manager = TableIndexManager::new("users".to_string());
875
876        manager.create_index(
877            "idx_name".to_string(),
878            vec!["name".to_string()],
879            false,
880            IndexType::Hash,
881        ).unwrap();
882
883        manager.create_index(
884            "idx_age_city".to_string(),
885            vec!["age".to_string(), "city".to_string()],
886            false,
887            IndexType::BTree,
888        ).unwrap();
889
890        // Exact match for hash index
891        let result = manager.find_index_for_columns(&["name".to_string()]);
892        assert_eq!(result, Some(("idx_name".to_string(), IndexType::Hash)));
893
894        // Prefix match for B-tree index
895        let result = manager.find_index_for_columns(&["age".to_string()]);
896        assert_eq!(result, Some(("idx_age_city".to_string(), IndexType::BTree)));
897
898        // No matching index
899        let result = manager.find_index_for_columns(&["email".to_string()]);
900        assert_eq!(result, None);
901    }
902
903    #[test]
904    fn test_index_remove() {
905        let index = BTreeIndex::new(
906            "idx_test".to_string(),
907            "users".to_string(),
908            vec!["id".to_string()],
909            false,
910        );
911
912        let key = IndexKey::single(IndexValue::Int(1));
913        index.insert(key.clone(), 0).unwrap();
914        index.insert(key.clone(), 1).unwrap();
915
916        assert_eq!(index.get(&key).len(), 2);
917
918        index.remove(&key, 0);
919        assert_eq!(index.get(&key).len(), 1);
920
921        index.remove(&key, 1);
922        assert_eq!(index.get(&key).len(), 0);
923    }
924}
aegis_query/index.rs

aegis_query/
index.rs
