xlog_runtime/executor/
join_cache.rs

1use std::collections::HashMap;
2use xlog_core::{RelId, ScalarType, Schema};
3use xlog_cuda::{CudaBuffer, JoinIndexV2};
4
5#[derive(Debug, Clone, PartialEq, Eq, Hash)]
6pub(crate) struct JoinIndexKey {
7    pub(crate) rel: RelId,
8    pub(crate) version: u64,
9    pub(crate) key_cols: Vec<usize>,
10    pub(crate) schema: JoinIndexSchemaSignature,
11    pub(crate) device_ordinal: u32,
12}
13
14#[derive(Debug, Clone, PartialEq, Eq, Hash)]
15pub(crate) struct JoinIndexSchemaSignature {
16    column_types: Vec<ScalarType>,
17    row_size_bytes: usize,
18}
19
20impl JoinIndexSchemaSignature {
21    fn from_schema(schema: &Schema) -> Self {
22        Self {
23            column_types: (0..schema.arity())
24                .filter_map(|idx| schema.column_type(idx))
25                .collect(),
26            row_size_bytes: schema.row_size_bytes(),
27        }
28    }
29}
30
31impl JoinIndexKey {
32    pub(crate) fn new(
33        rel: RelId,
34        version: u64,
35        key_cols: Vec<usize>,
36        schema: &Schema,
37        device_ordinal: u32,
38    ) -> Self {
39        Self {
40            rel,
41            version,
42            key_cols,
43            schema: JoinIndexSchemaSignature::from_schema(schema),
44            device_ordinal,
45        }
46    }
47}
48
49struct CachedJoinIndex {
50    index: Option<JoinIndexV2>,
51    bytes: u64,
52    last_used: u64,
53}
54
55/// Persistent join-index manager telemetry.
56#[derive(Clone, Debug, Default, PartialEq, Eq)]
57pub struct JoinIndexCacheStats {
58    /// Lookup attempts.
59    pub lookups: u64,
60    /// Successful index reuses.
61    pub hits: u64,
62    /// Lookup misses.
63    pub misses: u64,
64    /// Successful index builds inserted into the cache.
65    pub builds: u64,
66    /// LRU/budget evictions.
67    pub evictions: u64,
68    /// Entries invalidated because a relation changed.
69    pub invalidations: u64,
70    /// Stale entries rejected by provider validation.
71    pub stale_rejections: u64,
72    /// Background-build mode requests.
73    pub background_build_requests: u64,
74    /// Background-build mode completions.
75    pub background_builds_completed: u64,
76    /// Background builds whose indexed reuse was deferred until a later evaluation.
77    pub background_builds_deferred: u64,
78    /// Current retained index count.
79    pub entries: usize,
80    /// Current retained index bytes.
81    pub total_bytes: u64,
82}
83
84pub(crate) struct JoinIndexCache {
85    entries: HashMap<JoinIndexKey, CachedJoinIndex>,
86    clock: u64,
87    total_bytes: u64,
88    pub(crate) max_bytes: u64,
89    stats: JoinIndexCacheStats,
90}
91
92/// Estimate the GPU memory footprint of a join index built on `right` with `right_keys`.
93///
94/// Returns u64::MAX if keys are empty or column types are missing (signals "don't build").
95pub(crate) fn estimate_join_index_bytes(right: &CudaBuffer, right_keys: &[usize]) -> u64 {
96    if right_keys.is_empty() {
97        return u64::MAX;
98    }
99
100    let mut key_bytes_per_row: u64 = 0;
101    for &k in right_keys {
102        let Some(ty) = right.schema().column_type(k) else {
103            return u64::MAX;
104        };
105        key_bytes_per_row = key_bytes_per_row.saturating_add(ty.size_bytes() as u64);
106    }
107
108    let num_rows = right.num_rows();
109    let packed_bytes = num_rows.saturating_mul(key_bytes_per_row);
110    let target = num_rows.saturating_mul(2).max(1024);
111    let num_buckets = target.next_power_of_two();
112
113    // Stored index bytes: packed keys + (counts+offsets) + (entry row ids + entry hashes)
114    packed_bytes
115        .saturating_add(num_buckets.saturating_mul(8))
116        .saturating_add(num_rows.saturating_mul(12))
117}
118
119impl JoinIndexCache {
120    pub(crate) fn new(max_bytes: u64) -> Self {
121        Self {
122            entries: HashMap::new(),
123            clock: 0,
124            total_bytes: 0,
125            max_bytes,
126            stats: JoinIndexCacheStats::default(),
127        }
128    }
129
130    /// Decide whether to build a new join index for a relation.
131    ///
132    /// Heuristic: require higher "heat" for larger indexes, and avoid building under
133    /// memory pressure. Always skip if the estimated index cannot fit in the cache budget.
134    pub(crate) fn should_build(
135        &self,
136        est_index_bytes: u64,
137        build_heat: f32,
138        remaining_device_bytes: u64,
139        device_budget_bytes: u64,
140    ) -> bool {
141        let heat_threshold = if self.max_bytes > 0 && est_index_bytes > self.max_bytes / 2 {
142            0.6
143        } else {
144            0.3
145        };
146        let has_room =
147            remaining_device_bytes >= est_index_bytes.saturating_add(device_budget_bytes / 10);
148
149        build_heat >= heat_threshold && est_index_bytes <= self.max_bytes && has_room
150    }
151
152    pub(crate) fn clear(&mut self) {
153        let removed = self.entries.len() as u64;
154        self.entries.clear();
155        self.clock = 0;
156        self.total_bytes = 0;
157        self.stats.invalidations = self.stats.invalidations.saturating_add(removed);
158    }
159
160    pub(crate) fn get(&mut self, key: &JoinIndexKey) -> Option<&JoinIndexV2> {
161        self.stats.lookups = self.stats.lookups.saturating_add(1);
162        let Some(entry) = self.entries.get_mut(key) else {
163            self.stats.misses = self.stats.misses.saturating_add(1);
164            return None;
165        };
166        self.clock = self.clock.saturating_add(1);
167        entry.last_used = self.clock;
168        if let Some(index) = &entry.index {
169            self.stats.hits = self.stats.hits.saturating_add(1);
170            Some(index)
171        } else {
172            self.stats.misses = self.stats.misses.saturating_add(1);
173            None
174        }
175    }
176
177    pub(crate) fn insert(&mut self, key: JoinIndexKey, index: JoinIndexV2) {
178        let bytes = index.estimated_bytes();
179        if bytes > self.max_bytes {
180            return;
181        }
182
183        self.evict_until_fits(bytes);
184
185        self.clock = self.clock.saturating_add(1);
186        let last_used = self.clock;
187
188        if let Some(prev) = self.entries.remove(&key) {
189            self.total_bytes = self.total_bytes.saturating_sub(prev.bytes);
190        }
191
192        self.total_bytes = self.total_bytes.saturating_add(bytes);
193        self.entries.insert(
194            key,
195            CachedJoinIndex {
196                index: Some(index),
197                bytes,
198                last_used,
199            },
200        );
201        self.stats.builds = self.stats.builds.saturating_add(1);
202    }
203
204    pub(crate) fn remove(&mut self, key: &JoinIndexKey) {
205        if let Some(prev) = self.entries.remove(key) {
206            self.total_bytes = self.total_bytes.saturating_sub(prev.bytes);
207        }
208    }
209
210    pub(crate) fn remove_stale(&mut self, key: &JoinIndexKey) {
211        let before = self.entries.len();
212        self.remove(key);
213        if self.entries.len() < before {
214            self.stats.stale_rejections = self.stats.stale_rejections.saturating_add(1);
215        }
216    }
217
218    pub(crate) fn invalidate_rel(&mut self, rel: RelId) {
219        let keys: Vec<JoinIndexKey> = self
220            .entries
221            .keys()
222            .filter(|k| k.rel == rel)
223            .cloned()
224            .collect();
225        for key in keys {
226            if let Some(entry) = self.entries.remove(&key) {
227                self.total_bytes = self.total_bytes.saturating_sub(entry.bytes);
228                self.stats.invalidations = self.stats.invalidations.saturating_add(1);
229            }
230        }
231    }
232
233    pub(crate) fn evict_until_fits(&mut self, additional_bytes: u64) {
234        while !self.entries.is_empty()
235            && self.total_bytes.saturating_add(additional_bytes) > self.max_bytes
236        {
237            let mut oldest_key: Option<JoinIndexKey> = None;
238            let mut oldest_clock = u64::MAX;
239
240            for (k, v) in &self.entries {
241                if v.last_used < oldest_clock {
242                    oldest_clock = v.last_used;
243                    oldest_key = Some(k.clone());
244                }
245            }
246
247            let Some(key) = oldest_key else {
248                break;
249            };
250            if let Some(entry) = self.entries.remove(&key) {
251                self.total_bytes = self.total_bytes.saturating_sub(entry.bytes);
252                self.stats.evictions = self.stats.evictions.saturating_add(1);
253            } else {
254                break;
255            }
256        }
257    }
258
259    pub(crate) fn record_background_build_request(&mut self) {
260        self.stats.background_build_requests =
261            self.stats.background_build_requests.saturating_add(1);
262    }
263
264    pub(crate) fn record_background_build_complete(&mut self) {
265        self.stats.background_builds_completed =
266            self.stats.background_builds_completed.saturating_add(1);
267    }
268
269    pub(crate) fn record_background_build_deferred(&mut self) {
270        self.stats.background_builds_deferred =
271            self.stats.background_builds_deferred.saturating_add(1);
272    }
273
274    pub(crate) fn stats(&self) -> JoinIndexCacheStats {
275        let mut stats = self.stats.clone();
276        stats.entries = self.entries.len();
277        stats.total_bytes = self.total_bytes;
278        stats
279    }
280
281    #[cfg(test)]
282    fn insert_test_entry(&mut self, key: JoinIndexKey, bytes: u64) {
283        if bytes > self.max_bytes {
284            return;
285        }
286        self.evict_until_fits(bytes);
287        self.clock = self.clock.saturating_add(1);
288        self.total_bytes = self.total_bytes.saturating_add(bytes);
289        self.entries.insert(
290            key,
291            CachedJoinIndex {
292                index: None,
293                bytes,
294                last_used: self.clock,
295            },
296        );
297    }
298}
299
300#[cfg(test)]
301mod tests {
302    use super::*;
303    use xlog_core::{ScalarType, Schema};
304
305    fn schema(cols: Vec<(&str, ScalarType)>) -> Schema {
306        Schema::new(
307            cols.into_iter()
308                .map(|(name, ty)| (name.to_string(), ty))
309                .collect(),
310        )
311    }
312
313    #[test]
314    fn persistent_key_includes_schema_generation_key_and_device() {
315        let u32_schema = schema(vec![("k", ScalarType::U32)]);
316        let u64_schema = schema(vec![("k", ScalarType::U64)]);
317
318        let key = JoinIndexKey::new(RelId(7), 3, vec![0], &u32_schema, 0);
319        assert_eq!(key.rel, RelId(7));
320        assert_eq!(key.version, 3);
321        assert_eq!(key.key_cols, vec![0]);
322        assert_eq!(key.device_ordinal, 0);
323
324        assert_ne!(
325            key,
326            JoinIndexKey::new(RelId(7), 4, vec![0], &u32_schema, 0),
327            "generation/version must partition stale indexes"
328        );
329        assert_ne!(
330            key,
331            JoinIndexKey::new(RelId(7), 3, vec![0], &u64_schema, 0),
332            "schema changes must partition indexes"
333        );
334        assert_ne!(
335            key,
336            JoinIndexKey::new(RelId(7), 3, vec![0], &u32_schema, 1),
337            "device ordinal must partition indexes"
338        );
339    }
340
341    #[test]
342    fn persistent_cache_budget_evicts_lru_and_records_stats() {
343        let schema = schema(vec![("k", ScalarType::U32)]);
344        let key_a = JoinIndexKey::new(RelId(1), 1, vec![0], &schema, 0);
345        let key_b = JoinIndexKey::new(RelId(2), 1, vec![0], &schema, 0);
346        let mut cache = JoinIndexCache::new(100);
347
348        cache.insert_test_entry(key_a, 60);
349        cache.insert_test_entry(key_b, 60);
350
351        let stats = cache.stats();
352        assert_eq!(stats.entries, 1);
353        assert_eq!(stats.total_bytes, 60);
354        assert_eq!(stats.evictions, 1);
355    }
356
357    #[test]
358    fn persistent_cache_invalidation_records_removed_entries() {
359        let schema = schema(vec![("k", ScalarType::U32)]);
360        let key = JoinIndexKey::new(RelId(1), 1, vec![0], &schema, 0);
361        let mut cache = JoinIndexCache::new(100);
362
363        cache.insert_test_entry(key, 32);
364        cache.invalidate_rel(RelId(1));
365
366        let stats = cache.stats();
367        assert_eq!(stats.entries, 0);
368        assert_eq!(stats.total_bytes, 0);
369        assert_eq!(stats.invalidations, 1);
370    }
371}
xlog_runtime/executor/join_cache.rs

xlog_runtime/executor/
join_cache.rs
