Fix silent wrong data and libhdf5 interop found by the HDF5 audit #11
@@ -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<CacheAlignedBuffer>,
|
||||
@@ -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<Arc<HashMap<ChunkCoord, ChunkInfo>>>,
|
||||
/// Pre-built chunk index for O(1) coordinate lookups.
|
||||
chunk_index: Option<Arc<ChunkIndex>>,
|
||||
/// Pre-computed chunk layout for fast assembly.
|
||||
chunk_layout: Option<Arc<ChunkLayout>>,
|
||||
/// 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<HashMap<ChunkCoord, ChunkInfo>>,
|
||||
/// Per-dataset chunk indexes, keyed by chunk-index address.
|
||||
datasets: HashMap<u64, DatasetEntry>,
|
||||
|
||||
/// 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<u64>,
|
||||
/// Dataset the address-less methods act on (see `ensure_dataset`).
|
||||
current: Option<u64>,
|
||||
|
||||
/// LRU cache of decompressed chunk data.
|
||||
slots: Vec<CachedChunk>,
|
||||
|
||||
/// 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<ChunkCoord, usize>,
|
||||
slot_index: HashMap<SlotKey, usize>,
|
||||
|
||||
/// 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<ChunkCoord>,
|
||||
/// Last accessed chunk (for sequential detection).
|
||||
last_coord: Option<SlotKey>,
|
||||
|
||||
/// Access pattern statistics.
|
||||
stats: AccessStats,
|
||||
}
|
||||
|
||||
/// Pre-built chunk index for O(1) coordinate lookups.
|
||||
chunk_index: Option<ChunkIndex>,
|
||||
#[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<ChunkLayout>,
|
||||
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<Arc<CacheAlignedBuffer>> {
|
||||
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<CacheAlignedBuffer>,
|
||||
) -> Arc<CacheAlignedBuffer> {
|
||||
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.
|
||||
/// 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<E>(
|
||||
&self,
|
||||
addr: u64,
|
||||
rank: usize,
|
||||
build: impl FnOnce() -> Result<Vec<ChunkInfo>, E>,
|
||||
) -> Result<Vec<ChunkInfo>, E> {
|
||||
Ok(self
|
||||
.index_for(addr, rank, build)?
|
||||
.values()
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn index_for<E>(
|
||||
&self,
|
||||
addr: u64,
|
||||
rank: usize,
|
||||
build: impl FnOnce() -> Result<Vec<ChunkInfo>, E>,
|
||||
) -> Result<Arc<HashMap<ChunkCoord, ChunkInfo>>, E> {
|
||||
if let Some(index) = self.lock().touch(addr).index.clone() {
|
||||
return Ok(index);
|
||||
}
|
||||
let chunks = build()?;
|
||||
let map: HashMap<ChunkCoord, ChunkInfo> = 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)
|
||||
}
|
||||
|
||||
/// 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<E>(
|
||||
&self,
|
||||
addr: u64,
|
||||
rank: usize,
|
||||
build: impl FnOnce() -> Result<Vec<ChunkInfo>, E>,
|
||||
ds_dims: &[usize],
|
||||
chunk_dims: &[usize],
|
||||
elem_size: usize,
|
||||
) -> Result<Arc<ChunkLayout>, 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<ChunkInfo> = 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<Arc<CacheAlignedBuffer>> {
|
||||
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<u8>,
|
||||
) -> Arc<CacheAlignedBuffer> {
|
||||
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<CacheAlignedBuffer> {
|
||||
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.inner.lock().unwrap_or_else(|e| e.into_inner());
|
||||
if inner.index_addr == Some(addr) {
|
||||
return false;
|
||||
}
|
||||
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 mut inner = self.lock();
|
||||
let changed = inner.current != Some(addr);
|
||||
inner.current = Some(addr);
|
||||
changed
|
||||
}
|
||||
|
||||
/// Returns `true` if the chunk index has been built.
|
||||
/// 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::<core::convert::Infallible>(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<ChunkInfo> {
|
||||
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<ChunkInfo>` (order unspecified).
|
||||
/// Return all of the bound dataset's indexed chunks (order unspecified).
|
||||
pub fn all_indexed_chunks(&self) -> Option<Vec<ChunkInfo>> {
|
||||
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<F, R>(&self, f: F) -> Option<R>
|
||||
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<u8>`; 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<Vec<u8>> {
|
||||
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<Arc<CacheAlignedBuffer>> {
|
||||
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<u8>) -> Arc<CacheAlignedBuffer> {
|
||||
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<CacheAlignedBuffer> {
|
||||
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;
|
||||
let mut inner = self.lock();
|
||||
let addr = inner.current();
|
||||
inner.put_decompressed((addr, coord), 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
|
||||
/// [`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);
|
||||
}
|
||||
|
||||
// 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;
|
||||
}
|
||||
// ----- Whole-cache operations -----
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
/// Clear the entire cache (index + decompressed data).
|
||||
/// 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<ChunkInfo> = (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();
|
||||
|
||||
@@ -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<u64> = 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<Vec<u8>, 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<u64> = 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<usize> = 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(
|
||||
// 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,
|
||||
)?;
|
||||
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()))?;
|
||||
let chunk_total_bytes = plan.chunk_total_bytes;
|
||||
|
||||
// Decompress chunks (using LRU cache where possible)
|
||||
let mut chunk_buffers: Vec<Arc<CacheAlignedBuffer>> = 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<Arc<CacheAlignedBuffer>> = 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]
|
||||
|
||||
@@ -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));
|
||||
|
||||
@@ -53,6 +53,11 @@ pub enum DataLayout {
|
||||
single_chunk_filtered_size: Option<u64>,
|
||||
/// Filter mask for v4 single chunk with filters.
|
||||
single_chunk_filter_mask: Option<u32>,
|
||||
/// 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,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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<Vec<u8>, 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<Vec<u8>, 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<Vec<u8>, 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<Vec<u8>, Forma
|
||||
}
|
||||
}
|
||||
}
|
||||
result.extend_from_slice(tail);
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
@@ -965,19 +1032,19 @@ fn shuffle_compress(data: &[u8], element_size: usize) -> Result<Vec<u8>, 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<u8> = (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<u8> = (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<u8> = (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() {
|
||||
|
||||
@@ -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()
|
||||
};
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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<f64> {
|
||||
(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<f64> {
|
||||
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<String> = 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)]
|
||||
);
|
||||
}
|
||||
BIN
Binary file not shown.
@@ -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="<i4").reshape(35, 13),
|
||||
chunks=(8, 5))
|
||||
f.create_dataset("grid_gzip", data=np.arange(35 * 13, dtype="<i4").reshape(35, 13),
|
||||
chunks=(8, 5), compression="gzip", shuffle=True)
|
||||
"#
|
||||
));
|
||||
|
||||
let expect: Vec<f64> = (0..100).map(f64::from).collect();
|
||||
let grid: Vec<i32> = (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="<i4",
|
||||
compression="gzip", shuffle=True)
|
||||
for i in range(4):
|
||||
raw = np.arange(1000 + i * 8, 1008 + i * 8, dtype="<i4").tobytes()
|
||||
if i == 3:
|
||||
ds.id.write_direct_chunk((i * 8,), raw, filter_mask=0b11)
|
||||
elif i % 2:
|
||||
ds.id.write_direct_chunk((i * 8,), shuffle(raw, 4), filter_mask=0b10)
|
||||
else:
|
||||
ds.id.write_direct_chunk((i * 8,), zlib.compress(shuffle(raw, 4)), filter_mask=0)
|
||||
# shuffle (0) + gzip (1) on a 2-D dataset, chunk (0, 1) skips shuffle only.
|
||||
ds = f.create_dataset("grid", shape=(8, 8), chunks=(4, 4), dtype="<f8",
|
||||
compression="gzip", shuffle=True)
|
||||
full = np.arange(64, dtype="<f8").reshape(8, 8)
|
||||
for r in (0, 4):
|
||||
for c in (0, 4):
|
||||
raw = np.ascontiguousarray(full[r:r + 4, c:c + 4]).tobytes()
|
||||
if (r, c) == (0, 4):
|
||||
ds.id.write_direct_chunk((r, c), zlib.compress(raw), filter_mask=0b01)
|
||||
else:
|
||||
ds.id.write_direct_chunk((r, c), zlib.compress(shuffle(raw, 8)), filter_mask=0)
|
||||
assert (f["grid"][...] == full).all()
|
||||
assert (f["shuf_gzip"][...] == np.arange(1000, 1032)).all()
|
||||
"#
|
||||
));
|
||||
|
||||
let file = File::open(&path).unwrap();
|
||||
let want: Vec<i32> = (1000..1032).collect();
|
||||
assert_eq!(file.dataset("shuf_gzip").unwrap().read_i32().unwrap(), want);
|
||||
let grid: Vec<f64> = (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<u8> = (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="<i4").reshape(37, 21)
|
||||
ids = {{"fletcher32": 3, "shuffle": 2, "deflate": 1}}
|
||||
with h5py.File("{p}", "w") as f:
|
||||
for name, steps, data, chunk in (
|
||||
("fl_shuf_gzip", ("fletcher32", "shuffle", "deflate"), arr, (500,)),
|
||||
("fl_gzip", ("fletcher32", "deflate"), arr, (500,)),
|
||||
("shuf_fl_gzip", ("shuffle", "fletcher32", "deflate"), arr, (500,)),
|
||||
("grid_fl_shuf_gzip", ("fletcher32", "shuffle", "deflate"), grid, (8, 5)),
|
||||
):
|
||||
dcpl = h5py.h5p.create(h5py.h5p.DATASET_CREATE)
|
||||
dcpl.set_chunk(chunk)
|
||||
for s in steps:
|
||||
if s == "deflate":
|
||||
dcpl.set_deflate(4)
|
||||
elif s == "shuffle":
|
||||
dcpl.set_shuffle()
|
||||
else:
|
||||
dcpl.set_fletcher32()
|
||||
tid = h5py.h5t.py_create(data.dtype)
|
||||
space = h5py.h5s.create_simple(data.shape)
|
||||
d = h5py.h5d.create(f.id, name.encode(), tid, space, dcpl=dcpl)
|
||||
d.write(h5py.h5s.ALL, h5py.h5s.ALL, np.ascontiguousarray(data))
|
||||
order = [d.get_create_plist().get_filter(i)[0] for i in range(len(steps))]
|
||||
assert order == [ids[s] for s in steps], order
|
||||
"#
|
||||
));
|
||||
|
||||
let file = File::open(&path).unwrap();
|
||||
let arr: Vec<f64> = (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<i32> = (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="<f8").reshape(37, 53) * 0.5
|
||||
make(f, "fixed_1d", line, (64,), None, False)
|
||||
make(f, "ea_1d", line, (64,), (h5py.h5s.UNLIMITED,), True)
|
||||
make(f, "bt2_2d", grid, (8, 8), (h5py.h5s.UNLIMITED,) * 2, True)
|
||||
make(f, "fixed_2d", grid, (8, 8), None, True)
|
||||
print("OK")
|
||||
"#
|
||||
);
|
||||
let out = Command::new(python())
|
||||
.args(["-c", &script])
|
||||
.output()
|
||||
.expect("failed to run python3");
|
||||
assert!(
|
||||
out.status.success(),
|
||||
"{}",
|
||||
String::from_utf8_lossy(&out.stderr)
|
||||
);
|
||||
if String::from_utf8_lossy(&out.stdout).contains("NO_LIBHDF5") {
|
||||
eprintln!("SKIP: h5py's bundled libhdf5 not found (needed for H5Pset_chunk_opts)");
|
||||
return;
|
||||
}
|
||||
|
||||
let file = File::open(&path).unwrap();
|
||||
let line: Vec<f64> = (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<f64> = (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<u8> = line[990..].iter().flat_map(|v| v.to_le_bytes()).collect();
|
||||
assert_eq!(tail, want);
|
||||
}
|
||||
Reference in New Issue
Block a user