Skip to main content

oxirs_tdb/index/
adaptive.rs

1//! Adaptive index strategies for OxiRS TDB.
2//!
3//! Provides smart index selection based on observed query patterns, and
4//! a tiered-storage abstraction with hot/warm/cold layers.
5
6use std::collections::HashMap;
7
8// ─── Index Types ──────────────────────────────────────────────────────────────
9
10/// Category of index structure
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum IndexType {
13    /// Hash-based index — O(1) equality lookups
14    HashIndex,
15    /// B-Tree index — O(log n), supports range scans
16    BTreeIndex,
17    /// Bitmap index — efficient for low-cardinality columns
18    BitmapIndex,
19    /// Compressed bitmap index — roaring-bitmap style, very low cardinality
20    CompressedBitmapIndex,
21}
22
23// ─── IndexStats ──────────────────────────────────────────────────────────────
24
25/// Runtime statistics for a single index instance
26#[derive(Debug, Clone)]
27pub struct IndexStats {
28    /// Type of this index
29    pub index_type: IndexType,
30    /// Number of distinct indexed entries
31    pub entry_count: usize,
32    /// Approximate memory/disk usage in bytes
33    pub size_bytes: usize,
34    /// Average lookup time in milliseconds
35    pub avg_lookup_ms: f64,
36    /// Cache / bloom-filter hit rate in [0, 1]
37    pub hit_rate: f64,
38}
39
40// ─── QueryPattern ────────────────────────────────────────────────────────────
41
42/// Observed access pattern for a particular predicate
43#[derive(Debug, Clone)]
44pub struct QueryPattern {
45    /// RDF predicate URI
46    pub predicate: String,
47    /// Whether equality lookups are used (e.g. `FILTER(?x = <iri>)`)
48    pub is_equality: bool,
49    /// Whether range queries are used (e.g. FILTER(?x > 5))
50    pub is_range: bool,
51    /// Number of times this predicate has been queried
52    pub frequency: u64,
53    /// Fraction of triples selected per query (lower = more selective)
54    pub selectivity: f64,
55}
56
57// ─── Thresholds ──────────────────────────────────────────────────────────────
58
59/// Thresholds that govern which index type is selected
60#[derive(Debug, Clone)]
61pub struct IndexSelectionThresholds {
62    /// Maximum cardinality to use bitmap index (e.g. boolean columns)
63    pub bitmap_max_cardinality: usize,
64    /// Minimum query frequency to prefer hash index for equality
65    pub hash_min_frequency: u64,
66    /// Whether to use B-Tree for range queries
67    pub btree_for_ranges: bool,
68}
69
70impl Default for IndexSelectionThresholds {
71    fn default() -> Self {
72        Self {
73            bitmap_max_cardinality: 2,
74            hash_min_frequency: 10,
75            btree_for_ranges: true,
76        }
77    }
78}
79
80// ─── AdaptiveIndexSelector ───────────────────────────────────────────────────
81
82/// Tracks query patterns and recommends the optimal index type per predicate.
83pub struct AdaptiveIndexSelector {
84    patterns: HashMap<String, QueryPattern>,
85    thresholds: IndexSelectionThresholds,
86}
87
88impl AdaptiveIndexSelector {
89    /// Create a new selector with default thresholds.
90    pub fn new() -> Self {
91        Self {
92            patterns: HashMap::new(),
93            thresholds: IndexSelectionThresholds::default(),
94        }
95    }
96
97    /// Create a selector with custom thresholds.
98    pub fn with_thresholds(thresholds: IndexSelectionThresholds) -> Self {
99        Self {
100            patterns: HashMap::new(),
101            thresholds,
102        }
103    }
104
105    /// Record a query observation for `predicate`.
106    pub fn observe_query(&mut self, predicate: &str, is_equality: bool, is_range: bool) {
107        let entry = self
108            .patterns
109            .entry(predicate.to_string())
110            .or_insert_with(|| QueryPattern {
111                predicate: predicate.to_string(),
112                is_equality: false,
113                is_range: false,
114                frequency: 0,
115                selectivity: 1.0,
116            });
117        entry.frequency += 1;
118        if is_equality {
119            entry.is_equality = true;
120        }
121        if is_range {
122            entry.is_range = true;
123        }
124    }
125
126    /// Recommend an index type for `predicate` based on observed patterns.
127    ///
128    /// Decision rules (evaluated in order):
129    /// 1. If range queries observed and `btree_for_ranges` → [`IndexType::BTreeIndex`]
130    /// 2. If equality queries and frequency ≥ `hash_min_frequency` → [`IndexType::HashIndex`]
131    /// 3. If cardinality (approximated by frequency) ≤ `bitmap_max_cardinality` → [`IndexType::CompressedBitmapIndex`]
132    /// 4. Default → [`IndexType::BTreeIndex`]
133    pub fn recommend(&self, predicate: &str) -> IndexType {
134        match self.patterns.get(predicate) {
135            None => IndexType::BTreeIndex,
136            Some(p) => {
137                if p.is_range && self.thresholds.btree_for_ranges {
138                    return IndexType::BTreeIndex;
139                }
140                if p.is_equality && p.frequency >= self.thresholds.hash_min_frequency {
141                    return IndexType::HashIndex;
142                }
143                if p.frequency <= self.thresholds.bitmap_max_cardinality as u64 {
144                    return IndexType::CompressedBitmapIndex;
145                }
146                IndexType::BTreeIndex
147            }
148        }
149    }
150
151    /// Return recommendations for every tracked predicate.
152    pub fn all_recommendations(&self) -> Vec<(String, IndexType)> {
153        let mut recs: Vec<_> = self
154            .patterns
155            .keys()
156            .map(|p| (p.clone(), self.recommend(p)))
157            .collect();
158        recs.sort_by(|a, b| a.0.cmp(&b.0));
159        recs
160    }
161
162    /// Return the top-`n` predicates by query frequency, highest first.
163    pub fn top_predicates_by_frequency(&self, n: usize) -> Vec<(&str, u64)> {
164        let mut entries: Vec<(&str, u64)> = self
165            .patterns
166            .iter()
167            .map(|(k, v)| (k.as_str(), v.frequency))
168            .collect();
169        entries.sort_by_key(|b| std::cmp::Reverse(b.1));
170        entries.truncate(n);
171        entries
172    }
173}
174
175impl Default for AdaptiveIndexSelector {
176    fn default() -> Self {
177        Self::new()
178    }
179}
180
181// ─── TieredStorage ───────────────────────────────────────────────────────────
182
183/// A simple three-tier storage abstraction: hot (in-memory), warm (in-memory,
184/// compressed), and cold (disk-backed, path only tracked here).
185///
186/// Eviction policy: when a layer exceeds its threshold, the oldest-inserted
187/// entries (by insertion order) are evicted to the next layer.
188pub struct TieredStorage {
189    hot_layer: HashMap<String, Vec<u8>>,
190    hot_order: Vec<String>,
191    warm_layer: HashMap<String, Vec<u8>>,
192    warm_order: Vec<String>,
193    cold_layer_path: Option<std::path::PathBuf>,
194    hot_threshold: usize,
195    warm_threshold: usize,
196}
197
198impl TieredStorage {
199    /// Create a new tiered storage without a cold-layer path.
200    pub fn new(hot_threshold: usize, warm_threshold: usize) -> Self {
201        Self {
202            hot_layer: HashMap::new(),
203            hot_order: Vec::new(),
204            warm_layer: HashMap::new(),
205            warm_order: Vec::new(),
206            cold_layer_path: None,
207            hot_threshold,
208            warm_threshold,
209        }
210    }
211
212    /// Create a new tiered storage with a cold-layer path.
213    pub fn with_cold_path(
214        hot_threshold: usize,
215        warm_threshold: usize,
216        cold_path: std::path::PathBuf,
217    ) -> Self {
218        let mut s = Self::new(hot_threshold, warm_threshold);
219        s.cold_layer_path = Some(cold_path);
220        s
221    }
222
223    /// Retrieve a value, checking hot then warm layers.
224    pub fn get(&self, key: &str) -> Option<Vec<u8>> {
225        if let Some(v) = self.hot_layer.get(key) {
226            return Some(v.clone());
227        }
228        self.warm_layer.get(key).cloned()
229    }
230
231    /// Insert or update a value. The entry always lands in the hot layer.
232    pub fn put(&mut self, key: &str, value: Vec<u8>) {
233        if !self.hot_layer.contains_key(key) {
234            self.hot_order.push(key.to_string());
235        }
236        self.hot_layer.insert(key.to_string(), value);
237    }
238
239    /// Evict entries from the hot layer to the warm layer when hot exceeds
240    /// its threshold.  Returns the number of entries evicted.
241    pub fn evict_hot_to_warm(&mut self) -> usize {
242        let mut evicted = 0;
243        while self.hot_layer.len() > self.hot_threshold {
244            if self.hot_order.is_empty() {
245                break;
246            }
247            let key = self.hot_order.remove(0);
248            if let Some(val) = self.hot_layer.remove(&key) {
249                if !self.warm_layer.contains_key(&key) {
250                    self.warm_order.push(key.clone());
251                }
252                // In a real implementation we would compress here; for the
253                // purposes of this abstraction we store as-is.
254                self.warm_layer.insert(key, val);
255                evicted += 1;
256            }
257        }
258        evicted
259    }
260
261    /// Evict entries from the warm layer to the cold layer (disk) when warm
262    /// exceeds its threshold.  Returns the number of entries evicted.
263    pub fn evict_warm_to_cold(&mut self) -> usize {
264        let mut evicted = 0;
265        while self.warm_layer.len() > self.warm_threshold {
266            if self.warm_order.is_empty() {
267                break;
268            }
269            let key = self.warm_order.remove(0);
270            if let Some(val) = self.warm_layer.remove(&key) {
271                // Write to disk if a cold-layer path is configured.
272                if let Some(ref cold_path) = self.cold_layer_path {
273                    let file_path = cold_path.join(sanitize_key(&key));
274                    let _ = std::fs::create_dir_all(cold_path);
275                    let _ = std::fs::write(file_path, &val);
276                }
277                evicted += 1;
278            }
279        }
280        evicted
281    }
282
283    /// Number of entries currently in the hot layer.
284    pub fn hot_count(&self) -> usize {
285        self.hot_layer.len()
286    }
287
288    /// Number of entries currently in the warm layer.
289    pub fn warm_count(&self) -> usize {
290        self.warm_layer.len()
291    }
292}
293
294/// Convert a key string to a safe filename component.
295fn sanitize_key(key: &str) -> String {
296    key.chars()
297        .map(|c| {
298            if c.is_alphanumeric() || c == '_' || c == '-' {
299                c
300            } else {
301                '_'
302            }
303        })
304        .collect()
305}
306
307// ─── Tests ───────────────────────────────────────────────────────────────────
308
309#[cfg(test)]
310mod tests {
311    use super::*;
312
313    // ── AdaptiveIndexSelector ────────────────────────────────────────────────
314
315    #[test]
316    fn test_selector_new() {
317        let sel = AdaptiveIndexSelector::new();
318        assert!(sel.patterns.is_empty());
319    }
320
321    #[test]
322    fn test_recommend_unknown_predicate_is_btree() {
323        let sel = AdaptiveIndexSelector::new();
324        assert_eq!(sel.recommend("http://unknown"), IndexType::BTreeIndex);
325    }
326
327    #[test]
328    fn test_observe_query_increments_frequency() {
329        let mut sel = AdaptiveIndexSelector::new();
330        sel.observe_query("http://p", true, false);
331        sel.observe_query("http://p", true, false);
332        assert_eq!(sel.patterns["http://p"].frequency, 2);
333    }
334
335    #[test]
336    fn test_observe_query_sets_equality_flag() {
337        let mut sel = AdaptiveIndexSelector::new();
338        sel.observe_query("http://p", true, false);
339        assert!(sel.patterns["http://p"].is_equality);
340        assert!(!sel.patterns["http://p"].is_range);
341    }
342
343    #[test]
344    fn test_observe_query_sets_range_flag() {
345        let mut sel = AdaptiveIndexSelector::new();
346        sel.observe_query("http://p", false, true);
347        assert!(sel.patterns["http://p"].is_range);
348    }
349
350    #[test]
351    fn test_recommend_hash_for_frequent_equality() {
352        let mut sel = AdaptiveIndexSelector::new();
353        for _ in 0..15 {
354            sel.observe_query("http://p", true, false);
355        }
356        // frequency=15 >= hash_min_frequency=10, equality=true, no range
357        assert_eq!(sel.recommend("http://p"), IndexType::HashIndex);
358    }
359
360    #[test]
361    fn test_recommend_btree_for_range_overrides_hash() {
362        let mut sel = AdaptiveIndexSelector::new();
363        for _ in 0..20 {
364            sel.observe_query("http://p", true, true); // both equality and range
365        }
366        // Range takes priority when btree_for_ranges = true
367        assert_eq!(sel.recommend("http://p"), IndexType::BTreeIndex);
368    }
369
370    #[test]
371    fn test_recommend_compressed_bitmap_low_frequency() {
372        let mut sel = AdaptiveIndexSelector::new();
373        sel.observe_query("http://p", false, false); // frequency=1, no range, no equality
374        assert_eq!(sel.recommend("http://p"), IndexType::CompressedBitmapIndex);
375    }
376
377    #[test]
378    fn test_recommend_btree_default_medium_frequency_no_flags() {
379        let mut sel = AdaptiveIndexSelector::new();
380        for _ in 0..5 {
381            sel.observe_query("http://p", false, false);
382        }
383        // frequency=5 > bitmap_max_cardinality=2, not equality-frequent, no range
384        assert_eq!(sel.recommend("http://p"), IndexType::BTreeIndex);
385    }
386
387    #[test]
388    fn test_all_recommendations_returns_all() {
389        let mut sel = AdaptiveIndexSelector::new();
390        sel.observe_query("http://a", true, false);
391        sel.observe_query("http://b", false, true);
392        let recs = sel.all_recommendations();
393        assert_eq!(recs.len(), 2);
394    }
395
396    #[test]
397    fn test_all_recommendations_sorted() {
398        let mut sel = AdaptiveIndexSelector::new();
399        sel.observe_query("http://z", true, false);
400        sel.observe_query("http://a", true, false);
401        sel.observe_query("http://m", true, false);
402        let recs = sel.all_recommendations();
403        let keys: Vec<_> = recs.iter().map(|(k, _)| k.as_str()).collect();
404        assert!(keys.windows(2).all(|w| w[0] <= w[1]));
405    }
406
407    #[test]
408    fn test_top_predicates_by_frequency() {
409        let mut sel = AdaptiveIndexSelector::new();
410        for _ in 0..10 {
411            sel.observe_query("http://frequent", true, false);
412        }
413        for _ in 0..2 {
414            sel.observe_query("http://rare", false, false);
415        }
416        for _ in 0..5 {
417            sel.observe_query("http://medium", false, false);
418        }
419        let top = sel.top_predicates_by_frequency(2);
420        assert_eq!(top.len(), 2);
421        assert_eq!(top[0].0, "http://frequent");
422        assert_eq!(top[0].1, 10);
423    }
424
425    #[test]
426    fn test_top_predicates_by_frequency_empty() {
427        let sel = AdaptiveIndexSelector::new();
428        assert!(sel.top_predicates_by_frequency(5).is_empty());
429    }
430
431    #[test]
432    fn test_with_thresholds() {
433        let thresh = IndexSelectionThresholds {
434            bitmap_max_cardinality: 100,
435            hash_min_frequency: 1,
436            btree_for_ranges: false,
437        };
438        let mut sel = AdaptiveIndexSelector::with_thresholds(thresh);
439        sel.observe_query("http://p", true, true); // range, but btree_for_ranges=false
440                                                   // Now: not range-preferred, equality=true, frequency=1 >= 1 → HashIndex
441        assert_eq!(sel.recommend("http://p"), IndexType::HashIndex);
442    }
443
444    // ── TieredStorage ────────────────────────────────────────────────────────
445
446    #[test]
447    fn test_tiered_storage_new() {
448        let ts = TieredStorage::new(10, 5);
449        assert_eq!(ts.hot_count(), 0);
450        assert_eq!(ts.warm_count(), 0);
451    }
452
453    #[test]
454    fn test_put_and_get_from_hot() {
455        let mut ts = TieredStorage::new(10, 5);
456        ts.put("key1", b"value1".to_vec());
457        let got = ts.get("key1").unwrap();
458        assert_eq!(got, b"value1");
459    }
460
461    #[test]
462    fn test_get_missing_returns_none() {
463        let ts = TieredStorage::new(10, 5);
464        assert!(ts.get("missing").is_none());
465    }
466
467    #[test]
468    fn test_hot_count_after_puts() {
469        let mut ts = TieredStorage::new(10, 5);
470        ts.put("k1", vec![1]);
471        ts.put("k2", vec![2]);
472        assert_eq!(ts.hot_count(), 2);
473    }
474
475    #[test]
476    fn test_evict_hot_to_warm() {
477        let mut ts = TieredStorage::new(2, 10);
478        ts.put("k1", vec![1]);
479        ts.put("k2", vec![2]);
480        ts.put("k3", vec![3]); // hot_count = 3 > threshold 2
481        let evicted = ts.evict_hot_to_warm();
482        assert!(evicted > 0);
483        assert!(ts.hot_count() <= 2);
484    }
485
486    #[test]
487    fn test_evict_hot_to_warm_no_eviction_needed() {
488        let mut ts = TieredStorage::new(10, 5);
489        ts.put("k1", vec![1]);
490        let evicted = ts.evict_hot_to_warm();
491        assert_eq!(evicted, 0);
492        assert_eq!(ts.hot_count(), 1);
493    }
494
495    #[test]
496    fn test_evicted_entries_accessible_from_warm() {
497        let mut ts = TieredStorage::new(1, 10);
498        ts.put("k1", b"first".to_vec());
499        ts.put("k2", b"second".to_vec()); // hot_count = 2 > 1
500        ts.evict_hot_to_warm();
501        // k1 should have been evicted to warm
502        assert!(ts.warm_count() > 0);
503        assert!(ts.get("k1").is_some() || ts.get("k2").is_some());
504    }
505
506    #[test]
507    fn test_evict_warm_to_cold_no_path() {
508        let mut ts = TieredStorage::new(1, 2);
509        ts.put("k1", vec![1]);
510        ts.put("k2", vec![2]);
511        ts.evict_hot_to_warm();
512        ts.put("k3", vec![3]);
513        ts.put("k4", vec![4]);
514        ts.evict_hot_to_warm(); // warm now has entries
515                                // Manually push warm over limit
516        ts.warm_layer.insert("extra1".to_string(), vec![9]);
517        ts.warm_layer.insert("extra2".to_string(), vec![9]);
518        ts.warm_order.push("extra1".to_string());
519        ts.warm_order.push("extra2".to_string());
520        let evicted = ts.evict_warm_to_cold();
521        // Should have evicted some warm entries
522        assert!(evicted > 0 || ts.warm_count() <= 2);
523    }
524
525    #[test]
526    fn test_evict_warm_to_cold_with_path() {
527        let cold_dir = std::env::temp_dir().join("oxirs_tdb_tiered_cold_test");
528        std::fs::remove_dir_all(&cold_dir).ok();
529        let mut ts = TieredStorage::with_cold_path(1, 1, cold_dir.clone());
530        ts.put("k1", b"v1".to_vec());
531        ts.put("k2", b"v2".to_vec());
532        ts.evict_hot_to_warm();
533        ts.warm_layer.insert("w1".to_string(), b"warm1".to_vec());
534        ts.warm_layer.insert("w2".to_string(), b"warm2".to_vec());
535        ts.warm_order.push("w1".to_string());
536        ts.warm_order.push("w2".to_string());
537        ts.evict_warm_to_cold();
538        std::fs::remove_dir_all(&cold_dir).ok();
539    }
540}