Skip to main content

oxigdal_query/index/
mod.rs

1//! Index metadata and selection.
2
3pub mod selector;
4
5use crate::optimizer::cost_model::{IndexStatistics, IndexType};
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8
9/// Index metadata registry.
10pub struct IndexRegistry {
11    /// Indexes by table name.
12    indexes: HashMap<String, Vec<Index>>,
13}
14
15impl IndexRegistry {
16    /// Create a new index registry.
17    pub fn new() -> Self {
18        Self {
19            indexes: HashMap::new(),
20        }
21    }
22
23    /// Register an index.
24    pub fn register_index(&mut self, table: String, index: Index) {
25        self.indexes.entry(table).or_default().push(index);
26    }
27
28    /// Get indexes for a table.
29    pub fn get_indexes(&self, table: &str) -> Option<&[Index]> {
30        self.indexes.get(table).map(|v| v.as_slice())
31    }
32
33    /// Get index by name.
34    pub fn get_index(&self, table: &str, name: &str) -> Option<&Index> {
35        self.indexes.get(table)?.iter().find(|idx| idx.name == name)
36    }
37}
38
39impl Default for IndexRegistry {
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45/// Index metadata.
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct Index {
48    /// Index name.
49    pub name: String,
50    /// Table name.
51    pub table: String,
52    /// Indexed columns.
53    pub columns: Vec<String>,
54    /// Index type.
55    pub index_type: IndexType,
56    /// Index statistics.
57    pub statistics: IndexStatistics,
58    /// Whether the index is unique.
59    pub is_unique: bool,
60    /// Whether the index is primary key.
61    pub is_primary: bool,
62}
63
64impl Index {
65    /// Create a new index.
66    pub fn new(
67        name: String,
68        table: String,
69        columns: Vec<String>,
70        index_type: IndexType,
71        statistics: IndexStatistics,
72    ) -> Self {
73        Self {
74            name,
75            table,
76            columns,
77            index_type,
78            statistics,
79            is_unique: false,
80            is_primary: false,
81        }
82    }
83
84    /// Mark as unique index.
85    pub fn with_unique(mut self) -> Self {
86        self.is_unique = true;
87        self
88    }
89
90    /// Mark as primary key.
91    pub fn with_primary(mut self) -> Self {
92        self.is_primary = true;
93        self.is_unique = true;
94        self
95    }
96
97    /// Check if index covers columns.
98    pub fn covers_columns(&self, columns: &[String]) -> bool {
99        columns.iter().all(|col| self.columns.contains(col))
100    }
101
102    /// Check if index can be used for prefix match.
103    pub fn supports_prefix(&self, columns: &[String]) -> bool {
104        if columns.is_empty() {
105            return false;
106        }
107
108        for (i, col) in columns.iter().enumerate() {
109            if i >= self.columns.len() || &self.columns[i] != col {
110                return false;
111            }
112        }
113
114        true
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    #[test]
123    fn test_index_registry() {
124        let mut registry = IndexRegistry::new();
125
126        let idx_stats = IndexStatistics::new(
127            "idx_id".to_string(),
128            vec!["id".to_string()],
129            IndexType::BTree,
130            10000,
131        );
132
133        let index = Index::new(
134            "idx_id".to_string(),
135            "users".to_string(),
136            vec!["id".to_string()],
137            IndexType::BTree,
138            idx_stats,
139        );
140
141        registry.register_index("users".to_string(), index);
142
143        let indexes = registry.get_indexes("users");
144        assert!(indexes.is_some());
145        assert_eq!(indexes.as_ref().map(|i| i.len()), Some(1));
146    }
147
148    #[test]
149    fn test_index_covers_columns() {
150        let idx_stats = IndexStatistics::new(
151            "idx_name_age".to_string(),
152            vec!["name".to_string(), "age".to_string()],
153            IndexType::BTree,
154            10000,
155        );
156
157        let index = Index::new(
158            "idx_name_age".to_string(),
159            "users".to_string(),
160            vec!["name".to_string(), "age".to_string()],
161            IndexType::BTree,
162            idx_stats,
163        );
164
165        assert!(index.covers_columns(&["name".to_string(), "age".to_string()]));
166        assert!(index.covers_columns(&["name".to_string()]));
167        assert!(!index.covers_columns(&["email".to_string()]));
168    }
169
170    #[test]
171    fn test_index_supports_prefix() {
172        let idx_stats = IndexStatistics::new(
173            "idx_a_b_c".to_string(),
174            vec!["a".to_string(), "b".to_string(), "c".to_string()],
175            IndexType::BTree,
176            10000,
177        );
178
179        let index = Index::new(
180            "idx_a_b_c".to_string(),
181            "table".to_string(),
182            vec!["a".to_string(), "b".to_string(), "c".to_string()],
183            IndexType::BTree,
184            idx_stats,
185        );
186
187        assert!(index.supports_prefix(&["a".to_string()]));
188        assert!(index.supports_prefix(&["a".to_string(), "b".to_string()]));
189        assert!(!index.supports_prefix(&["b".to_string()]));
190        assert!(!index.supports_prefix(&["a".to_string(), "c".to_string()]));
191    }
192}