aurora_db/
index.rs

1use crate::error::Result;
2use crate::search::FullTextIndex;
3use crate::types::{Document, Value};
4use dashmap::DashMap;
5use serde::{Deserialize, Serialize};
6use std::sync::Arc;
7
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub enum IndexType {
10    BTree,
11    Hash,
12    FullText,
13    Custom(String),
14}
15
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct IndexDefinition {
18    pub name: String,
19    pub collection: String,
20    pub fields: Vec<String>,
21    pub index_type: IndexType,
22    pub unique: bool,
23}
24
25pub struct Index {
26    definition: IndexDefinition,
27    data: Arc<DashMap<Value, Vec<String>>>, // Value -> Document IDs
28    full_text: Option<Arc<FullTextIndex>>,
29}
30
31impl Index {
32    pub fn new(definition: IndexDefinition) -> Self {
33        let full_text = if matches!(definition.index_type, IndexType::FullText) {
34            Some(Arc::new(FullTextIndex::new(
35                &definition.collection,
36                &definition.fields[0],
37            )))
38        } else {
39            None
40        };
41
42        Self {
43            definition,
44            data: Arc::new(DashMap::new()),
45            full_text,
46        }
47    }
48
49    pub fn insert(&self, doc: &Document) -> Result<()> {
50        let key = self.extract_key(doc)?;
51        let doc_id = doc.id.clone();
52
53        match self.data.get_mut(&key) {
54            Some(mut ids) => {
55                if self.definition.unique && !ids.is_empty() {
56                    return Err(crate::error::AqlError::invalid_operation(
57                        "Unique constraint violation".to_string(),
58                    ));
59                }
60                ids.push(doc_id.clone());
61            }
62            None => {
63                self.data.insert(key, vec![doc_id.clone()]);
64            }
65        }
66
67        if let Some(ft_index) = &self.full_text {
68            ft_index.index_document(doc)?;
69        }
70
71        Ok(())
72    }
73
74    pub fn search(&self, value: &Value) -> Option<Vec<String>> {
75        self.data.get(value).map(|ids| ids.clone())
76    }
77
78    pub fn search_text(&self, query: &str) -> Option<Vec<(String, f32)>> {
79        self.full_text.as_ref().map(|ft| ft.search(query))
80    }
81
82    fn extract_key(&self, doc: &Document) -> Result<Value> {
83        if self.definition.fields.len() == 1 {
84            Ok(doc
85                .data
86                .get(&self.definition.fields[0])
87                .cloned()
88                .unwrap_or(Value::Null))
89        } else {
90            let values: Vec<Value> = self
91                .definition
92                .fields
93                .iter()
94                .map(|f| doc.data.get(f).cloned().unwrap_or(Value::Null))
95                .collect();
96            Ok(Value::Array(values))
97        }
98    }
99
100    #[allow(dead_code)]
101    pub fn full_text(&self) -> Option<Arc<FullTextIndex>> {
102        self.full_text.clone()
103    }
104}
aurora_db/index.rs

aurora_db/
index.rs
