1use std::collections::HashMap;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum IndexType {
13 HashIndex,
15 BTreeIndex,
17 BitmapIndex,
19 CompressedBitmapIndex,
21}
22
23#[derive(Debug, Clone)]
27pub struct IndexStats {
28 pub index_type: IndexType,
30 pub entry_count: usize,
32 pub size_bytes: usize,
34 pub avg_lookup_ms: f64,
36 pub hit_rate: f64,
38}
39
40#[derive(Debug, Clone)]
44pub struct QueryPattern {
45 pub predicate: String,
47 pub is_equality: bool,
49 pub is_range: bool,
51 pub frequency: u64,
53 pub selectivity: f64,
55}
56
57#[derive(Debug, Clone)]
61pub struct IndexSelectionThresholds {
62 pub bitmap_max_cardinality: usize,
64 pub hash_min_frequency: u64,
66 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
80pub struct AdaptiveIndexSelector {
84 patterns: HashMap<String, QueryPattern>,
85 thresholds: IndexSelectionThresholds,
86}
87
88impl AdaptiveIndexSelector {
89 pub fn new() -> Self {
91 Self {
92 patterns: HashMap::new(),
93 thresholds: IndexSelectionThresholds::default(),
94 }
95 }
96
97 pub fn with_thresholds(thresholds: IndexSelectionThresholds) -> Self {
99 Self {
100 patterns: HashMap::new(),
101 thresholds,
102 }
103 }
104
105 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 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 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 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
181pub 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 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 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 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 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 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 self.warm_layer.insert(key, val);
255 evicted += 1;
256 }
257 }
258 evicted
259 }
260
261 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 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 pub fn hot_count(&self) -> usize {
285 self.hot_layer.len()
286 }
287
288 pub fn warm_count(&self) -> usize {
290 self.warm_layer.len()
291 }
292}
293
294fn 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#[cfg(test)]
310mod tests {
311 use super::*;
312
313 #[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 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); }
366 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); 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 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); assert_eq!(sel.recommend("http://p"), IndexType::HashIndex);
442 }
443
444 #[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]); 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()); ts.evict_hot_to_warm();
501 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(); 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 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}