From 9ea44d473d0a4bd8047fc081afed314561562605 Mon Sep 17 00:00:00 2001 From: osobh Date: Fri, 25 Sep 2026 21:01:57 -0500 Subject: [PATCH 1/5] fix(format): read v1 chunk B-tree key offsets as 8 bytes A type-1 (raw data chunk) B-tree key holds the chunk size, the filter mask and one offset per dimension, and those offsets are always 8 bytes: they are dataset coordinates, not file addresses. The reader used the superblock's size-of-offsets for them, so in a file with 4-byte offsets every key was misparsed. Unfiltered chunked datasets read as zeros (with stray bytes where a misread address landed on data) and filtered ones failed with "deflate: truncated stream". Only the sibling and child addresses follow size-of-offsets now. The unit-test B-tree builder wrote keys the same wrong way, which is why its tests passed; it now matches the format. Regression: h5py_four_byte_offsets_chunked_reads (h5py, set_sizes(4, 4) and (4, 8); 1-D and 2-D, unfiltered and gzip) and the unit test collect_chunks_with_four_byte_addresses. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/clawhdf5-format/src/chunked_read.rs | 54 ++++++++- .../clawhdf5/tests/h5py_chunked_read_tests.rs | 105 ++++++++++++++++++ 2 files changed, 153 insertions(+), 6 deletions(-) create mode 100644 crates/clawhdf5/tests/h5py_chunked_read_tests.rs diff --git a/crates/clawhdf5-format/src/chunked_read.rs b/crates/clawhdf5-format/src/chunked_read.rs index 8dade54..b0f5141 100644 --- a/crates/clawhdf5-format/src/chunked_read.rs +++ b/crates/clawhdf5-format/src/chunked_read.rs @@ -223,6 +223,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 +264,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 +297,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; @@ -1592,7 +1602,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); @@ -1602,7 +1613,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 @@ -1680,6 +1691,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; 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..d46cff5 --- /dev/null +++ b/crates/clawhdf5/tests/h5py_chunked_read_tests.rs @@ -0,0 +1,105 @@ +//! 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; + +/// 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}" + ); + } + } +} From 585e14d5e213bac5c4a271ece5136d6a0371f14f Mon Sep 17 00:00:00 2001 From: osobh Date: Fri, 25 Sep 2026 21:05:07 -0500 Subject: [PATCH 2/5] fix(format): honour each bit of a chunk's filter mask A chunk's filter mask has one bit per pipeline filter; bit i set means filter i was not applied to that chunk (an optional filter that declined, or a direct chunk write). Every read path treated any nonzero mask as "no filters applied" and returned the stored bytes, so a chunk that skipped only gzip in a shuffle+gzip pipeline came back still shuffled (h5py write_direct_chunk with filter_mask=0b10: 8 of 32 values wrong). decompress_chunk_masked undoes the filters the mask leaves set and skips the rest; an unsupported filter is no longer an error when the chunk skipped it. The full, cached, sweep, indexed, parallel and selection (partial_read) paths all use it, and a chunk is copied straight from the file only when every filter was skipped. decompress_chunk is the mask-0 case. Regression: h5py_partial_filter_mask_skips_only_masked_filters (1-D shuffle+gzip with masks 0, 0b10 and 0b11; 2-D with 0b01; full and hyperslab reads) and filter_mask_skips_only_the_masked_filters. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/clawhdf5-format/src/chunked_read.rs | 55 +++++--- crates/clawhdf5-format/src/filters.rs | 124 +++++++++++++++--- crates/clawhdf5-format/src/parallel_read.rs | 38 +++--- crates/clawhdf5-format/src/partial_read.rs | 16 ++- .../clawhdf5/tests/h5py_chunked_read_tests.rs | 70 ++++++++++ 5 files changed, 244 insertions(+), 59 deletions(-) diff --git a/crates/clawhdf5-format/src/chunked_read.rs b/crates/clawhdf5-format/src/chunked_read.rs index b0f5141..2dab147 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() }; @@ -886,10 +888,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) { @@ -911,7 +914,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")] @@ -1194,11 +1203,13 @@ 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() }; @@ -1335,11 +1346,13 @@ 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() }; diff --git a/crates/clawhdf5-format/src/filters.rs b/crates/clawhdf5-format/src/filters.rs index c10cea9..31c024a 100644 --- a/crates/clawhdf5-format/src/filters.rs +++ b/crates/clawhdf5-format/src/filters.rs @@ -19,33 +19,64 @@ 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() { +/// 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); it caps each +/// stage's output. +pub fn decompress_chunk_masked( + compressed: &[u8], + pipeline: &FilterPipeline, + chunk_size: usize, + element_size: u32, + filter_mask: u32, +) -> Result, FormatError> { + let mut data = compressed.to_vec(); + for (i, filter) in pipeline.filters.iter().enumerate().rev() { + if filter_skipped(filter_mask, i) { + continue; + } + let bound = chunk_size; 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)), }; } @@ -1413,6 +1444,65 @@ 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 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/h5py_chunked_read_tests.rs b/crates/clawhdf5/tests/h5py_chunked_read_tests.rs index d46cff5..58fe290 100644 --- a/crates/clawhdf5/tests/h5py_chunked_read_tests.rs +++ b/crates/clawhdf5/tests/h5py_chunked_read_tests.rs @@ -8,6 +8,7 @@ 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`). @@ -103,3 +104,72 @@ for name, lengths in (("{p}", 4), ("{p}.l8", 8)): } } } + +// --------------------------------------------------------------------------- +// 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); +} From e162c013fdd92c80e95e092ba42e7b9e5c2c2628 Mon Sep 17 00:00:00 2001 From: osobh Date: Fri, 25 Sep 2026 21:06:22 -0500 Subject: [PATCH 3/5] fix(format): bound each filter stage by what the stages before it produce Every decode stage was capped at the chunk's decoded size. That holds only when every filter ahead of the codec preserves size; Fletcher32 does not (it appends a 4-byte checksum), so a pipeline with Fletcher32 before deflate (NetCDF-4's fletcher32 -> shuffle -> deflate ordering, h5repack's "all filters") failed with "deflate: output exceeds size limit" on every chunk. decompress_chunk_masked now computes each stage's bound by running the chunk size forward through the filters that precede it in write order (and that the chunk's mask did not skip): shuffle keeps the size, Fletcher32 adds 4, any codec adds at most n/8 + 64. The cap is still a small constant factor of the chunk, so a decompression bomb is rejected as before (tested). Shuffle also had to learn libhdf5's handling of a length that is not a whole number of elements (chunk + checksum): shuffle the whole elements and leave the trailing bytes in place, in both directions. It used to refuse such data. Regression: h5py_fletcher32_before_deflate_reads (fletcher->shuffle-> gzip, fletcher->gzip, shuffle->fletcher->gzip, and a 2-D i32 grid), fletcher32_ahead_of_deflate_stays_bounded and shuffle_leaves_a_partial_trailing_element_in_place. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/clawhdf5-format/src/filters.rs | 111 +++++++++++++++--- .../clawhdf5/tests/h5py_chunked_read_tests.rs | 64 ++++++++++ 2 files changed, 160 insertions(+), 15 deletions(-) diff --git a/crates/clawhdf5-format/src/filters.rs b/crates/clawhdf5-format/src/filters.rs index 31c024a..ac17725 100644 --- a/crates/clawhdf5-format/src/filters.rs +++ b/crates/clawhdf5-format/src/filters.rs @@ -31,6 +31,26 @@ pub fn decompress_chunk( decompress_chunk_masked(compressed, pipeline, chunk_size, element_size, 0) } +/// 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 { @@ -48,8 +68,12 @@ pub fn all_filters_skipped(pipeline: &FilterPipeline, filter_mask: u32) -> bool /// 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); it caps each -/// stage's output. +/// `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, @@ -57,12 +81,23 @@ pub fn decompress_chunk_masked( 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 = chunk_size; + let bound = bounds[i]; data = match filter.filter_id { FILTER_SHUFFLE => shuffle_decompress(&data, element_size as usize)?, // `bound` caps the decoded size so these decoders can't be forced @@ -944,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 @@ -981,6 +1016,7 @@ fn shuffle_decompress(data: &[u8], element_size: usize) -> Result, Forma } } } + result.extend_from_slice(tail); Ok(result) } @@ -996,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) } @@ -1503,6 +1539,51 @@ mod tests { ); } + #[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/tests/h5py_chunked_read_tests.rs b/crates/clawhdf5/tests/h5py_chunked_read_tests.rs index 58fe290..deda24f 100644 --- a/crates/clawhdf5/tests/h5py_chunked_read_tests.rs +++ b/crates/clawhdf5/tests/h5py_chunked_read_tests.rs @@ -173,3 +173,67 @@ with h5py.File("{p}", "w") as f: 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 + ); +} From d074385944c8c49f1325a39de6208227e90e581b Mon Sep 17 00:00:00 2001 From: osobh Date: Fri, 25 Sep 2026 21:08:36 -0500 Subject: [PATCH 4/5] fix(format): read partial edge chunks stored unfiltered MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Layout message v4 flag bit 0 (H5D_CHUNK_DONT_FILTER_PARTIAL_CHUNKS, set with H5Pset_chunk_opts) makes libhdf5 store every chunk that extends past the dataset's extent without the filter pipeline, while its filter mask still reads 0. The parser ignored the flag, so readers tried to inflate raw bytes: libhdf5's own h5fc_edge_v3.h5 failed with "deflate: ... unknown compression method". DataLayout::Chunked gains dont_filter_partial_edge_chunks (always false for v3), and list_chunks — the one place every read path gets its chunk list from — marks such partial chunks as having skipped every filter, so the full, cached, indexed, parallel and selection readers all copy them as-is. chunked_write.rs gets `..` in one exhaustive test pattern for the new field. Regression: libhdf5_edge_chunk_fixture_reads (h5fc_edge_v3.h5 from the HDF5 tools test files, committed as a 2.5 KB fixture), and h5py_unfiltered_partial_edge_chunks_read (the flag set through h5py's bundled libhdf5 via ctypes, as h5py has no binding for it: fixed array, extensible array and B-tree v2 indexes, 1-D and 2-D, plus a hyperslab of the last chunk), and v4_chunked_dont_filter_partial_edge_chunks_flag. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/clawhdf5-format/src/chunked_read.rs | 27 ++++- crates/clawhdf5-format/src/chunked_write.rs | 1 + crates/clawhdf5-format/src/data_layout.rs | 34 ++++++ .../clawhdf5/tests/fixtures/h5fc_edge_v3.h5 | Bin 0 -> 2478 bytes .../clawhdf5/tests/h5py_chunked_read_tests.rs | 106 ++++++++++++++++++ 5 files changed, 167 insertions(+), 1 deletion(-) create mode 100644 crates/clawhdf5/tests/fixtures/h5fc_edge_v3.h5 diff --git a/crates/clawhdf5-format/src/chunked_read.rs b/crates/clawhdf5-format/src/chunked_read.rs index 2dab147..bc006df 100644 --- a/crates/clawhdf5-format/src/chunked_read.rs +++ b/crates/clawhdf5-format/src/chunked_read.rs @@ -519,6 +519,7 @@ pub fn list_chunks( addr_opt, single_filtered_size, single_filter_mask, + unfiltered_edges, ) = match layout { DataLayout::Chunked { chunk_dimensions, @@ -527,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, @@ -534,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( @@ -566,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)? @@ -645,6 +648,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)) } @@ -1829,6 +1849,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 { @@ -1852,6 +1873,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, @@ -2009,6 +2031,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, @@ -2091,6 +2114,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, @@ -2253,6 +2277,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, diff --git a/crates/clawhdf5-format/src/chunked_write.rs b/crates/clawhdf5-format/src/chunked_write.rs index 6816546..04d2b3e 100644 --- a/crates/clawhdf5-format/src/chunked_write.rs +++ b/crates/clawhdf5-format/src/chunked_write.rs @@ -1314,6 +1314,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/tests/fixtures/h5fc_edge_v3.h5 b/crates/clawhdf5/tests/fixtures/h5fc_edge_v3.h5 new file mode 100644 index 0000000000000000000000000000000000000000..6f92057cf653c4dc47cad6133e2ecdc425e861ae GIT binary patch literal 2478 zcmeD5aB<`1lHy|G;9!7(|4^`w6Cz>&mAE-AN6_EHC5TC3cZXdBGOl0}0;z#%gwb3~ z42%p6Y$6Omiqj?7H6-5E#ohHV)HntQ7=wKJL%-5Nu#*&#odaSsFfytzFfcJOGV(x; zXM@t{d}a{_MnMI61{Scq0)sdQgLplV%?flqGdDL610&E7pbvn40E%(|eagbb$jr#f z%EHWbALu_Ym;j{*5q0yk(ib;J4;KbTJ}#Kwe?#4S_Jl%}o1=>p10&1_DbV=mhxh=> zO9AtWQRJFXiz2}1{Kt`Q18UmvsFd71*A%GkLC!COZ5H@wiDsCkw pSUf!9kj%rtD=7IJjmyA5UE6S*2eO9{#taMMMTqqW@g}0H1ONf9YWx5I literal 0 HcmV?d00001 diff --git a/crates/clawhdf5/tests/h5py_chunked_read_tests.rs b/crates/clawhdf5/tests/h5py_chunked_read_tests.rs index deda24f..2b9e902 100644 --- a/crates/clawhdf5/tests/h5py_chunked_read_tests.rs +++ b/crates/clawhdf5/tests/h5py_chunked_read_tests.rs @@ -237,3 +237,109 @@ with h5py.File("{p}", "w") as f: 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); +} From 9066d34eaa3032c4eca6be956c857a7296d1cdf4 Mon Sep 17 00:00:00 2001 From: osobh Date: Fri, 25 Sep 2026 21:12:52 -0500 Subject: [PATCH 5/5] fix(format): key the shared chunk cache by dataset A File is Send + Sync and keeps one ChunkCache for all its datasets. The cached readers bound that cache to "the current dataset" with ensure_dataset(addr), then checked, built and read its index and its decompressed chunks in separate lock acquisitions. Two threads reading two chunked datasets interleaved those steps, so one could store its chunk index under the other's binding, or get the other's decompressed chunk for the same coordinate: wrong data, or an index-out-of-bounds panic when the ranks differed (16 threads x 40 reads over 24 datasets panicked on every run). The cache now keeps per-dataset state keyed by chunk-index address: the chunk index, ChunkIndex and ChunkLayout per dataset (held as Arcs, built outside the lock, first writer wins), and decompressed chunks keyed by (address, coordinate). The chunked readers use the new addr-taking methods (chunks_for, chunk_layout_for, get/put_decompressed_in, prefetch_hint_in) exclusively. Memory stays bounded: decompressed data by the existing byte/slot budget across datasets, indexes by at most 64 datasets and 2^20 index entries in total, dropping the least recently used dataset's index first. Switching datasets no longer throws away the other datasets' cached chunks. The address-less methods remain and act on the dataset last bound with ensure_dataset; they are documented as not for concurrent readers. Regression: threads_reading_different_datasets_get_their_own_chunks (crates/clawhdf5/tests/concurrent_chunk_cache.rs), plus cache unit tests datasets_sharing_coordinates_stay_separate, dataset_indexes_are_bounded and concurrent_readers_of_different_datasets_see_their_own_chunks. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/clawhdf5-format/src/chunk_cache.rs | 793 ++++++++++++------ crates/clawhdf5-format/src/chunked_read.rs | 131 ++- .../clawhdf5/tests/concurrent_chunk_cache.rs | 87 ++ 3 files changed, 660 insertions(+), 351 deletions(-) create mode 100644 crates/clawhdf5/tests/concurrent_chunk_cache.rs 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 bc006df..90bc584 100644 --- a/crates/clawhdf5-format/src/chunked_read.rs +++ b/crates/clawhdf5-format/src/chunked_read.rs @@ -836,24 +836,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)?; @@ -920,7 +916,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), } @@ -957,7 +953,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); @@ -1161,24 +1157,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)?; @@ -1209,12 +1201,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 @@ -1233,7 +1225,7 @@ pub fn read_chunked_data_sweep( } else { raw_chunk.to_vec() }; - cache.put_decompressed(coord, dec) + cache.put_decompressed_in(addr, coord, dec) }; let chunk_offsets: Vec = chunk_info @@ -1317,48 +1309,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; @@ -1377,17 +1355,15 @@ pub fn read_chunked_data_indexed( 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) } @@ -2307,12 +2283,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()); @@ -2334,7 +2310,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 @@ -2343,6 +2319,7 @@ mod tests { ) .unwrap(); assert_eq!(raw1, raw2); + assert_eq!(cache.indexed_dataset_count(), 1); } #[test] 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)] + ); +}