use std::{
collections::HashMap,
sync::{Arc, Mutex},
};
use object_store::path::Path;
use parquet::file::metadata::ParquetMetaData;
use tokio::sync::OnceCell;
type ParquetMetadataCell = OnceCell<Arc<ParquetMetaData>>;
#[derive(Default)]
pub(crate) struct ParquetMetadataCache {
entries: Mutex<HashMap<(Path, u64), Arc<ParquetMetadataCell>>>,
}
impl ParquetMetadataCache {
pub(super) fn entry(&self, path: &Path, file_size: u64) -> Arc<ParquetMetadataCell> {
Arc::clone(
self.entries
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.entry((path.clone(), file_size))
.or_insert_with(|| Arc::new(OnceCell::new())),
)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn entries_are_reused_only_for_the_same_path_and_size() {
let cache = ParquetMetadataCache::default();
let path = Path::from("part.parquet");
let entry = cache.entry(&path, 100);
assert!(Arc::ptr_eq(&entry, &cache.entry(&path, 100)));
assert!(!Arc::ptr_eq(&entry, &cache.entry(&path, 101)));
assert!(!Arc::ptr_eq(
&entry,
&cache.entry(&Path::from("other.parquet"), 100)
));
}
}