Skip to main content

tenflowers_dataset/cache/
extension.rs

1//! Extension traits for adding caching to datasets
2//!
3//! This module provides convenient extension methods for any dataset type.
4
5use crate::cache::dataset::{CachedDataset, WarmingStrategy};
6use crate::Dataset;
7use tenflowers_core::Result;
8
9/// Extension trait for adding caching to any dataset
10pub trait CacheExt<T>: Dataset<T> + Sized {
11    /// Wrap this dataset with caching
12    fn cached(self, capacity: usize) -> CachedDataset<T, Self>
13    where
14        T: Clone + Send + Sync + 'static,
15    {
16        CachedDataset::new(self, capacity)
17    }
18
19    /// Wrap this dataset with caching and pre-warm the cache
20    fn cached_with_warming(
21        self,
22        capacity: usize,
23        strategy: WarmingStrategy,
24    ) -> Result<CachedDataset<T, Self>>
25    where
26        T: Clone + Send + Sync + 'static,
27    {
28        let cached = CachedDataset::new(self, capacity);
29        let indices = strategy.generate_indices(cached.len());
30        cached.warm_cache(&indices)?;
31        Ok(cached)
32    }
33}
34
35impl<T, D: Dataset<T>> CacheExt<T> for D {}