diff --git a/crates/clawhdf5-format/src/chunk_cache.rs b/crates/clawhdf5-format/src/chunk_cache.rs index 89703f4..aedb602 100644 --- a/crates/clawhdf5-format/src/chunk_cache.rs +++ b/crates/clawhdf5-format/src/chunk_cache.rs @@ -223,13 +223,32 @@ pub const DEFAULT_CACHE_BYTES: usize = 16 * 1024 * 1024; // 16 MiB /// coordinate map and reduces collision chains compared to power-of-two sizes. pub const DEFAULT_MAX_SLOTS: usize = 521; +/// Most datasets whose chunk index a [`ChunkCache`] keeps at once. +pub const MAX_INDEXED_DATASETS: usize = 64; + +/// Most chunk-index entries, summed over all datasets, a [`ChunkCache`] keeps. +/// Least-recently-used datasets' indexes are dropped past this (the dataset +/// being read is always kept), so a file with many or huge chunked datasets +/// cannot grow the cache without bound. +pub const MAX_INDEXED_CHUNKS: usize = 1 << 20; + +/// The dataset key the address-less (legacy) methods use when +/// [`ChunkCache::ensure_dataset`] has not been called. +#[cfg(feature = "std")] +const UNBOUND_DATASET: u64 = u64::MAX; + // --------------------------------------------------------------------------- // LRU entry // --------------------------------------------------------------------------- +/// Decompressed chunks are keyed by dataset *and* coordinate: every chunked +/// dataset has a chunk at (0, 0, ...), so the coordinate alone is ambiguous. +#[cfg(feature = "std")] +type SlotKey = (u64, ChunkCoord); + #[cfg(feature = "std")] struct CachedChunk { - coord: ChunkCoord, + key: SlotKey, /// Shared so a cache hit is a refcount bump, not a copy of the whole /// (potentially large) decompressed chunk. data: Arc, @@ -237,21 +256,48 @@ struct CachedChunk { last_access: u64, } +/// Per-dataset index state. +#[cfg(feature = "std")] +#[derive(Default)] +struct DatasetEntry { + /// Chunk coordinate -> ChunkInfo (offset + size in file). + index: Option>>, + /// Pre-built chunk index for O(1) coordinate lookups. + chunk_index: Option>, + /// Pre-computed chunk layout for fast assembly. + chunk_layout: Option>, + /// Tick of the last use, for dropping the least recently used dataset. + last_used: u64, +} + +#[cfg(feature = "std")] +impl DatasetEntry { + fn weight(&self) -> usize { + self.index.as_ref().map_or(0, |m| m.len()) + + self.chunk_index.as_ref().map_or(0, |c| c.num_chunks()) + } +} + // --------------------------------------------------------------------------- // ChunkCache // --------------------------------------------------------------------------- -/// A per-dataset chunk cache with hash-based index and LRU eviction. +/// A per-file chunk cache: chunk indexes per dataset, plus an LRU of +/// decompressed chunks, all keyed by dataset. /// -/// # Usage +/// A dataset is identified by the address of its chunk index (B-tree, fixed +/// or extensible array, ...), which is unique within a file. Every method +/// that takes an `addr` works on that dataset only, so threads reading +/// different datasets through one shared cache never see each other's +/// chunks. The address-less methods (`has_index`, `populate_index`, +/// `get_decompressed`, ...) act on the dataset last bound with +/// [`Self::ensure_dataset`]; that binding is shared state, so concurrent +/// readers must use the `*_in` / `*_for` methods instead (the chunked +/// readers in [`crate::chunked_read`] do). /// -/// ```ignore -/// let cache = ChunkCache::new(); -/// // Pass &cache to read_chunked_data — it will populate the index lazily. -/// ``` -/// -/// The cache is wrapped in `Mutex` internally so it can be mutated through -/// shared references (thread-safe). +/// Memory is bounded: decompressed data by `max_bytes`/`max_slots` across +/// all datasets, indexes by [`MAX_INDEXED_DATASETS`] and +/// [`MAX_INDEXED_CHUNKS`]. /// /// Only available with the `std` feature because it requires `std::sync::Mutex`. #[cfg(feature = "std")] @@ -261,26 +307,20 @@ pub struct ChunkCache { #[cfg(feature = "std")] struct CacheInner { - /// Hash index: chunk coordinate -> ChunkInfo (offset + size in file). - /// Populated once per dataset on first access. - index: Option>, + /// Per-dataset chunk indexes, keyed by chunk-index address. + datasets: HashMap, - /// Address of the dataset (its chunk-index base address) that the cached - /// index, chunk index, layout, and decompressed slots currently belong to. - /// The cache is shared per file across datasets, so every cached-read entry - /// checks this and resets the per-dataset state when the dataset changes — - /// otherwise one dataset's chunk index (with its own rank) would be reused - /// for another, corrupting reads. - index_addr: Option, + /// Dataset the address-less methods act on (see `ensure_dataset`). + current: Option, /// LRU cache of decompressed chunk data. slots: Vec, - /// Coordinate -> index into `slots`, for O(1) lookup instead of a linear + /// Key -> index into `slots`, for O(1) lookup instead of a linear /// scan. Kept in sync with `slots` on every insert/evict/clear — in /// particular, `slots.swap_remove(i)` moves the last element into slot /// `i`, so the moved element's index entry must be updated too. - slot_index: HashMap, + slot_index: HashMap, /// Current total bytes of cached decompressed data. current_bytes: usize, @@ -294,17 +334,145 @@ struct CacheInner { /// Monotonic counter for LRU ordering. tick: u64, - /// Last accessed chunk coordinate (for sequential detection). - last_coord: Option, + /// Last accessed chunk (for sequential detection). + last_coord: Option, /// Access pattern statistics. stats: AccessStats, +} - /// Pre-built chunk index for O(1) coordinate lookups. - chunk_index: Option, +#[cfg(feature = "std")] +impl CacheInner { + fn current(&self) -> u64 { + self.current.unwrap_or(UNBOUND_DATASET) + } - /// Pre-computed chunk layout for fast assembly. - chunk_layout: Option, + fn touch(&mut self, addr: u64) -> &mut DatasetEntry { + self.tick += 1; + let tick = self.tick; + let entry = self.datasets.entry(addr).or_default(); + entry.last_used = tick; + entry + } + + fn entry(&self, addr: u64) -> Option<&DatasetEntry> { + self.datasets.get(&addr) + } + + /// Drop least-recently-used datasets' indexes (never `keep`'s) until the + /// dataset and chunk-entry budgets hold. + fn trim_datasets(&mut self, keep: u64) { + loop { + let total: usize = self.datasets.values().map(DatasetEntry::weight).sum(); + if self.datasets.len() <= MAX_INDEXED_DATASETS && total <= MAX_INDEXED_CHUNKS { + return; + } + let victim = self + .datasets + .iter() + .filter(|(a, _)| **a != keep) + .min_by_key(|(_, e)| e.last_used) + .map(|(a, _)| *a); + match victim { + Some(a) => { + self.datasets.remove(&a); + } + None => return, + } + } + } + + fn get_decompressed(&mut self, addr: u64, coord: &[u64]) -> Option> { + self.tick += 1; + let tick = self.tick; + + // Track sequential vs random access + let is_sequential = self.last_coord.as_ref().is_some_and(|(prev_addr, prev)| { + // Sequential if exactly one dimension changed + let changes: usize = prev + .iter() + .zip(coord.iter()) + .filter(|(a, b)| a != b) + .count(); + *prev_addr == addr && changes <= 1 + }); + if is_sequential { + self.stats.sequential_count += 1; + } else if self.last_coord.is_some() { + self.stats.random_count += 1; + } + let key: SlotKey = (addr, coord.to_vec()); + let found = if let Some(&idx) = self.slot_index.get(&key) { + self.slots[idx].last_access = tick; + Some(Arc::clone(&self.slots[idx].data)) + } else { + None + }; + self.last_coord = Some(key); + if let Some(ref data) = found { + self.stats.hits += 1; + self.stats.bytes_read += data.len() as u64; + } else { + self.stats.misses += 1; + } + found + } + + fn put_decompressed( + &mut self, + key: SlotKey, + data: Arc, + ) -> Arc { + let data_len = data.len(); + + // Don't cache if single chunk exceeds budget — still return the data + // to the caller, just don't retain it. + if data_len > self.max_bytes { + return data; + } + + // Check if already present + self.tick += 1; + let tick = self.tick; + if let Some(&idx) = self.slot_index.get(&key) { + self.slots[idx].last_access = tick; + return Arc::clone(&self.slots[idx].data); // already cached + } + + // Evict until we have room + while self.slots.len() >= self.max_slots + || (self.current_bytes + data_len > self.max_bytes && !self.slots.is_empty()) + { + // Find LRU slot + let lru_idx = self + .slots + .iter() + .enumerate() + .min_by_key(|(_, s)| s.last_access) + .map(|(i, _)| i) + .unwrap(); + let removed = self.slots.swap_remove(lru_idx); + self.slot_index.remove(&removed.key); + // swap_remove moved the former last element into `lru_idx` (unless + // it *was* the last element) — fix up that element's index entry. + if lru_idx < self.slots.len() { + let moved_key = self.slots[lru_idx].key.clone(); + self.slot_index.insert(moved_key, lru_idx); + } + self.current_bytes -= removed.data.len(); + self.stats.evictions += 1; + } + + self.current_bytes += data_len; + let new_idx = self.slots.len(); + self.slot_index.insert(key.clone(), new_idx); + self.slots.push(CachedChunk { + key, + data: Arc::clone(&data), + last_access: tick, + }); + data + } } /// Access pattern statistics tracked by the chunk cache. @@ -356,8 +524,8 @@ impl ChunkCache { pub fn with_capacity(max_bytes: usize, max_slots: usize) -> Self { Self { inner: std::sync::Mutex::new(CacheInner { - index: None, - index_addr: None, + datasets: HashMap::new(), + current: None, slots: Vec::with_capacity(max_slots.min(64)), slot_index: HashMap::with_capacity(max_slots.min(64)), current_bytes: 0, @@ -366,340 +534,331 @@ impl ChunkCache { tick: 0, last_coord: None, stats: AccessStats::default(), - chunk_index: None, - chunk_layout: None, }), } } - // ----- Index operations ----- + fn lock(&self) -> std::sync::MutexGuard<'_, CacheInner> { + self.inner.lock().unwrap_or_else(|e| e.into_inner()) + } /// The most decompressed bytes this cache will hold. pub fn max_bytes(&self) -> usize { - self.inner.lock().map(|g| g.max_bytes).unwrap_or(0) + self.lock().max_bytes } - /// Bind the cache to the dataset at chunk-index address `addr`. + // ----- Dataset-keyed operations (safe to use concurrently) ----- + + /// The chunk list of the dataset whose chunk index is at `addr`. /// - /// The cache is shared per file across all of its datasets. If the cache - /// currently holds state for a different dataset, all per-dataset state - /// (chunk index, chunk-index map, layout, and decompressed slots) is - /// dropped so the next access rebuilds it for this dataset. Reading the - /// same dataset again is a no-op, preserving the cache's benefit for - /// repeated/sequential access. Returns `true` if a reset occurred. - pub fn ensure_dataset(&self, addr: u64) -> bool { - let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - if inner.index_addr == Some(addr) { - return false; + /// On the first call for a dataset, `build` scans its chunk index; the + /// result is kept (offsets truncated to `rank` for the lookup key), so + /// later calls skip the scan. `build` runs without the cache lock held; + /// if two threads race to build the same dataset's index, the first + /// stored one wins and both return equivalent lists. + pub fn chunks_for( + &self, + addr: u64, + rank: usize, + build: impl FnOnce() -> Result, E>, + ) -> Result, E> { + Ok(self + .index_for(addr, rank, build)? + .values() + .cloned() + .collect()) + } + + fn index_for( + &self, + addr: u64, + rank: usize, + build: impl FnOnce() -> Result, E>, + ) -> Result>, E> { + if let Some(index) = self.lock().touch(addr).index.clone() { + return Ok(index); } - inner.index = None; - inner.chunk_index = None; - inner.chunk_layout = None; - inner.slots.clear(); - inner.slot_index.clear(); - inner.current_bytes = 0; - inner.last_coord = None; - inner.index_addr = Some(addr); - true + let chunks = build()?; + let map: HashMap = chunks + .into_iter() + .map(|ci| (ci.offsets.iter().take(rank).copied().collect(), ci)) + .collect(); + let mut inner = self.lock(); + let entry = inner.touch(addr); + let index = Arc::clone(entry.index.get_or_insert_with(|| Arc::new(map))); + inner.trim_datasets(addr); + Ok(index) } - /// Returns `true` if the chunk index has been built. + /// The pre-computed assembly layout of the dataset at `addr`, building + /// its chunk index (via `build`, as in [`Self::chunks_for`]) and layout on + /// first use. + pub fn chunk_layout_for( + &self, + addr: u64, + rank: usize, + build: impl FnOnce() -> Result, E>, + ds_dims: &[usize], + chunk_dims: &[usize], + elem_size: usize, + ) -> Result, E> { + let (layout, chunk_index) = { + let mut inner = self.lock(); + let entry = inner.touch(addr); + (entry.chunk_layout.clone(), entry.chunk_index.clone()) + }; + if let Some(layout) = layout { + return Ok(layout); + } + let chunk_index = match chunk_index { + Some(ci) => ci, + None => { + let index = self.index_for(addr, rank, build)?; + let chunks: Vec = index.values().cloned().collect(); + Arc::new(ChunkIndex::build(&chunks, rank)) + } + }; + let layout = ChunkLayout::build(&chunk_index, ds_dims, chunk_dims, elem_size); + let mut inner = self.lock(); + let entry = inner.touch(addr); + entry.chunk_index.get_or_insert(chunk_index); + let layout = Arc::clone(entry.chunk_layout.get_or_insert_with(|| Arc::new(layout))); + inner.trim_datasets(addr); + Ok(layout) + } + + /// Cached decompressed chunk at `coord` of the dataset at `addr`. + /// + /// O(1) lookup; the clone is an `Arc` refcount bump, not a copy of the + /// underlying decompressed data. + pub fn get_decompressed_in(&self, addr: u64, coord: &[u64]) -> Option> { + self.lock().get_decompressed(addr, coord) + } + + /// Cache decompressed chunk data for `coord` of the dataset at `addr`. + /// Returns the `Arc`-shared buffer now cached (or already cached). + pub fn put_decompressed_in( + &self, + addr: u64, + coord: ChunkCoord, + data: Vec, + ) -> Arc { + self.put_decompressed_aligned_in(addr, coord, CacheAlignedBuffer::from_vec(data)) + } + + /// [`Self::put_decompressed_in`] for an already-aligned buffer. + pub fn put_decompressed_aligned_in( + &self, + addr: u64, + coord: ChunkCoord, + data: CacheAlignedBuffer, + ) -> Arc { + let data = Arc::new(data); + self.lock().put_decompressed((addr, coord), data) + } + + /// Record that the given chunk coordinates of the dataset at `addr` are + /// predicted to be accessed soon (bookkeeping only). + /// + /// This does **not** prefetch or pre-decompress anything — it only + /// checks whether each coordinate is already in the chunk index and + /// updates access-pattern stats accordingly. + pub fn prefetch_hint_in(&self, addr: u64, next_coords: &[ChunkCoord]) { + let mut inner = self.lock(); + let Some(index) = inner.entry(addr).and_then(|e| e.index.clone()) else { + return; + }; + let known = next_coords + .iter() + .filter(|c| index.contains_key(*c)) + .count(); + inner.stats.sequential_count += known as u64; + } + + // ----- Address-less operations on the bound dataset ----- + + /// Bind the address-less methods to the dataset at chunk-index address + /// `addr`. Returns `true` if this changed the bound dataset. + /// + /// Each dataset's state is kept separately, so switching loses nothing + /// and never exposes one dataset's index or chunks to another. The + /// binding itself is shared, though: concurrent readers should use the + /// `addr`-taking methods rather than bind and then call these. + pub fn ensure_dataset(&self, addr: u64) -> bool { + let mut inner = self.lock(); + let changed = inner.current != Some(addr); + inner.current = Some(addr); + changed + } + + /// Returns `true` if the bound dataset's chunk index has been built. pub fn has_index(&self) -> bool { - self.inner - .lock() - .unwrap_or_else(|e| e.into_inner()) - .index - .is_some() + let inner = self.lock(); + inner + .entry(inner.current()) + .is_some_and(|e| e.index.is_some()) } - /// Build the chunk index from a pre-collected list of `ChunkInfo`. + /// Build the bound dataset's chunk index from a pre-collected list of + /// `ChunkInfo`. /// /// The `rank` parameter is used to truncate offsets to spatial dims only /// (B-tree v1 stores rank+1 offsets). pub fn populate_index(&self, chunks: &[ChunkInfo], rank: usize) { - let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - if inner.index.is_some() { - return; // already populated - } - let mut map = HashMap::with_capacity(chunks.len()); - - for ci in chunks { - let coord: ChunkCoord = ci.offsets.iter().take(rank).copied().collect(); - map.insert(coord, ci.clone()); - } - inner.index = Some(map); + let addr = self.lock().current(); + let _ = self.index_for::(addr, rank, || Ok(chunks.to_vec())); } - /// Look up a chunk by its spatial coordinate in the index. + /// Look up a chunk by its spatial coordinate in the bound dataset's index. pub fn lookup_index(&self, coord: &[u64]) -> Option { - let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - inner.index.as_ref()?.get(coord).cloned() + let inner = self.lock(); + inner + .entry(inner.current())? + .index + .as_ref()? + .get(coord) + .cloned() } - /// Return all indexed chunks as a `Vec` (order unspecified). + /// Return all of the bound dataset's indexed chunks (order unspecified). pub fn all_indexed_chunks(&self) -> Option> { - let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - inner.index.as_ref().map(|m| m.values().cloned().collect()) + let inner = self.lock(); + let index = inner.entry(inner.current())?.index.as_ref()?; + Some(index.values().cloned().collect()) } - // ----- Chunk index (pre-built coordinate → ChunkInfo map) ----- - - /// Returns `true` if the chunk B-tree index has been built. + /// Returns `true` if the bound dataset's `ChunkIndex` has been built. pub fn has_chunk_index(&self) -> bool { - self.inner - .lock() - .unwrap_or_else(|e| e.into_inner()) - .chunk_index - .is_some() + let inner = self.lock(); + inner + .entry(inner.current()) + .is_some_and(|e| e.chunk_index.is_some()) } - /// Build and store the chunk B-tree index from a pre-collected list of `ChunkInfo`. + /// Build and store the bound dataset's `ChunkIndex`. pub fn populate_chunk_index(&self, chunks: &[ChunkInfo], rank: usize) { - let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - if inner.chunk_index.is_some() { - return; - } - inner.chunk_index = Some(ChunkIndex::build(chunks, rank)); + let built = Arc::new(ChunkIndex::build(chunks, rank)); + let mut inner = self.lock(); + let addr = inner.current(); + inner.touch(addr).chunk_index.get_or_insert(built); + inner.trim_datasets(addr); } - // ----- Chunk layout (pre-computed assembly plan) ----- - - /// Returns `true` if the chunk layout has been computed. + /// Returns `true` if the bound dataset's chunk layout has been computed. pub fn has_chunk_layout(&self) -> bool { - self.inner - .lock() - .unwrap_or_else(|e| e.into_inner()) - .chunk_layout - .is_some() + let inner = self.lock(); + inner + .entry(inner.current()) + .is_some_and(|e| e.chunk_layout.is_some()) } - /// Build and store the pre-computed chunk layout for fast assembly. + /// Build and store the bound dataset's chunk layout (needs its + /// `ChunkIndex`; does nothing without one). pub fn populate_chunk_layout(&self, ds_dims: &[usize], chunk_dims: &[usize], elem_size: usize) { - let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - if inner.chunk_layout.is_some() { + let mut inner = self.lock(); + let addr = inner.current(); + let entry = inner.touch(addr); + if entry.chunk_layout.is_some() { return; } - if let Some(ref idx) = inner.chunk_index { - inner.chunk_layout = Some(ChunkLayout::build(idx, ds_dims, chunk_dims, elem_size)); + if let Some(idx) = entry.chunk_index.clone() { + entry.chunk_layout = Some(Arc::new(ChunkLayout::build( + &idx, ds_dims, chunk_dims, elem_size, + ))); } } - /// Execute a function with a reference to the chunk layout. - /// - /// Returns `None` if the layout hasn't been computed yet. + /// Execute a function with a reference to the bound dataset's chunk + /// layout. Returns `None` if the layout hasn't been computed yet. pub fn with_chunk_layout(&self, f: F) -> Option where F: FnOnce(&ChunkLayout) -> R, { - let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - inner.chunk_layout.as_ref().map(f) + let layout = { + let inner = self.lock(); + inner.entry(inner.current())?.chunk_layout.clone()? + }; + Some(f(&layout)) } - // ----- Decompressed data cache (LRU) ----- - - /// Try to get cached decompressed data for a chunk coordinate. + /// Try to get cached decompressed data for a chunk of the bound dataset. /// - /// O(1) lookup. Returns an owned copy for API compatibility with callers - /// that need a `Vec`; prefer [`Self::get_decompressed_aligned`] when - /// an `Arc`-shared buffer works for the caller, since that avoids the - /// copy entirely. + /// Returns an owned copy; prefer [`Self::get_decompressed_aligned`] when + /// an `Arc`-shared buffer works for the caller. pub fn get_decompressed(&self, coord: &[u64]) -> Option> { self.get_decompressed_aligned(coord) .map(|arc| arc.as_slice().to_vec()) } - /// Try to get a reference-counted clone of the aligned buffer for a chunk. - /// - /// O(1) index lookup; the clone is an `Arc` refcount bump, not a copy of - /// the underlying decompressed data. + /// Reference-counted cached buffer for a chunk of the bound dataset. pub fn get_decompressed_aligned(&self, coord: &[u64]) -> Option> { - let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - inner.tick += 1; - let tick = inner.tick; - - // Track sequential vs random access - let is_sequential = inner.last_coord.as_ref().is_some_and(|prev| { - // Sequential if exactly one dimension changed - let changes: usize = prev - .iter() - .zip(coord.iter()) - .filter(|(a, b)| a != b) - .count(); - changes <= 1 - }); - if is_sequential { - inner.stats.sequential_count += 1; - } else if inner.last_coord.is_some() { - inner.stats.random_count += 1; - } - inner.last_coord = Some(coord.to_vec()); - - let found = if let Some(&idx) = inner.slot_index.get(coord) { - inner.slots[idx].last_access = tick; - Some(Arc::clone(&inner.slots[idx].data)) - } else { - None - }; - if let Some(ref data) = found { - inner.stats.hits += 1; - inner.stats.bytes_read += data.len() as u64; - } else { - inner.stats.misses += 1; - } - found + let mut inner = self.lock(); + let addr = inner.current(); + inner.get_decompressed(addr, coord) } - /// Insert decompressed chunk data into the LRU cache. - /// - /// The data is stored in a [`CacheAlignedBuffer`] so subsequent reads - /// return cache-line-aligned memory. Returns the `Arc`-shared buffer that - /// is now cached (or already was), so the caller can reuse it directly - /// instead of holding a separate copy of the same data. + /// Insert decompressed chunk data for the bound dataset into the LRU + /// cache, returning the `Arc`-shared buffer now cached. pub fn put_decompressed(&self, coord: ChunkCoord, data: Vec) -> Arc { - let aligned = CacheAlignedBuffer::from_vec(data); - self.put_decompressed_aligned(coord, aligned) + self.put_decompressed_aligned(coord, CacheAlignedBuffer::from_vec(data)) } - /// Insert an already-aligned buffer into the LRU cache. - /// - /// Returns the `Arc`-shared buffer now held by the cache (the one just - /// inserted, or the existing cached copy if `coord` was already present). + /// Insert an already-aligned buffer for the bound dataset. pub fn put_decompressed_aligned( &self, coord: ChunkCoord, data: CacheAlignedBuffer, ) -> Arc { let data = Arc::new(data); - let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - let data_len = data.len(); - - // Don't cache if single chunk exceeds budget — still return the data - // to the caller, just don't retain it. - if data_len > inner.max_bytes { - return data; - } - - // Check if already present - inner.tick += 1; - let tick = inner.tick; - if let Some(&idx) = inner.slot_index.get(&coord) { - inner.slots[idx].last_access = tick; - return Arc::clone(&inner.slots[idx].data); // already cached - } - - // Evict until we have room - while inner.slots.len() >= inner.max_slots - || (inner.current_bytes + data_len > inner.max_bytes && !inner.slots.is_empty()) - { - // Find LRU slot - let lru_idx = inner - .slots - .iter() - .enumerate() - .min_by_key(|(_, s)| s.last_access) - .map(|(i, _)| i) - .unwrap(); - let removed = inner.slots.swap_remove(lru_idx); - inner.slot_index.remove(&removed.coord); - // swap_remove moved the former last element into `lru_idx` (unless - // it *was* the last element) — fix up that element's index entry. - if lru_idx < inner.slots.len() { - let moved_coord = inner.slots[lru_idx].coord.clone(); - inner.slot_index.insert(moved_coord, lru_idx); - } - inner.current_bytes -= removed.data.len(); - inner.stats.evictions += 1; - } - - inner.current_bytes += data_len; - let new_idx = inner.slots.len(); - inner.slot_index.insert(coord.clone(), new_idx); - inner.slots.push(CachedChunk { - coord, - data: Arc::clone(&data), - last_access: tick, - }); - data + let mut inner = self.lock(); + let addr = inner.current(); + inner.put_decompressed((addr, coord), data) } - /// Clear the entire cache (index + decompressed data). + /// [`Self::prefetch_hint_in`] for the bound dataset. + pub fn prefetch_hint(&self, next_coords: &[ChunkCoord]) { + let addr = self.lock().current(); + self.prefetch_hint_in(addr, next_coords); + } + + // ----- Whole-cache operations ----- + + /// Clear the entire cache (indexes + decompressed data + stats). pub fn clear(&self) { - let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - inner.index = None; - inner.index_addr = None; + let mut inner = self.lock(); + inner.datasets.clear(); + inner.current = None; inner.slots.clear(); inner.slot_index.clear(); inner.current_bytes = 0; inner.tick = 0; inner.last_coord = None; inner.stats = AccessStats::default(); - inner.chunk_index = None; - inner.chunk_layout = None; - } - - /// Record that the given chunk coordinates are predicted to be accessed - /// soon (bookkeeping only). - /// - /// This does **not** prefetch or pre-decompress anything — it only - /// checks whether each coordinate is already in the chunk index and - /// updates access-pattern stats accordingly. Real prefetching (e.g. - /// background pre-decompression) is not implemented. - pub fn prefetch_hint(&self, next_coords: &[ChunkCoord]) { - let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - if inner.index.is_none() { - return; - } - drop(inner); - // For each predicted coordinate, verify it exists in the index. - // The index is already populated, so this is a no-op for known chunks. - // The purpose is to signal intent — callers can pre-decompress if needed. - // We touch the stats to record that prefetch hints were issued. - let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - for coord in next_coords { - let exists = inner - .index - .as_ref() - .map(|idx| idx.contains_key(coord)) - .unwrap_or(false); - if exists { - inner.stats.sequential_count += 1; - } - } } /// Return the current access pattern statistics. pub fn access_stats(&self) -> AccessStats { - self.inner - .lock() - .unwrap_or_else(|e| e.into_inner()) - .stats - .clone() + self.lock().stats.clone() } /// Update the sweep direction label in the access stats. pub fn set_sweep_direction(&self, direction: &'static str) { - self.inner - .lock() - .unwrap_or_else(|e| e.into_inner()) - .stats - .sweep_direction = Some(direction); + self.lock().stats.sweep_direction = Some(direction); } - /// Number of decompressed chunks currently cached. + /// Number of decompressed chunks currently cached (all datasets). pub fn cached_chunk_count(&self) -> usize { - self.inner - .lock() - .unwrap_or_else(|e| e.into_inner()) - .slots - .len() + self.lock().slots.len() } - /// Total bytes of decompressed data currently cached. + /// Total bytes of decompressed data currently cached (all datasets). pub fn cached_bytes(&self) -> usize { - self.inner - .lock() - .unwrap_or_else(|e| e.into_inner()) - .current_bytes + self.lock().current_bytes + } + + /// Number of datasets whose chunk index is currently kept. + pub fn indexed_dataset_count(&self) -> usize { + self.lock().datasets.len() } } @@ -808,6 +967,92 @@ mod tests { assert_eq!(cache.cached_bytes(), 0); } + #[test] + fn datasets_sharing_coordinates_stay_separate() { + let cache = ChunkCache::new(); + let a = vec![make_chunk(vec![0, 0], 0x100, 8)]; + let b = vec![make_chunk(vec![0, 0], 0x900, 8)]; + let got_a = cache.chunks_for::<()>(1, 1, || Ok(a.clone())).unwrap(); + let got_b = cache.chunks_for::<()>(2, 1, || Ok(b.clone())).unwrap(); + assert_eq!(got_a[0].address, 0x100); + assert_eq!(got_b[0].address, 0x900); + // Built once per dataset: a second lookup doesn't call the builder. + let again = cache + .chunks_for::<()>(1, 1, || panic!("index rebuilt")) + .unwrap(); + assert_eq!(again[0].address, 0x100); + + cache.put_decompressed_in(1, vec![0], vec![1; 4]); + cache.put_decompressed_in(2, vec![0], vec![2; 4]); + assert_eq!( + cache.get_decompressed_in(1, &[0]).unwrap().as_slice(), + &[1; 4] + ); + assert_eq!( + cache.get_decompressed_in(2, &[0]).unwrap().as_slice(), + &[2; 4] + ); + assert!(cache.get_decompressed_in(3, &[0]).is_none()); + assert_eq!(cache.cached_chunk_count(), 2); + + // The bound-dataset methods see only the bound dataset. + cache.ensure_dataset(2); + assert_eq!(cache.lookup_index(&[0]).unwrap().address, 0x900); + assert_eq!(cache.get_decompressed(&[0]).unwrap(), vec![2; 4]); + } + + #[test] + fn dataset_indexes_are_bounded() { + let cache = ChunkCache::new(); + for addr in 0..(MAX_INDEXED_DATASETS as u64 + 10) { + cache + .chunks_for::<()>(addr, 1, || Ok(vec![make_chunk(vec![0], addr, 8)])) + .unwrap(); + } + assert_eq!(cache.indexed_dataset_count(), MAX_INDEXED_DATASETS); + + // One huge index evicts the others but is itself kept. + let huge: Vec = (0..MAX_INDEXED_CHUNKS as u64) + .map(|i| make_chunk(vec![i], i, 8)) + .collect(); + let got = cache.chunks_for::<()>(9999, 1, || Ok(huge)).unwrap(); + assert_eq!(got.len(), MAX_INDEXED_CHUNKS); + assert_eq!(cache.indexed_dataset_count(), 1); + } + + #[test] + fn concurrent_readers_of_different_datasets_see_their_own_chunks() { + let cache = std::sync::Arc::new(ChunkCache::with_capacity(1 << 20, 64)); + let handles: Vec<_> = (0..8u64) + .map(|t| { + let cache = std::sync::Arc::clone(&cache); + std::thread::spawn(move || { + for round in 0..500u64 { + let addr = (t + round) % 16; + let coord = vec![round % 4]; + let chunks = cache + .chunks_for::<()>(addr, 1, || { + Ok((0..4).map(|c| make_chunk(vec![c], addr, 8)).collect()) + }) + .unwrap(); + assert!(chunks.iter().all(|c| c.address == addr)); + let want = vec![addr as u8; 8]; + let got = match cache.get_decompressed_in(addr, &coord) { + Some(hit) => hit.to_vec(), + None => cache + .put_decompressed_in(addr, coord, want.clone()) + .to_vec(), + }; + assert_eq!(got, want); + } + }) + }) + .collect(); + for h in handles { + h.join().unwrap(); + } + } + #[test] fn duplicate_insert_is_noop() { let cache = ChunkCache::new(); diff --git a/crates/clawhdf5-format/src/chunked_read.rs b/crates/clawhdf5-format/src/chunked_read.rs index f07ae27..3a3dc05 100644 --- a/crates/clawhdf5-format/src/chunked_read.rs +++ b/crates/clawhdf5-format/src/chunked_read.rs @@ -15,7 +15,7 @@ use crate::datatype::Datatype; use crate::error::FormatError; use crate::extensible_array::{ExtensibleArrayHeader, read_extensible_array_chunks}; use crate::filter_pipeline::FilterPipeline; -use crate::filters::decompress_chunk; +use crate::filters::{all_filters_skipped, decompress_chunk_masked}; use crate::fixed_array::{FixedArrayHeader, read_fixed_array_chunks}; #[cfg(feature = "std")] use std::sync::Arc; @@ -65,11 +65,13 @@ fn decompress_all_chunks( let raw_chunk = &file_data[c_addr..c_addr + size]; let decompressed = if let Some(pl) = pipeline { - if chunk_info.filter_mask == 0 { - decompress_chunk(raw_chunk, pl, chunk_total_bytes, element_size)? - } else { - raw_chunk.to_vec() - } + decompress_chunk_masked( + raw_chunk, + pl, + chunk_total_bytes, + element_size, + chunk_info.filter_mask, + )? } else { raw_chunk.to_vec() }; @@ -223,6 +225,10 @@ pub fn collect_chunk_info( collect_chunk_info_inner(file_data, btree_address, ndims, offset_size, length_size, 0) } +/// Width of each chunk offset in a v1 chunk B-tree key, independent of the +/// file's size-of-offsets. +const CHUNK_KEY_OFFSET_SIZE: u8 = 8; + /// Maximum recursion depth for chunk B-tree traversal (malformed/cyclic data /// protection), matching `btree_v1.rs`'s `MAX_BTREE_DEPTH`. const MAX_CHUNK_BTREE_DEPTH: usize = 64; @@ -260,8 +266,14 @@ fn collect_chunk_info_inner( let mut pos = offset + 8 + os * 2; // skip left/right sibling - // Key size: chunk_size(4) + filter_mask(4) + ndims * offset_size - let key_size = 4 + 4 + ndims * os; + // Key: chunk_size(4) + filter_mask(4) + one offset per dimension. The + // offsets are always 8 bytes each — they are dataset coordinates, not file + // addresses, so they do not follow the superblock's size-of-offsets (only + // the sibling and child addresses do). + let key_size = ndims + .checked_mul(CHUNK_KEY_OFFSET_SIZE as usize) + .and_then(|n| n.checked_add(8)) + .ok_or_else(|| FormatError::ChunkedReadError("chunk key too large".into()))?; if node_level == 0 { // Leaf node: keys and children interleaved @@ -287,8 +299,8 @@ fn collect_chunk_info_inner( let mut offsets = Vec::with_capacity(ndims); let mut kp = pos + 8; for _ in 0..ndims { - offsets.push(read_offset(file_data, kp, offset_size)?); - kp += os; + offsets.push(read_offset(file_data, kp, CHUNK_KEY_OFFSET_SIZE)?); + kp += CHUNK_KEY_OFFSET_SIZE as usize; } pos += key_size; @@ -507,6 +519,7 @@ pub fn list_chunks( addr_opt, single_filtered_size, single_filter_mask, + unfiltered_edges, ) = match layout { DataLayout::Chunked { chunk_dimensions, @@ -515,6 +528,7 @@ pub fn list_chunks( chunk_index_type, single_chunk_filtered_size, single_chunk_filter_mask, + dont_filter_partial_edge_chunks, } => ( chunk_dimensions, *version, @@ -522,6 +536,7 @@ pub fn list_chunks( *btree_address, *single_chunk_filtered_size, *single_chunk_filter_mask, + *dont_filter_partial_edge_chunks, ), _ => { return Err(FormatError::ChunkedReadError( @@ -554,7 +569,7 @@ pub fn list_chunks( } // Collect chunks based on version and index type - let chunks = match (version, chunk_index_type) { + let mut chunks = match (version, chunk_index_type) { (3, _) => { let ndims = chunk_dimensions.len(); // rank+1 collect_chunk_info(file_data, addr, ndims, offset_size, length_size)? @@ -635,6 +650,23 @@ pub fn list_chunks( } }; + // With "don't filter partial edge chunks", a chunk that extends past the + // dataset's extent is stored raw while its filter mask still reads 0. + // Mark every filter skipped so all read paths copy it as-is. + if unfiltered_edges { + for chunk in &mut chunks { + let partial = chunk + .offsets + .iter() + .zip(&chunk_dims) + .zip(&ds_dims) + .any(|((&off, &cd), &dd)| off.saturating_add(cd as u64) > dd as u64); + if partial { + chunk.filter_mask = u32::MAX; + } + } + } + Ok((chunks, chunk_dims)) } @@ -806,24 +838,20 @@ pub fn read_chunked_data_cached( ))); } - // The per-file cache is shared across datasets; bind it to this one so a - // different dataset's chunk index is never reused for this read. - cache.ensure_dataset(addr); - - // Populate chunk index on first access - if !cache.has_index() { - let (chunks, _) = list_chunks( + // The per-file cache is shared across datasets (and threads); every + // lookup is keyed by this dataset's chunk-index address, so another + // dataset's index or chunks are never used for this read. + let chunks = cache.chunks_for(addr, rank, || { + list_chunks( file_data, layout, dataspace, elem_size, offset_size, length_size, - )?; - cache.populate_index(&chunks, rank); - } - - let chunks = cache.all_indexed_chunks().unwrap_or_default(); + ) + .map(|(chunks, _)| chunks) + })?; // Assemble output let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?; @@ -878,10 +906,11 @@ pub fn read_chunked_data_cached( }; // Chunks stored as-is (no pipeline, or the filter mask says this chunk - // skipped it) are copied straight from the file bytes: they are already in - // memory, so routing them through a Vec and then an aligned cache buffer - // was two extra copies of the whole dataset for nothing. - let stored_raw = |c: &ChunkInfo| pipeline.is_none() || c.filter_mask != 0; + // skipped every filter) are copied straight from the file bytes: they are + // already in memory, so routing them through a Vec and then an aligned + // cache buffer was two extra copies of the whole dataset for nothing. + let stored_raw = + |c: &ChunkInfo| pipeline.is_none_or(|pl| all_filters_skipped(pl, c.filter_mask)); let mut misses: Vec<&ChunkInfo> = Vec::new(); for chunk_info in &chunks { if stored_raw(chunk_info) { @@ -889,7 +918,7 @@ pub fn read_chunked_data_cached( continue; } let coord: Vec = chunk_info.offsets.iter().take(rank).copied().collect(); - match cache.get_decompressed_aligned(&coord) { + match cache.get_decompressed_in(addr, &coord) { Some(cached) => place(&cached, chunk_info), None => misses.push(chunk_info), } @@ -903,7 +932,13 @@ pub fn read_chunked_data_cached( let cache_them = total_bytes <= cache.max_bytes(); if let Some(pl) = pipeline { let decode = |c: &&ChunkInfo| -> Result, FormatError> { - decompress_chunk(raw_bytes(c)?, pl, chunk_total_bytes, elem_size as u32) + decompress_chunk_masked( + raw_bytes(c)?, + pl, + chunk_total_bytes, + elem_size as u32, + c.filter_mask, + ) }; for batch in misses.chunks(DECODE_BATCH) { #[cfg(feature = "parallel")] @@ -920,7 +955,7 @@ pub fn read_chunked_data_cached( let data = data?; if cache_them { let coord: Vec = chunk_info.offsets.iter().take(rank).copied().collect(); - let cached = cache.put_decompressed(coord, data); + let cached = cache.put_decompressed_in(addr, coord, data); place(&cached, chunk_info); } else { place(&data, chunk_info); @@ -1124,24 +1159,20 @@ pub fn read_chunked_data_sweep( ))); } - // The per-file cache is shared across datasets; bind it to this one so a - // different dataset's chunk index is never reused for this read. - cache.ensure_dataset(addr); - - // Populate chunk index on first access - if !cache.has_index() { - let (chunks, _) = list_chunks( + // The per-file cache is shared across datasets (and threads); every + // lookup is keyed by this dataset's chunk-index address, so another + // dataset's index or chunks are never used for this read. + let chunks = cache.chunks_for(addr, rank, || { + list_chunks( file_data, layout, dataspace, elem_size, offset_size, length_size, - )?; - cache.populate_index(&chunks, rank); - } - - let chunks = cache.all_indexed_chunks().unwrap_or_default(); + ) + .map(|(chunks, _)| chunks) + })?; // Assemble output let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?; @@ -1172,12 +1203,12 @@ pub fn read_chunked_data_sweep( // Issue prefetch hint for predicted next chunks if !sweep.predicted_next.is_empty() { - cache.prefetch_hint(&sweep.predicted_next); + cache.prefetch_hint_in(addr, &sweep.predicted_next); cache.set_sweep_direction(sweep.direction); } // Try decompressed cache first - let decompressed = if let Some(cached) = cache.get_decompressed_aligned(&coord) { + let decompressed = if let Some(cached) = cache.get_decompressed_in(addr, &coord) { cached } else { // Decompress from file @@ -1186,15 +1217,17 @@ pub fn read_chunked_data_sweep( ensure_len(file_data, c_addr, size)?; let raw_chunk = &file_data[c_addr..c_addr + size]; let dec = if let Some(pl) = pipeline { - if chunk_info.filter_mask == 0 { - decompress_chunk(raw_chunk, pl, chunk_total_bytes, elem_size as u32)? - } else { - raw_chunk.to_vec() - } + decompress_chunk_masked( + raw_chunk, + pl, + chunk_total_bytes, + elem_size as u32, + chunk_info.filter_mask, + )? } else { raw_chunk.to_vec() }; - cache.put_decompressed(coord, dec) + cache.put_decompressed_in(addr, coord, dec) }; let chunk_offsets: Vec = chunk_info @@ -1278,48 +1311,34 @@ pub fn read_chunked_data_indexed( ))); } - // The per-file cache is shared across datasets; bind it to this one so a - // different dataset's chunk index is never reused for this read. - cache.ensure_dataset(addr); - - // Build chunk index on first access - if !cache.has_chunk_index() { - let (chunks, _) = list_chunks( - file_data, - layout, - dataspace, - elem_size, - offset_size, - length_size, - )?; - cache.populate_chunk_index(&chunks, rank); - // Also populate the legacy index for compatibility - if !cache.has_index() { - cache.populate_index(&chunks, rank); - } - } - - // Build chunk layout on first access - if !cache.has_chunk_layout() { - cache.populate_chunk_layout(&ds_dims, &chunk_dims, elem_size); - } - - // Get the layout info (mappings, output size, chunk total bytes) - let (mappings_info, output_bytes, chunk_total_bytes) = cache - .with_chunk_layout(|layout| { - let info: Vec<_> = layout - .mappings - .iter() - .map(|m| (m.coord.clone(), m.file_offset, m.file_size, m.filter_mask)) - .collect(); - (info, layout.output_bytes, layout.chunk_total_bytes) - }) - .ok_or_else(|| FormatError::ChunkedReadError("chunk layout not available".into()))?; + // Chunk index and assembly plan for this dataset, built on first access + // and kept per dataset (keyed by chunk-index address) in the shared cache. + let plan = cache.chunk_layout_for( + addr, + rank, + || { + list_chunks( + file_data, + layout, + dataspace, + elem_size, + offset_size, + length_size, + ) + .map(|(chunks, _)| chunks) + }, + &ds_dims, + &chunk_dims, + elem_size, + )?; + let chunk_total_bytes = plan.chunk_total_bytes; // Decompress chunks (using LRU cache where possible) - let mut chunk_buffers: Vec> = Vec::with_capacity(mappings_info.len()); - for (coord, file_offset, file_size, filter_mask) in &mappings_info { - if let Some(cached) = cache.get_decompressed_aligned(coord) { + let mut chunk_buffers: Vec> = Vec::with_capacity(plan.mappings.len()); + for m in &plan.mappings { + let (coord, file_offset, file_size, filter_mask) = + (&m.coord, &m.file_offset, &m.file_size, &m.filter_mask); + if let Some(cached) = cache.get_decompressed_in(addr, coord) { chunk_buffers.push(cached); } else { let c_addr = *file_offset as usize; @@ -1327,26 +1346,26 @@ pub fn read_chunked_data_indexed( ensure_len(file_data, c_addr, size)?; let raw_chunk = &file_data[c_addr..c_addr + size]; let decompressed = if let Some(pl) = pipeline { - if *filter_mask == 0 { - decompress_chunk(raw_chunk, pl, chunk_total_bytes, elem_size as u32)? - } else { - raw_chunk.to_vec() - } + decompress_chunk_masked( + raw_chunk, + pl, + chunk_total_bytes, + elem_size as u32, + *filter_mask, + )? } else { raw_chunk.to_vec() }; let aligned = CacheAlignedBuffer::from_vec(decompressed); - let arc = cache.put_decompressed_aligned(coord.clone(), aligned); + let arc = cache.put_decompressed_aligned_in(addr, coord.clone(), aligned); chunk_buffers.push(arc); } } // Assemble using pre-computed layout - let mut output = vec![0u8; output_bytes]; + let mut output = vec![0u8; plan.output_bytes]; let data_refs: Vec<&[u8]> = chunk_buffers.iter().map(|b| b.as_slice()).collect(); - cache.with_chunk_layout(|layout| { - layout.assemble(&data_refs, &mut output); - }); + plan.assemble(&data_refs, &mut output); Ok(output) } @@ -1594,7 +1613,8 @@ mod tests { } else { 0 }; - write_offset(&mut buf, off, offset_size); + // Key offsets are always 8 bytes (they are coordinates). + write_offset(&mut buf, off, 8); } // Child: address write_offset(&mut buf, chunk.address, offset_size); @@ -1604,7 +1624,7 @@ mod tests { buf.extend_from_slice(&0u32.to_le_bytes()); // chunk_size buf.extend_from_slice(&0u32.to_le_bytes()); // filter_mask for _ in 0..ndims { - write_offset(&mut buf, u64::MAX, offset_size); + write_offset(&mut buf, u64::MAX, 8); } buf @@ -1682,6 +1702,37 @@ mod tests { assert_eq!(result[2].address, 0x300); } + #[test] + fn collect_chunks_with_four_byte_addresses() { + // Sibling and child addresses are 4 bytes; the key offsets stay 8. + let ndims = 3; + let os: u8 = 4; + let chunks = vec![ + ChunkInfo { + chunk_size: 80, + filter_mask: 2, + offsets: vec![0, 5, 0], + address: 0x1000, + }, + ChunkInfo { + chunk_size: 96, + filter_mask: 0, + offsets: vec![8, 10, 0], + address: 0x2000, + }, + ]; + let btree = build_chunk_btree_leaf(&chunks, ndims, os); + assert_eq!(btree.len(), 8 + 2 * 4 + 2 * (8 + 3 * 8 + 4) + (8 + 3 * 8)); + let result = collect_chunk_info(&btree, 0, ndims, os, os).unwrap(); + assert_eq!(result.len(), 2); + for (got, want) in result.iter().zip(&chunks) { + assert_eq!(got.offsets, want.offsets); + assert_eq!(got.address, want.address); + assert_eq!(got.chunk_size, want.chunk_size); + assert_eq!(got.filter_mask, want.filter_mask); + } + } + #[test] fn collect_empty_btree() { let ndims = 2; @@ -1776,6 +1827,7 @@ mod tests { chunk_index_type: None, single_chunk_filtered_size: None, single_chunk_filter_mask: None, + dont_filter_partial_edge_chunks: false, }; let dataspace = Dataspace { @@ -1799,6 +1851,7 @@ mod tests { chunk_index_type: None, single_chunk_filtered_size: None, single_chunk_filter_mask: None, + dont_filter_partial_edge_chunks: false, }; let dataspace = Dataspace { space_type: DataspaceType::Simple, @@ -1956,6 +2009,7 @@ mod tests { chunk_index_type: None, single_chunk_filtered_size: None, single_chunk_filter_mask: None, + dont_filter_partial_edge_chunks: false, }; let dataspace = Dataspace { space_type: DataspaceType::Simple, @@ -2038,6 +2092,7 @@ mod tests { chunk_index_type: None, single_chunk_filtered_size: None, single_chunk_filter_mask: None, + dont_filter_partial_edge_chunks: false, }; let dataspace = Dataspace { space_type: DataspaceType::Simple, @@ -2200,6 +2255,7 @@ mod tests { chunk_index_type: Some(1), single_chunk_filtered_size: None, single_chunk_filter_mask: None, + dont_filter_partial_edge_chunks: false, }; let dataspace = Dataspace { space_type: DataspaceType::Simple, @@ -2229,12 +2285,12 @@ mod tests { let datatype = make_f64_type(); let cache = ChunkCache::new(); - assert!(!cache.has_index()); + assert_eq!(cache.indexed_dataset_count(), 0); let raw = read_chunked_data_cached( &file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache, ) .unwrap(); - assert!(cache.has_index()); + assert_eq!(cache.indexed_dataset_count(), 1); assert_eq!(raw.len(), 20 * 8); for i in 0..20 { let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap()); @@ -2256,7 +2312,7 @@ mod tests { &file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache, ) .unwrap(); - assert!(cache.has_index()); + assert_eq!(cache.indexed_dataset_count(), 1); assert_eq!(cache.cached_chunk_count(), 0); // Second read — reuses the cached index @@ -2265,6 +2321,7 @@ mod tests { ) .unwrap(); assert_eq!(raw1, raw2); + assert_eq!(cache.indexed_dataset_count(), 1); } #[test] diff --git a/crates/clawhdf5-format/src/chunked_write.rs b/crates/clawhdf5-format/src/chunked_write.rs index 682a1c6..2ebc902 100644 --- a/crates/clawhdf5-format/src/chunked_write.rs +++ b/crates/clawhdf5-format/src/chunked_write.rs @@ -1609,6 +1609,7 @@ mod tests { chunk_index_type, single_chunk_filtered_size, single_chunk_filter_mask, + .. } => { assert_eq!(version, 4); assert_eq!(chunk_index_type, Some(1)); diff --git a/crates/clawhdf5-format/src/data_layout.rs b/crates/clawhdf5-format/src/data_layout.rs index dd931ff..8429c5d 100644 --- a/crates/clawhdf5-format/src/data_layout.rs +++ b/crates/clawhdf5-format/src/data_layout.rs @@ -53,6 +53,11 @@ pub enum DataLayout { single_chunk_filtered_size: Option, /// Filter mask for v4 single chunk with filters. single_chunk_filter_mask: Option, + /// Layout v4 flag bit 0 (`H5D_CHUNK_DONT_FILTER_PARTIAL_CHUNKS`): + /// partial edge chunks — those extending past the dataset's current + /// extent in some dimension — are stored without the filter pipeline, + /// even though their filter mask is 0. Always `false` for v3. + dont_filter_partial_edge_chunks: bool, }, /// Virtual dataset layout (v4 only). Virtual { @@ -322,6 +327,7 @@ impl DataLayout { chunk_index_type: None, single_chunk_filtered_size: None, single_chunk_filter_mask: None, + dont_filter_partial_edge_chunks: false, }) } _ => Err(FormatError::InvalidLayoutClass(layout_class)), @@ -505,6 +511,7 @@ impl DataLayout { chunk_index_type: Some(chunk_index_type), single_chunk_filtered_size, single_chunk_filter_mask, + dont_filter_partial_edge_chunks: flags & 0x01 != 0, }) } 3 => { @@ -602,6 +609,7 @@ mod tests { chunk_index_type: None, single_chunk_filtered_size: None, single_chunk_filter_mask: None, + dont_filter_partial_edge_chunks: false, } ); } @@ -679,10 +687,35 @@ mod tests { chunk_index_type: Some(1), single_chunk_filtered_size: None, single_chunk_filter_mask: None, + dont_filter_partial_edge_chunks: false, } ); } + #[test] + fn v4_chunked_dont_filter_partial_edge_chunks_flag() { + let mut buf = vec![4u8, 2]; // version=4, class=2 + buf.push(0x01); // flags bit 0 = don't filter partial edge chunks + buf.push(2); // dimensionality=2 + buf.push(4); // dim_size_encoded_length=4 + buf.extend_from_slice(&5u32.to_le_bytes()); + buf.extend_from_slice(&4u32.to_le_bytes()); + buf.push(3); // Fixed Array + buf.push(10); // max_dblk_page_nelmts_bits + buf.extend_from_slice(&0x3000u64.to_le_bytes()); + match DataLayout::parse(&buf, 8, 8).unwrap() { + DataLayout::Chunked { + dont_filter_partial_edge_chunks, + btree_address, + .. + } => { + assert!(dont_filter_partial_edge_chunks); + assert_eq!(btree_address, Some(0x3000)); + } + other => panic!("expected Chunked, got {other:?}"), + } + } + #[test] fn v4_chunked_single_chunk_with_filters() { let mut buf = vec![4u8, 2]; // version=4, class=2 @@ -705,6 +738,7 @@ mod tests { chunk_index_type: Some(1), single_chunk_filtered_size: Some(1024), single_chunk_filter_mask: Some(0), + dont_filter_partial_edge_chunks: false, } ); } diff --git a/crates/clawhdf5-format/src/filters.rs b/crates/clawhdf5-format/src/filters.rs index c10cea9..ac17725 100644 --- a/crates/clawhdf5-format/src/filters.rs +++ b/crates/clawhdf5-format/src/filters.rs @@ -19,33 +19,99 @@ pub(crate) const MAX_DECOMPRESS_SIZE: usize = 256 * 1024 * 1024; /// Apply a filter pipeline to decompress a chunk. /// Filters are applied in REVERSE order for decompression. +/// +/// Equivalent to [`decompress_chunk_masked`] with a filter mask of 0 (every +/// filter was applied when the chunk was written). pub fn decompress_chunk( compressed: &[u8], pipeline: &FilterPipeline, chunk_size: usize, element_size: u32, ) -> Result, FormatError> { - let mut data = compressed.to_vec(); + decompress_chunk_masked(compressed, pipeline, chunk_size, element_size, 0) +} - for filter in pipeline.filters.iter().rev() { +/// Upper bound on the output of filter `filter_id` applied (in the write +/// direction) to `input` bytes. 0 means "unknown" and stays unknown. +/// +/// Shuffle preserves the size and Fletcher32 appends a 4-byte checksum. Any +/// other filter is a codec whose output can exceed its input on +/// incompressible data (deflate's stored blocks, LZ4's and zstd's literal +/// runs, codec headers); `n + n/8 + 64` covers every supported codec's worst +/// case while still bounding a decompression bomb to a small multiple of the +/// chunk. +fn filter_output_bound(filter_id: u16, input: usize) -> usize { + if input == 0 { + return 0; + } + match filter_id { + FILTER_SHUFFLE => input, + FILTER_FLETCHER32 => input.saturating_add(4), + _ => input.saturating_add(input / 8).saturating_add(64), + } +} + +/// Whether bit `index` of a chunk's filter mask says filter `index` was +/// skipped when the chunk was written. +fn filter_skipped(filter_mask: u32, index: usize) -> bool { + index < 32 && filter_mask & (1u32 << index) != 0 +} + +/// Whether `filter_mask` says none of `pipeline`'s filters were applied, so +/// the stored bytes are the chunk itself. +pub fn all_filters_skipped(pipeline: &FilterPipeline, filter_mask: u32) -> bool { + (0..pipeline.filters.len()).all(|i| filter_skipped(filter_mask, i)) +} + +/// Decompress a chunk whose filter mask is `filter_mask`: bit *i* set means +/// filter *i* of the pipeline was not applied when the chunk was written (an +/// optional filter that declined, or a direct chunk write), so only that +/// filter is skipped here; the others are still undone, in reverse order. +/// +/// `chunk_size` is the chunk's decoded size (0 if unknown). Each stage's +/// output is capped at what the filters before it (in write order) can have +/// produced from `chunk_size` bytes — e.g. a Fletcher32 checksum placed +/// before deflate (NetCDF-4's ordering) makes deflate's output 4 bytes +/// larger than the chunk — so the decompression-bomb limit stays tight +/// without rejecting valid pipelines. +pub fn decompress_chunk_masked( + compressed: &[u8], + pipeline: &FilterPipeline, + chunk_size: usize, + element_size: u32, + filter_mask: u32, +) -> Result, FormatError> { + // bounds[i]: the most bytes that entered filter i on the write side, and + // so the most that undoing filter i may produce. + let mut bounds = Vec::with_capacity(pipeline.filters.len()); + let mut size = chunk_size; + for (i, filter) in pipeline.filters.iter().enumerate() { + bounds.push(size); + if !filter_skipped(filter_mask, i) { + size = filter_output_bound(filter.filter_id, size); + } + } + + let mut data = compressed.to_vec(); + for (i, filter) in pipeline.filters.iter().enumerate().rev() { + if filter_skipped(filter_mask, i) { + continue; + } + let bound = bounds[i]; data = match filter.filter_id { FILTER_SHUFFLE => shuffle_decompress(&data, element_size as usize)?, - // `chunk_size` is the expected decompressed size (shuffle/fletcher32 - // are size-preserving, so it bounds these too); pass it so these - // decoders can't be forced into unbounded allocation by a hostile - // or corrupted compressed payload. - FILTER_DEFLATE => deflate_decompress(&data, chunk_size)?, - FILTER_LZ4 => lz4_decompress(&data, chunk_size)?, - FILTER_ZSTD => zstd_decompress(&data, chunk_size)?, + // `bound` caps the decoded size so these decoders can't be forced + // into unbounded allocation by a hostile or corrupted payload. + FILTER_DEFLATE => deflate_decompress(&data, bound)?, + FILTER_LZ4 => lz4_decompress(&data, bound)?, + FILTER_ZSTD => zstd_decompress(&data, bound)?, FILTER_FLETCHER32 => fletcher32_verify(&data)?, - FILTER_PCODEC => pcodec_decompress(&data, element_size as usize, chunk_size)?, - // `chunk_size` is the expected decompressed size; pass it so these - // decoders can reject an element count that would over-allocate. - FILTER_SCALEOFFSET => scaleoffset_decompress(&data, &filter.client_data, chunk_size)?, - FILTER_NBIT => nbit_decompress(&data, &filter.client_data, chunk_size)?, - FILTER_SZIP => { - crate::filters_szip::szip_decompress(&data, &filter.client_data, chunk_size)? - } + FILTER_PCODEC => pcodec_decompress(&data, element_size as usize, bound)?, + // These decoders also reject an element count that would + // over-allocate past `bound`. + FILTER_SCALEOFFSET => scaleoffset_decompress(&data, &filter.client_data, bound)?, + FILTER_NBIT => nbit_decompress(&data, &filter.client_data, bound)?, + FILTER_SZIP => crate::filters_szip::szip_decompress(&data, &filter.client_data, bound)?, other => return Err(FormatError::UnsupportedFilter(other)), }; } @@ -913,13 +979,13 @@ fn shuffle_decompress(data: &[u8], element_size: usize) -> Result, Forma if element_size <= 1 { return Ok(data.to_vec()); } - if !data.len().is_multiple_of(element_size) { - return Err(FormatError::FilterError( - "shuffle: data length not a multiple of element size".into(), - )); - } + // Like libhdf5, only whole elements are shuffled; trailing bytes (e.g. a + // Fletcher32 checksum appended before the shuffle) are stored as-is. + let whole = data.len() - data.len() % element_size; + let (data, tail) = data.split_at(whole); let num_elements = data.len() / element_size; - let mut result = vec![0u8; data.len()]; + let mut result = vec![0u8; whole]; + result.reserve_exact(tail.len()); // The shuffled stream is `element_size` byte planes of `num_elements` // bytes each; un-shuffling interleaves them. This is on the read path of @@ -950,6 +1016,7 @@ fn shuffle_decompress(data: &[u8], element_size: usize) -> Result, Forma } } } + result.extend_from_slice(tail); Ok(result) } @@ -965,19 +1032,19 @@ fn shuffle_compress(data: &[u8], element_size: usize) -> Result, FormatE if element_size <= 1 { return Ok(data.to_vec()); } - if !data.len().is_multiple_of(element_size) { - return Err(FormatError::FilterError( - "shuffle: data length not a multiple of element size".into(), - )); - } + // Trailing bytes that don't make a whole element are left in place, as + // libhdf5 does. + let whole = data.len() - data.len() % element_size; + let (data, tail) = data.split_at(whole); let num_elements = data.len() / element_size; - let mut result = vec![0u8; data.len()]; + let mut result = vec![0u8; whole]; match element_size { 4 => shuffle_compress_4(data, num_elements, &mut result), 8 => shuffle_compress_general(data, num_elements, element_size, &mut result), _ => shuffle_compress_general(data, num_elements, element_size, &mut result), } + result.extend_from_slice(tail); Ok(result) } @@ -1413,6 +1480,110 @@ mod tests { assert_eq!(decompressed, data); } + fn filter(filter_id: u16) -> FilterDescription { + FilterDescription { + filter_id, + name: None, + flags: 0, + client_data: vec![], + } + } + + #[test] + #[cfg(feature = "deflate")] + fn filter_mask_skips_only_the_masked_filters() { + let pipeline = FilterPipeline { + version: 2, + filters: vec![filter(FILTER_SHUFFLE), filter(FILTER_DEFLATE)], + }; + let only_shuffle = FilterPipeline { + version: 2, + filters: vec![filter(FILTER_SHUFFLE)], + }; + let only_deflate = FilterPipeline { + version: 2, + filters: vec![filter(FILTER_DEFLATE)], + }; + let data: Vec = (0..200).map(|i| (i * 7 % 256) as u8).collect(); + let n = data.len(); + + let shuffled = compress_chunk(&data, &only_shuffle, 8).unwrap(); + assert_ne!(shuffled, data); + let deflated = compress_chunk(&data, &only_deflate, 8).unwrap(); + + // Bit 1: deflate skipped, shuffle still undone. + assert_eq!( + decompress_chunk_masked(&shuffled, &pipeline, n, 8, 0b10).unwrap(), + data + ); + // Bit 0: shuffle skipped, deflate still undone. + assert_eq!( + decompress_chunk_masked(&deflated, &pipeline, n, 8, 0b01).unwrap(), + data + ); + // Both bits (and bits past the pipeline): stored as-is. + assert_eq!( + decompress_chunk_masked(&data, &pipeline, n, 8, u32::MAX).unwrap(), + data + ); + assert!(all_filters_skipped(&pipeline, 0b11)); + assert!(!all_filters_skipped(&pipeline, 0b10)); + // An unsupported filter is fine when the chunk skipped it. + let unknown = FilterPipeline { + version: 2, + filters: vec![filter(32000), filter(FILTER_DEFLATE)], + }; + assert_eq!( + decompress_chunk_masked(&deflated, &unknown, n, 8, 0b01).unwrap(), + data + ); + } + + #[test] + #[cfg(feature = "deflate")] + fn fletcher32_ahead_of_deflate_stays_bounded() { + // NetCDF-4 order: the checksum is appended before shuffle and deflate, + // so deflate decodes chunk + 4 bytes. + let pipeline = FilterPipeline { + version: 2, + filters: vec![ + filter(FILTER_FLETCHER32), + filter(FILTER_SHUFFLE), + filter(FILTER_DEFLATE), + ], + }; + let data: Vec = (0..400).map(|i| (i * 13 % 251) as u8).collect(); + let n = data.len(); + let stored = compress_chunk(&data, &pipeline, 8).unwrap(); + assert_eq!(decompress_chunk(&stored, &pipeline, n, 8).unwrap(), data); + + // The cap still bites: a stream that inflates past chunk + 4 bytes + // is rejected rather than allocated. + let only_deflate = FilterPipeline { + version: 2, + filters: vec![filter(FILTER_DEFLATE)], + }; + let oversized = compress_chunk(&vec![0u8; n + 5], &only_deflate, 8).unwrap(); + let err = decompress_chunk(&oversized, &pipeline, n, 8).unwrap_err(); + assert!( + matches!(err, FormatError::DecompressionError(_)), + "expected a size-limit error, got {err:?}" + ); + // A bomb is still stopped near the chunk size. + let bomb = compress_chunk(&vec![0u8; 64 * n], &only_deflate, 8).unwrap(); + assert!(decompress_chunk(&bomb, &pipeline, n, 8).is_err()); + } + + #[test] + fn shuffle_leaves_a_partial_trailing_element_in_place() { + // libhdf5 shuffles whole elements and copies the remainder as-is. + let data: Vec = (0..20).collect(); + let shuffled = shuffle_compress(&data, 8).unwrap(); + assert_eq!(&shuffled[16..], &data[16..]); + assert_eq!(&shuffled[..4], &[0, 8, 1, 9]); + assert_eq!(shuffle_decompress(&shuffled, 8).unwrap(), data); + } + #[test] #[cfg(feature = "deflate")] fn pipeline_compress_decompress_roundtrip() { diff --git a/crates/clawhdf5-format/src/parallel_read.rs b/crates/clawhdf5-format/src/parallel_read.rs index 14f5613..0bb2785 100644 --- a/crates/clawhdf5-format/src/parallel_read.rs +++ b/crates/clawhdf5-format/src/parallel_read.rs @@ -10,7 +10,7 @@ use crate::chunked_read::ChunkInfo; use crate::error::FormatError; use crate::filter_pipeline::FilterPipeline; -use crate::filters::decompress_chunk; +use crate::filters::decompress_chunk_masked; use crate::lane_partition::{self, LaneStats, PartitionStats}; /// Threshold: only use parallel decompression when chunk count exceeds this. @@ -84,11 +84,13 @@ pub fn decompress_chunks_lane_partitioned( } let raw_chunk = &file_data[c_addr..c_addr + size]; - let decompressed = if chunk_info.filter_mask == 0 { - decompress_chunk(raw_chunk, pipeline, chunk_total_bytes, element_size)? - } else { - raw_chunk.to_vec() - }; + let decompressed = decompress_chunk_masked( + raw_chunk, + pipeline, + chunk_total_bytes, + element_size, + chunk_info.filter_mask, + )?; stats.chunks_processed += 1; stats.compressed_bytes += size as u64; @@ -158,11 +160,13 @@ pub fn decompress_chunks_parallel( } let raw_chunk = &file_data[c_addr..c_addr + size]; - let decompressed = if chunk_info.filter_mask == 0 { - decompress_chunk(raw_chunk, pipeline, chunk_total_bytes, element_size)? - } else { - raw_chunk.to_vec() - }; + let decompressed = decompress_chunk_masked( + raw_chunk, + pipeline, + chunk_total_bytes, + element_size, + chunk_info.filter_mask, + )?; Ok(DecompressedChunk { index, @@ -200,11 +204,13 @@ pub fn decompress_chunks_sequential( let raw_chunk = &file_data[c_addr..c_addr + size]; let decompressed = if let Some(pl) = pipeline { - if chunk_info.filter_mask == 0 { - decompress_chunk(raw_chunk, pl, chunk_total_bytes, element_size)? - } else { - raw_chunk.to_vec() - } + decompress_chunk_masked( + raw_chunk, + pl, + chunk_total_bytes, + element_size, + chunk_info.filter_mask, + )? } else { raw_chunk.to_vec() }; diff --git a/crates/clawhdf5-format/src/partial_read.rs b/crates/clawhdf5-format/src/partial_read.rs index 7cdd2de..5cc3aaa 100644 --- a/crates/clawhdf5-format/src/partial_read.rs +++ b/crates/clawhdf5-format/src/partial_read.rs @@ -22,7 +22,7 @@ use crate::data_read::extract_selection_from_buffer; use crate::dataspace::Dataspace; use crate::error::FormatError; use crate::filter_pipeline::FilterPipeline; -use crate::filters::decompress_chunk; +use crate::filters::{all_filters_skipped, decompress_chunk_masked}; use crate::selection::Selection; /// The smallest axis-aligned box containing every selected element, as @@ -325,12 +325,18 @@ pub fn read_selection( expected: at.saturating_add(chunk.chunk_size as usize), available: file_data.len(), })?; - // Mirrors the full-read path: a non-zero filter mask means the - // chunk was stored unfiltered. + // Mirrors the full-read path: filter-mask bit i set means + // filter i was not applied to this chunk. let decoded; let data: &[u8] = match pipeline { - Some(pl) if chunk.filter_mask == 0 => { - decoded = decompress_chunk(raw, pl, chunk_bytes, elem_size as u32)?; + Some(pl) if !all_filters_skipped(pl, chunk.filter_mask) => { + decoded = decompress_chunk_masked( + raw, + pl, + chunk_bytes, + elem_size as u32, + chunk.filter_mask, + )?; &decoded } _ => raw, diff --git a/crates/clawhdf5/tests/concurrent_chunk_cache.rs b/crates/clawhdf5/tests/concurrent_chunk_cache.rs new file mode 100644 index 0000000..275e2eb --- /dev/null +++ b/crates/clawhdf5/tests/concurrent_chunk_cache.rs @@ -0,0 +1,87 @@ +//! A `File` is `Send + Sync` and keeps one chunk cache for all its datasets. +//! Threads reading different chunked datasets through the same `File` must +//! each get their own dataset's data. + +use std::sync::Arc; + +use clawhdf5::{File, FileBuilder}; + +const DATASETS: usize = 24; +const THREADS: usize = 16; +const ROUNDS: usize = 40; + +/// Contents of dataset `k`: distinct from every other dataset's, element for +/// element, so any chunk served from the wrong dataset shows. +fn values(k: usize, n: usize) -> Vec { + (0..n).map(|i| (k * 100_000 + i) as f64).collect() +} + +fn build() -> File { + let mut b = FileBuilder::new(); + for k in 0..DATASETS { + let ds = b.create_dataset(&format!("d{k:02}")); + match k % 3 { + // 1-D, compressed: chunk offsets 0, 8, 16, ... in every dataset. + 0 => { + ds.with_f64_data(&values(k, 64)).with_shape(&[64]); + ds.with_chunks(&[8]).with_deflate(1); + } + // 1-D, shuffle + compressed, a different length. + 1 => { + ds.with_f64_data(&values(k, 40)).with_shape(&[40]); + ds.with_chunks(&[8]).with_shuffle().with_deflate(1); + } + // 2-D, compressed: coordinates (0,0), (0,4), (4,0), ... overlap + // the other datasets' in the first dimension. + _ => { + ds.with_f64_data(&values(k, 64)).with_shape(&[8, 8]); + ds.with_chunks(&[4, 4]).with_deflate(1); + } + } + } + File::from_bytes(b.finish().unwrap()).unwrap() +} + +fn expected(k: usize) -> Vec { + values(k, if k % 3 == 1 { 40 } else { 64 }) +} + +#[test] +fn threads_reading_different_datasets_get_their_own_chunks() { + let file = Arc::new(build()); + // Sequential sanity check first. + for k in 0..DATASETS { + let got = file.dataset(&format!("d{k:02}")).unwrap().read_f64(); + assert_eq!(got.unwrap(), expected(k), "sequential d{k:02}"); + } + + let handles: Vec<_> = (0..THREADS) + .map(|t| { + let file = Arc::clone(&file); + std::thread::spawn(move || { + let mut wrong = Vec::new(); + for round in 0..ROUNDS { + let k = (t * 7 + round * 5) % DATASETS; + let name = format!("d{k:02}"); + match file.dataset(&name).unwrap().read_f64() { + Ok(v) if v == expected(k) => {} + Ok(v) => wrong.push(format!("{name}: wrong data, first {:?}", &v[..4])), + Err(e) => wrong.push(format!("{name}: {e}")), + } + } + wrong + }) + }) + .collect(); + let failures: Vec = handles + .into_iter() + .flat_map(|h| h.join().unwrap()) + .collect(); + assert!( + failures.is_empty(), + "{} of {} concurrent reads were wrong, e.g. {:?}", + failures.len(), + THREADS * ROUNDS, + &failures[..failures.len().min(5)] + ); +} diff --git a/crates/clawhdf5/tests/fixtures/h5fc_edge_v3.h5 b/crates/clawhdf5/tests/fixtures/h5fc_edge_v3.h5 new file mode 100644 index 0000000..6f92057 Binary files /dev/null and b/crates/clawhdf5/tests/fixtures/h5fc_edge_v3.h5 differ diff --git a/crates/clawhdf5/tests/h5py_chunked_read_tests.rs b/crates/clawhdf5/tests/h5py_chunked_read_tests.rs new file mode 100644 index 0000000..2b9e902 --- /dev/null +++ b/crates/clawhdf5/tests/h5py_chunked_read_tests.rs @@ -0,0 +1,345 @@ +//! Chunked-read regressions against files written by h5py / libhdf5. +//! +//! Each test builds its input with h5py (or uses a small committed fixture +//! when h5py cannot produce the feature) and compares clawhdf5's read with the +//! known contents. Tests are skipped if python3 with h5py is not available, +//! unless `CLAWHDF5_REQUIRE_INTEROP=1`. + +use std::process::Command; + +use clawhdf5::File; +use clawhdf5_format::selection::Selection; + +/// The Python interpreter to drive interop checks with (see +/// `h5py_interop_tests.rs`). +fn python() -> String { + std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string()) +} + +fn interop_required() -> bool { + std::env::var("CLAWHDF5_REQUIRE_INTEROP").is_ok_and(|v| v == "1") +} + +fn python_available() -> bool { + Command::new(python()) + .args(["-c", "import h5py, numpy"]) + .output() + .map(|o| o.status.success()) + .unwrap_or(false) +} + +macro_rules! skip_if_no_python { + () => { + if !python_available() { + assert!( + !interop_required(), + "CLAWHDF5_REQUIRE_INTEROP=1 but python3 with h5py is not available" + ); + eprintln!("SKIP: python3 with h5py not available"); + return; + } + }; +} + +fn run_python(script: &str) { + let output = Command::new(python()) + .args(["-c", script]) + .output() + .expect("failed to run python3"); + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + let stdout = String::from_utf8_lossy(&output.stdout); + panic!("Python script failed:\nSTDOUT: {stdout}\nSTDERR: {stderr}"); + } +} + +// --------------------------------------------------------------------------- +// Files with 4-byte addresses (superblock size-of-offsets = 4) +// --------------------------------------------------------------------------- + +/// The chunk B-tree (v1, type 1) stores each chunk offset in its keys as a +/// fixed 8-byte value whatever the file's size-of-offsets. Reading them with +/// the offset width misparsed every key in a 4-byte-offset file: unfiltered +/// datasets came back as zeros and filtered ones failed to inflate. +#[test] +fn h5py_four_byte_offsets_chunked_reads() { + skip_if_no_python!(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("sizes_4.h5"); + let p = path.display().to_string(); + run_python(&format!( + r#" +import h5py, numpy as np +for name, lengths in (("{p}", 4), ("{p}.l8", 8)): + fcpl = h5py.h5p.create(h5py.h5p.FILE_CREATE) + fcpl.set_sizes(4, lengths) + fid = h5py.h5f.create(name.encode(), h5py.h5f.ACC_TRUNC, fcpl=fcpl) + with h5py.File(fid) as f: + f.create_dataset("plain", data=np.arange(100.0), chunks=(10,)) + f.create_dataset("gzip", data=np.arange(100.0), chunks=(10,), compression="gzip") + f.create_dataset("grid", data=np.arange(35 * 13, dtype=" = (0..100).map(f64::from).collect(); + let grid: Vec = (0..35 * 13).collect(); + for name in [p.clone(), format!("{p}.l8")] { + let file = File::open(&name).unwrap(); + for ds in ["plain", "gzip"] { + assert_eq!( + file.dataset(ds).unwrap().read_f64().unwrap(), + expect, + "{name}:{ds}" + ); + } + for ds in ["grid", "grid_gzip"] { + assert_eq!( + file.dataset(ds).unwrap().read_i32().unwrap(), + grid, + "{name}:{ds}" + ); + } + } +} + +// --------------------------------------------------------------------------- +// Per-chunk filter masks +// --------------------------------------------------------------------------- + +/// A chunk's filter mask has one bit per pipeline filter: bit i set means +/// filter i was not applied to that chunk. Any nonzero mask used to skip the +/// whole pipeline, so a chunk that skipped only gzip was handed back still +/// shuffled. +#[test] +fn h5py_partial_filter_mask_skips_only_masked_filters() { + skip_if_no_python!(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("mask.h5"); + let p = path.display().to_string(); + run_python(&format!( + r#" +import h5py, numpy as np, zlib +def shuffle(b, es): + a = np.frombuffer(b, dtype=np.uint8).reshape(-1, es) + return a.T.copy().tobytes() +with h5py.File("{p}", "w") as f: + # shuffle (0) + gzip (1). Even chunks: both applied. Odd chunks: mask + # 0b10, gzip skipped, shuffle applied. Chunk 3: mask 0b11, raw. + ds = f.create_dataset("shuf_gzip", shape=(32,), chunks=(8,), dtype=" = (1000..1032).collect(); + assert_eq!(file.dataset("shuf_gzip").unwrap().read_i32().unwrap(), want); + let grid: Vec = (0..64).map(f64::from).collect(); + assert_eq!(file.dataset("grid").unwrap().read_f64().unwrap(), grid); + // The selection path decodes chunks on its own. + let part = file + .dataset("shuf_gzip") + .unwrap() + .read_selection(&Selection::Hyperslab { + start: vec![6], + stride: vec![1], + count: vec![1], + block: vec![20], + }) + .unwrap(); + let want: Vec = (1006..1026i32).flat_map(i32::to_le_bytes).collect(); + assert_eq!(part, want); +} + +// --------------------------------------------------------------------------- +// Size-changing filters ahead of a codec +// --------------------------------------------------------------------------- + +/// Fletcher32 placed before the compressor (NetCDF-4's ordering) makes the +/// codec's decoded output 4 bytes larger than the chunk. The decompression +/// cap was the chunk size for every stage, so these files failed with +/// "deflate: output exceeds size limit". +#[test] +fn h5py_fletcher32_before_deflate_reads() { + skip_if_no_python!(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("fletcher_first.h5"); + let p = path.display().to_string(); + run_python(&format!( + r#" +import h5py, numpy as np +arr = np.sin(np.arange(5000) / 50.0) +grid = np.arange(37 * 21, dtype=" = (0..5000).map(|i| (f64::from(i) / 50.0).sin()).collect(); + for name in ["fl_shuf_gzip", "fl_gzip", "shuf_fl_gzip"] { + let values = file.dataset(name).unwrap().read_f64().unwrap(); + assert_eq!(values.len(), arr.len(), "{name}"); + for (i, (v, w)) in values.iter().zip(&arr).enumerate() { + assert!((v - w).abs() < 1e-12, "{name}[{i}]: {v} vs {w}"); + } + } + let grid: Vec = (0..37 * 21).collect(); + assert_eq!( + file.dataset("grid_fl_shuf_gzip") + .unwrap() + .read_i32() + .unwrap(), + grid + ); +} + +// --------------------------------------------------------------------------- +// "Don't filter partial edge chunks" +// --------------------------------------------------------------------------- + +/// libhdf5's own test file for `H5Pset_chunk_opts(H5D_CHUNK_DONT_FILTER_PARTIAL_CHUNKS)`: +/// a 12x6 f32 dataset in 5x5 gzip chunks whose edge chunks are stored raw +/// with a filter mask of 0. Reading it tried to inflate the raw edge chunks +/// ("deflate: ... unknown compression method"). +#[test] +fn libhdf5_edge_chunk_fixture_reads() { + let path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/h5fc_edge_v3.h5" + ); + let file = File::open(path).unwrap(); + let ds = file.dataset("DSET_EDGE").unwrap(); + assert_eq!(ds.shape().unwrap(), vec![12, 6]); + assert_eq!(ds.read_f32().unwrap(), vec![100.0f32; 72]); +} + +/// The same layout flag on datasets with varied contents, set through the +/// libhdf5 that h5py ships (h5py has no binding for `H5Pset_chunk_opts`). +#[test] +fn h5py_unfiltered_partial_edge_chunks_read() { + skip_if_no_python!(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("edge.h5"); + let p = path.display().to_string(); + let script = format!( + r#" +import ctypes, glob, os, sys +import h5py, numpy as np +here = os.path.dirname(h5py.__file__) +libs = glob.glob(os.path.join(here, "..", "h5py.libs", "libhdf5-*.so*")) +libs += glob.glob(os.path.join(here, ".dylibs", "libhdf5*.dylib")) +if not libs: + print("NO_LIBHDF5") + sys.exit(0) +lib = ctypes.CDLL(libs[0]) +lib.H5Pset_chunk_opts.argtypes = [ctypes.c_int64, ctypes.c_uint] +def make(f, name, data, chunk, maxshape, shuffle): + dcpl = h5py.h5p.create(h5py.h5p.DATASET_CREATE) + dcpl.set_chunk(chunk) + if shuffle: + dcpl.set_shuffle() + dcpl.set_deflate(6) + assert lib.H5Pset_chunk_opts(dcpl.id, 0x0002) >= 0 + space = h5py.h5s.create_simple(data.shape, maxshape) + d = h5py.h5d.create(f.id, name.encode(), h5py.h5t.py_create(data.dtype), space, dcpl=dcpl) + d.write(h5py.h5s.ALL, h5py.h5s.ALL, np.ascontiguousarray(data)) +with h5py.File("{p}", "w") as f: + line = np.sin(np.arange(1000) / 7.0) + grid = np.arange(37 * 53, dtype=" = (0..1000).map(|i| (f64::from(i) / 7.0).sin()).collect(); + for name in ["fixed_1d", "ea_1d"] { + let values = file.dataset(name).unwrap().read_f64().unwrap(); + assert_eq!(values.len(), line.len(), "{name}"); + for (i, (v, w)) in values.iter().zip(&line).enumerate() { + assert!((v - w).abs() < 1e-12, "{name}[{i}]: {v} vs {w}"); + } + } + let grid: Vec = (0..37 * 53).map(|i| f64::from(i) * 0.5).collect(); + for name in ["bt2_2d", "fixed_2d"] { + assert_eq!( + file.dataset(name).unwrap().read_f64().unwrap(), + grid, + "{name}" + ); + } + // A selection touching only the last (partial, unfiltered) chunk. + let tail = file + .dataset("fixed_1d") + .unwrap() + .read_selection(&Selection::Hyperslab { + start: vec![990], + stride: vec![1], + count: vec![1], + block: vec![10], + }) + .unwrap(); + let want: Vec = line[990..].iter().flat_map(|v| v.to_le_bytes()).collect(); + assert_eq!(tail, want); +}