Skip to main content

vibesql_executor/cache/
query_result_cache.rs

1//! Thread-safe query result cache with LRU eviction
2//!
3//! Caches actual query results (rows + schema) to avoid re-executing
4//! identical read-only queries. This is particularly effective for
5//! SQLLogicTest workloads with repeated query patterns.
6
7use std::{
8    collections::{HashMap, HashSet},
9    sync::{
10        atomic::{AtomicUsize, Ordering},
11        RwLock,
12    },
13};
14
15use vibesql_storage::Row;
16
17use super::{CacheStats, QuerySignature};
18use crate::schema::CombinedSchema;
19
20/// Cached query result with metadata
21#[derive(Clone)]
22struct CachedResult {
23    /// Query result rows
24    rows: Vec<Row>,
25    /// Schema for the result
26    schema: CombinedSchema,
27    /// Tables accessed by this query (for invalidation)
28    tables: HashSet<String>,
29}
30
31/// Thread-safe cache for query results
32///
33/// Caches the actual result rows and schema from SELECT queries.
34/// Results are invalidated when any referenced table is modified.
35pub struct QueryResultCache {
36    cache: RwLock<HashMap<QuerySignature, CachedResult>>,
37    max_size: usize,
38    hits: AtomicUsize,
39    misses: AtomicUsize,
40    evictions: AtomicUsize,
41}
42
43impl QueryResultCache {
44    /// Create a new result cache with specified max size
45    pub fn new(max_size: usize) -> Self {
46        Self {
47            cache: RwLock::new(HashMap::new()),
48            max_size,
49            hits: AtomicUsize::new(0),
50            misses: AtomicUsize::new(0),
51            evictions: AtomicUsize::new(0),
52        }
53    }
54
55    /// Try to get cached result for a query
56    pub fn get(&self, signature: &QuerySignature) -> Option<(Vec<Row>, CombinedSchema)> {
57        let cache = self.cache.read().unwrap();
58        if let Some(entry) = cache.get(signature) {
59            self.hits.fetch_add(1, Ordering::Relaxed);
60            Some((entry.rows.clone(), entry.schema.clone()))
61        } else {
62            self.misses.fetch_add(1, Ordering::Relaxed);
63            None
64        }
65    }
66
67    /// Insert query result into cache with table dependencies
68    pub fn insert(
69        &self,
70        signature: QuerySignature,
71        rows: Vec<Row>,
72        schema: CombinedSchema,
73        tables: HashSet<String>,
74    ) {
75        let entry = CachedResult { rows, schema, tables };
76        let mut cache = self.cache.write().unwrap();
77
78        // Simple LRU: evict first entry if at capacity
79        if cache.len() >= self.max_size {
80            if let Some(key) = cache.keys().next().cloned() {
81                cache.remove(&key);
82                self.evictions.fetch_add(1, Ordering::Relaxed);
83            }
84        }
85
86        cache.insert(signature, entry);
87    }
88
89    /// Check if signature is cached
90    pub fn contains(&self, signature: &QuerySignature) -> bool {
91        self.cache.read().unwrap().contains_key(signature)
92    }
93
94    /// Clear all cached results
95    pub fn clear(&self) {
96        self.cache.write().unwrap().clear();
97    }
98
99    /// Invalidate all queries touching a specific table
100    ///
101    /// This should be called when a table is modified (INSERT/UPDATE/DELETE)
102    pub fn invalidate_table(&self, table: &str) {
103        let mut cache = self.cache.write().unwrap();
104        cache.retain(|_, entry| !entry.tables.iter().any(|t| t.eq_ignore_ascii_case(table)));
105    }
106
107    /// Invalidate entire cache (e.g., on any write when not tracking dependencies)
108    pub fn invalidate_all(&self) {
109        self.clear();
110    }
111
112    /// Get cache statistics
113    pub fn stats(&self) -> CacheStats {
114        let cache = self.cache.read().unwrap();
115        let hits = self.hits.load(Ordering::Relaxed);
116        let misses = self.misses.load(Ordering::Relaxed);
117        let total = hits + misses;
118        let hit_rate = if total > 0 { hits as f64 / total as f64 } else { 0.0 };
119
120        CacheStats {
121            hits,
122            misses,
123            evictions: self.evictions.load(Ordering::Relaxed),
124            size: cache.len(),
125            hit_rate,
126        }
127    }
128
129    /// Get maximum cache size
130    pub fn max_size(&self) -> usize {
131        self.max_size
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use vibesql_types::SqlValue;
138
139    use super::*;
140
141    fn make_test_row(values: Vec<SqlValue>) -> Row {
142        Row::new(values)
143    }
144
145    fn make_test_schema() -> CombinedSchema {
146        use vibesql_catalog::{ColumnSchema, TableSchema};
147        use vibesql_types::DataType;
148
149        let columns = vec![
150            ColumnSchema {
151                name: "id".to_string(),
152                data_type: DataType::Integer,
153                nullable: false,
154                default_value: None,
155                generated_expr: None,
156                is_exact_integer_type: false,
157                collation: None,
158            },
159            ColumnSchema {
160                name: "name".to_string(),
161                data_type: DataType::Varchar { max_length: Some(255) },
162                nullable: true,
163                default_value: None,
164                generated_expr: None,
165                is_exact_integer_type: false,
166                collation: None,
167            },
168        ];
169
170        let schema = TableSchema::new("users".to_string(), columns);
171        CombinedSchema::from_table("users".to_string(), schema)
172    }
173
174    #[test]
175    fn test_cache_hit() {
176        let cache = QueryResultCache::new(10);
177        let sig = QuerySignature::from_sql("SELECT * FROM users");
178        let rows = vec![
179            make_test_row(vec![
180                SqlValue::Integer(1),
181                SqlValue::Varchar(arcstr::ArcStr::from("Alice")),
182            ]),
183            make_test_row(vec![
184                SqlValue::Integer(2),
185                SqlValue::Varchar(arcstr::ArcStr::from("Bob")),
186            ]),
187        ];
188        let schema = make_test_schema();
189        let mut tables = HashSet::new();
190        tables.insert("users".to_string());
191
192        cache.insert(sig.clone(), rows.clone(), schema.clone(), tables);
193        let result = cache.get(&sig);
194
195        assert!(result.is_some());
196        let (cached_rows, _cached_schema) = result.unwrap();
197        assert_eq!(cached_rows.len(), 2);
198        assert_eq!(cached_rows[0].values.len(), 2);
199    }
200
201    #[test]
202    fn test_cache_miss() {
203        let cache = QueryResultCache::new(10);
204        let sig = QuerySignature::from_sql("SELECT * FROM users");
205
206        let result = cache.get(&sig);
207        assert!(result.is_none());
208    }
209
210    #[test]
211    fn test_lru_eviction() {
212        let cache = QueryResultCache::new(2);
213        let schema = make_test_schema();
214        let rows = vec![make_test_row(vec![SqlValue::Integer(1)])];
215        let tables = HashSet::new();
216
217        let sig1 = QuerySignature::from_sql("SELECT * FROM users");
218        let sig2 = QuerySignature::from_sql("SELECT * FROM orders");
219        let sig3 = QuerySignature::from_sql("SELECT * FROM products");
220
221        cache.insert(sig1, rows.clone(), schema.clone(), tables.clone());
222        cache.insert(sig2, rows.clone(), schema.clone(), tables.clone());
223        assert_eq!(cache.stats().size, 2);
224
225        cache.insert(sig3, rows, schema, tables);
226        assert_eq!(cache.stats().size, 2);
227        assert_eq!(cache.stats().evictions, 1);
228    }
229
230    #[test]
231    fn test_cache_clear() {
232        let cache = QueryResultCache::new(10);
233        let sig = QuerySignature::from_sql("SELECT * FROM users");
234        let rows = vec![make_test_row(vec![SqlValue::Integer(1)])];
235        let schema = make_test_schema();
236        let tables = HashSet::new();
237
238        cache.insert(sig.clone(), rows, schema, tables);
239        assert!(cache.contains(&sig));
240
241        cache.clear();
242        assert!(!cache.contains(&sig));
243    }
244
245    #[test]
246    fn test_table_invalidation() {
247        let cache = QueryResultCache::new(10);
248        let sig = QuerySignature::from_sql("SELECT * FROM users WHERE id = 1");
249        let rows = vec![make_test_row(vec![SqlValue::Integer(1)])];
250        let schema = make_test_schema();
251        let mut tables = HashSet::new();
252        tables.insert("users".to_string());
253
254        cache.insert(sig.clone(), rows, schema, tables);
255        assert!(cache.contains(&sig));
256
257        cache.invalidate_table("users");
258        assert!(!cache.contains(&sig));
259    }
260
261    #[test]
262    fn test_table_invalidation_case_insensitive() {
263        let cache = QueryResultCache::new(10);
264        let sig = QuerySignature::from_sql("SELECT * FROM users");
265        let rows = vec![make_test_row(vec![SqlValue::Integer(1)])];
266        let schema = make_test_schema();
267        let mut tables = HashSet::new();
268        tables.insert("users".to_string());
269
270        cache.insert(sig.clone(), rows, schema, tables);
271        assert!(cache.contains(&sig));
272
273        // Invalidate with different case
274        cache.invalidate_table("USERS");
275        assert!(!cache.contains(&sig));
276    }
277
278    #[test]
279    fn test_cache_stats() {
280        let cache = QueryResultCache::new(10);
281        let sig = QuerySignature::from_sql("SELECT * FROM users");
282        let rows = vec![make_test_row(vec![SqlValue::Integer(1)])];
283        let schema = make_test_schema();
284        let tables = HashSet::new();
285
286        cache.insert(sig.clone(), rows, schema, tables);
287
288        // Generate hits
289        cache.get(&sig);
290        cache.get(&sig);
291
292        // Generate miss
293        let other_sig = QuerySignature::from_sql("SELECT * FROM orders");
294        cache.get(&other_sig);
295
296        let stats = cache.stats();
297        assert_eq!(stats.hits, 2);
298        assert_eq!(stats.misses, 1);
299        assert!((stats.hit_rate - 2.0 / 3.0).abs() < 0.01);
300    }
301}