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

aegis_query/
index.rs
