From cc1c872a93b03db0046e5e5fe48d1afa5fc6b58a Mon Sep 17 00:00:00 2001 From: osobh Date: Sat, 26 Sep 2026 10:05:31 -0500 Subject: [PATCH 1/7] format: decode chunks into reusable scratch buffers decompress_chunk_exact_with decodes a chunk into a DecodeScratch the caller keeps between chunks, instead of a new Vec per chunk and per filter stage. Deflate inflates into a kept buffer with a reset (not rebuilt) inflater, shuffle interleaves into the other buffer, and Fletcher32 checks and drops its checksum in place (on the stored bytes when it is the first filter undone). Other filters go through the registry as before. Output and errors are those of decompress_chunk_exact; a unit test checks that for every pipeline shape and filter mask with one reused scratch. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/clawhdf5-format/src/filters.rs | 363 ++++++++++++++++++++++++-- 1 file changed, 342 insertions(+), 21 deletions(-) diff --git a/crates/clawhdf5-format/src/filters.rs b/crates/clawhdf5-format/src/filters.rs index 36f2e96..872b8d8 100644 --- a/crates/clawhdf5-format/src/filters.rs +++ b/crates/clawhdf5-format/src/filters.rs @@ -148,6 +148,180 @@ pub fn decompress_chunk_exact( Ok(data) } +/// Buffers a chunk decoder keeps between chunks, so decoding a dataset's +/// chunks one after another reuses the same memory instead of allocating +/// (and faulting in) fresh buffers for every chunk and every filter stage. +/// +/// Use one per thread with [`decompress_chunk_exact_with`]. Buffers larger +/// than [`DecodeScratch::RETAIN_BYTES`] are released by +/// [`DecodeScratch::trim`], so a scratch kept for a long time (a +/// thread-local, say) does not hold on to a huge chunk's memory. +#[derive(Default)] +pub struct DecodeScratch { + a: Vec, + b: Vec, + #[cfg(feature = "deflate")] + inflater: Option, +} + +impl core::fmt::Debug for DecodeScratch { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("DecodeScratch") + .field("a_capacity", &self.a.capacity()) + .field("b_capacity", &self.b.capacity()) + .finish_non_exhaustive() + } +} + +/// Which buffer holds the data between two filter stages. +#[derive(Clone, Copy)] +enum Stage { + /// `compressed[..len]`: still (a prefix of) the stored bytes. + Stored(usize), + A, + B, +} + +impl DecodeScratch { + /// Largest buffer [`trim`](Self::trim) keeps (4 MiB). + pub const RETAIN_BYTES: usize = 4 << 20; + + /// An empty scratch; buffers are allocated on first use. + pub fn new() -> Self { + Self::default() + } + + /// Release any buffer larger than [`Self::RETAIN_BYTES`]. + pub fn trim(&mut self) { + for buf in [&mut self.a, &mut self.b] { + if buf.capacity() > Self::RETAIN_BYTES { + *buf = Vec::new(); + } + } + } +} + +/// [`decompress_chunk_exact`] into reusable buffers: the decoded chunk is +/// returned as a slice of `scratch` (or of `compressed`, when every filter +/// that was applied only appended a checksum), valid until `scratch` is used +/// again. +/// +/// Deflate, shuffle and Fletcher32 — h5py's and libhdf5's usual pipeline — +/// decode without allocating once `scratch` has grown to the chunk size: the +/// inflater writes into a kept buffer (and its state is reset, not +/// rebuilt), shuffle interleaves into the other buffer, and Fletcher32 checks +/// the checksum and drops it in place. Every other filter goes through the +/// filter registry as [`decompress_chunk_masked`] does, and its output +/// replaces a scratch buffer. The result is byte for byte what +/// [`decompress_chunk_exact`] returns, with the same errors. +pub fn decompress_chunk_exact_with<'s>( + compressed: &'s [u8], + pipeline: &FilterPipeline, + chunk_size: usize, + element_size: u32, + filter_mask: u32, + coords: &[u64], + scratch: &'s mut DecodeScratch, +) -> Result<&'s [u8], FormatError> { + // Same per-stage bounds as `decompress_chunk_masked`. + let mut bounds = [0usize; 32]; + let mut bounds_vec = Vec::new(); + let bounds: &mut [usize] = if pipeline.filters.len() <= bounds.len() { + &mut bounds[..pipeline.filters.len()] + } else { + bounds_vec.resize(pipeline.filters.len(), 0); + &mut bounds_vec + }; + let mut size = chunk_size; + for (i, filter) in pipeline.filters.iter().enumerate() { + bounds[i] = size; + if !filter_skipped(filter_mask, i) { + size = filter_output_bound(filter.filter_id, size); + } + } + + let mut stage = Stage::Stored(compressed.len()); + for (i, filter) in pipeline.filters.iter().enumerate().rev() { + if filter_skipped(filter_mask, i) { + continue; + } + let ctx = FilterContext { + filter, + element_size: element_size as usize, + max_output: bounds[i], + }; + // The stage's input, and the buffer its output goes to (the one + // not holding the input). + let (input, out): (&[u8], &mut Vec) = match stage { + Stage::Stored(len) => (&compressed[..len], &mut scratch.a), + Stage::A => (&scratch.a, &mut scratch.b), + Stage::B => (&scratch.b, &mut scratch.a), + }; + let next = match stage { + Stage::Stored(_) | Stage::B => Stage::A, + Stage::A => Stage::B, + }; + match filter.filter_id { + // Built in and never overridable (`register_filter` refuses + // built-in IDs), so the registry would pick exactly these. + FILTER_FLETCHER32 => { + // Check and drop the checksum where the data is. + let payload = fletcher32_payload(input)?; + stage = match stage { + Stage::Stored(_) => Stage::Stored(payload), + Stage::A => { + scratch.a.truncate(payload); + Stage::A + } + Stage::B => { + scratch.b.truncate(payload); + Stage::B + } + }; + continue; + } + FILTER_SHUFFLE => shuffle_decompress_into(input, ctx.element_size, out), + #[cfg(all( + feature = "deflate", + not(all(target_os = "macos", feature = "system-zlib-decompress")) + ))] + FILTER_DEFLATE => { + let limit = if ctx.max_output != 0 { + ctx.max_output + } else { + MAX_DECOMPRESS_SIZE + }; + let size_hint = if ctx.max_output != 0 { + ctx.max_output + } else { + input.len().saturating_mul(4).min(1 << 20) + }; + let inflater = scratch + .inflater + .get_or_insert_with(|| flate2::Decompress::new(true)); + inflater.reset(true); + inflate_bounded_into(inflater, input, size_hint, limit, out) + .map_err(FormatError::DecompressionError)?; + } + _ => *out = filter_registry::decode(input, &ctx)?, + } + stage = next; + } + + let data: &[u8] = match stage { + Stage::Stored(len) => &compressed[..len], + Stage::A => &scratch.a, + Stage::B => &scratch.b, + }; + if chunk_size != 0 && data.len() != chunk_size { + return Err(FormatError::ChunkedReadError(format!( + "chunk at {coords:?} decoded to {} bytes, expected {chunk_size}", + data.len() + ))); + } + Ok(data) +} + /// Apply a filter pipeline to compress a chunk. /// Filters are applied in FORWARD order for compression. pub fn compress_chunk( @@ -892,33 +1066,55 @@ pub(crate) fn inflate_bounded( size_hint: usize, limit: usize, ) -> Result, String> { - use flate2::{Decompress, FlushDecompress, Status}; + let mut out = Vec::new(); + inflate_bounded_into( + &mut flate2::Decompress::new(true), + data, + size_hint, + limit, + &mut out, + )?; + Ok(out) +} + +/// [`inflate_bounded`] with a fresh or reset `inflater`, into `out`: its +/// contents are replaced and its allocation reused. +#[cfg(feature = "deflate")] +fn inflate_bounded_into( + inflater: &mut flate2::Decompress, + data: &[u8], + size_hint: usize, + limit: usize, + out: &mut Vec, +) -> Result<(), String> { + use flate2::{FlushDecompress, Status}; // One byte of headroom past the limit distinguishes an over-size stream // from one that legitimately ends exactly at the limit. let max_capacity = limit.saturating_add(1); - let mut out = Vec::new(); - out.try_reserve_exact(size_hint.clamp(1, max_capacity)) + // A kept buffer may already be larger than `max_capacity`; the decoder + // can then write past the limit, which the check below still refuses. + out.clear(); + let want = size_hint.clamp(1, max_capacity); + out.try_reserve_exact(want) .map_err(|e| format!("deflate: cannot allocate output: {e}"))?; - let mut inflater = Decompress::new(true); loop { let (in_before, out_before) = (inflater.total_in(), inflater.total_out()); let status = inflater - .decompress_vec( - &data[in_before as usize..], - &mut out, - FlushDecompress::Finish, - ) + .decompress_vec(&data[in_before as usize..], out, FlushDecompress::Finish) .map_err(|e| format!("deflate: {e}"))?; if out.len() > limit { return Err("deflate: output exceeds size limit".into()); } match status { - Status::StreamEnd => return Ok(out), + Status::StreamEnd => return Ok(()), Status::Ok | Status::BufError if out.len() == out.capacity() => { // Out of room: double, up to the limit. - let grow = out.capacity().min(max_capacity - out.capacity()).max(1); + let grow = out + .capacity() + .min(max_capacity.saturating_sub(out.capacity())) + .max(1); out.try_reserve_exact(grow) .map_err(|e| format!("deflate: cannot allocate output: {e}"))?; } @@ -1217,16 +1413,30 @@ fn zstd_compress(data: &[u8], level: u32) -> Result, FormatError> { /// On disk: all byte-0s of each element together, then all byte-1s, etc. /// Output: elements in natural order. fn shuffle_decompress(data: &[u8], element_size: usize) -> Result, FormatError> { + let mut result = Vec::new(); + shuffle_decompress_into(data, element_size, &mut result); + Ok(result) +} + +/// [`shuffle_decompress`] into `result`, replacing its contents and reusing +/// its allocation. +fn shuffle_decompress_into(data: &[u8], element_size: usize, result: &mut Vec) { if element_size <= 1 { - return Ok(data.to_vec()); + result.clear(); + result.extend_from_slice(data); + return; } // 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; whole]; - result.reserve_exact(tail.len()); + // Every byte of `result[..whole]` is overwritten below, so a reused + // buffer keeps its old bytes instead of being zeroed first; only growth + // is zero-filled. + result.truncate(whole); + result.reserve_exact(whole + tail.len() - result.len()); + result.resize(whole, 0); // The shuffled stream is `element_size` byte planes of `num_elements` // bytes each; un-shuffling interleaves them. This is on the read path of @@ -1245,10 +1455,10 @@ fn shuffle_decompress(data: &[u8], element_size: usize) -> Result, Forma } } match element_size { - 2 => interleave::<2>(data, num_elements, &mut result), - 4 => interleave::<4>(data, num_elements, &mut result), - 8 => interleave::<8>(data, num_elements, &mut result), - 16 => interleave::<16>(data, num_elements, &mut result), + 2 => interleave::<2>(data, num_elements, result), + 4 => interleave::<4>(data, num_elements, result), + 8 => interleave::<8>(data, num_elements, result), + 16 => interleave::<16>(data, num_elements, result), _ => { for (i, element) in result.chunks_exact_mut(element_size).enumerate() { for (j, byte) in element.iter_mut().enumerate() { @@ -1258,8 +1468,6 @@ fn shuffle_decompress(data: &[u8], element_size: usize) -> Result, Forma } } result.extend_from_slice(tail); - - Ok(result) } /// Shuffle (compress direction): group bytes by position within each element. @@ -1411,6 +1619,12 @@ fn fletcher32_compute(data: &[u8]) -> u32 { /// Verify Fletcher32 checksum and strip it from the data. /// The last 4 bytes are the stored checksum. fn fletcher32_verify(data: &[u8]) -> Result, FormatError> { + fletcher32_payload(data).map(|len| data[..len].to_vec()) +} + +/// Verify the Fletcher32 checksum that ends `data`; the length of the data +/// before it. +fn fletcher32_payload(data: &[u8]) -> Result { if data.len() < 4 { return Err(FormatError::FilterError( "fletcher32: data too short for checksum".into(), @@ -1430,7 +1644,7 @@ fn fletcher32_verify(data: &[u8]) -> Result, FormatError> { computed, }); } - Ok(payload.to_vec()) + Ok(payload.len()) } /// Append Fletcher32 checksum to data. @@ -1545,6 +1759,113 @@ fn pcodec_decompress( #[cfg(test)] mod tests { + /// `decompress_chunk_exact_with` returns exactly what + /// `decompress_chunk_exact` returns — data or error — for every pipeline + /// shape, filter mask and chunk size, with one scratch reused across all + /// of them in an order that grows, shrinks and swaps its buffers. + #[test] + fn decode_with_scratch_matches_the_allocating_decoder() { + let f = |filter_id: u16, client_data: Vec| FilterDescription { + filter_id, + name: None, + flags: 0, + client_data, + }; + let mut pipelines = vec![ + vec![f(FILTER_SHUFFLE, vec![4])], + vec![f(FILTER_FLETCHER32, vec![])], + // NetCDF-4's order: the checksum is taken before shuffle. + vec![f(FILTER_FLETCHER32, vec![]), f(FILTER_SHUFFLE, vec![4])], + ]; + #[cfg(feature = "deflate")] + pipelines.extend([ + vec![f(FILTER_DEFLATE, vec![4])], + vec![f(FILTER_SHUFFLE, vec![4]), f(FILTER_DEFLATE, vec![4])], + // h5py's order with `fletcher32=True`: checksum last. + vec![ + f(FILTER_SHUFFLE, vec![4]), + f(FILTER_DEFLATE, vec![4]), + f(FILTER_FLETCHER32, vec![]), + ], + vec![ + f(FILTER_FLETCHER32, vec![]), + f(FILTER_SHUFFLE, vec![4]), + f(FILTER_DEFLATE, vec![1]), + ], + ]); + #[cfg(feature = "lzf")] + pipelines.push(vec![ + f(FILTER_SHUFFLE, vec![4]), + f(crate::filter_pipeline::FILTER_LZF, vec![]), + ]); + + let mut scratch = DecodeScratch::new(); + for elements in [1usize, 7, 4096, 3, 65536, 100] { + let data: Vec = (0..elements as u32) + .flat_map(|i| (i.wrapping_mul(2654435761) >> (i % 13)).to_le_bytes()) + .collect(); + for filters in &pipelines { + let pipeline = FilterPipeline { + version: 2, + filters: filters.clone(), + }; + let n = filters.len() as u32; + for mask in 0..(1u32 << n) { + // Encode only the filters the mask says were applied. + let mut stored = data.clone(); + for (i, filter) in filters.iter().enumerate() { + if mask & (1 << i) == 0 { + let ctx = FilterContext { + filter, + element_size: 4, + max_output: 0, + }; + stored = filter_registry::encode(&stored, &ctx).unwrap(); + } + } + let mut cases = vec![(stored.clone(), data.len())]; + // Corrupt: last byte flipped, truncated, wrong size. + let mut flipped = stored.clone(); + *flipped.last_mut().unwrap() ^= 0x5a; + cases.push((flipped, data.len())); + cases.push((stored[..stored.len() / 2].to_vec(), data.len())); + cases.push((stored.clone(), data.len() + 4)); + cases.push((stored.clone(), 0)); + for (bytes, size) in cases { + let want = decompress_chunk_exact(&bytes, &pipeline, size, 4, mask, &[3]); + let got = decompress_chunk_exact_with( + &bytes, + &pipeline, + size, + 4, + mask, + &[3], + &mut scratch, + ) + .map(<[u8]>::to_vec); + match (&want, &got) { + (Ok(w), Ok(g)) => assert_eq!(w, g, "{filters:?} mask {mask}"), + (Err(w), Err(g)) => { + assert_eq!(w.to_string(), g.to_string(), "{filters:?}") + } + _ => panic!("{filters:?} mask {mask} size {size}: {want:?} vs {got:?}"), + } + } + } + } + } + // Long-lived scratch gives back a huge chunk's buffers. + let big = vec![0u8; DecodeScratch::RETAIN_BYTES + 8]; + let shuffle = FilterPipeline { + version: 2, + filters: vec![f(FILTER_SHUFFLE, vec![4])], + }; + decompress_chunk_exact_with(&big, &shuffle, big.len(), 4, 0, &[0], &mut scratch).unwrap(); + scratch.trim(); + assert!(scratch.a.capacity() <= DecodeScratch::RETAIN_BYTES); + assert!(scratch.b.capacity() <= DecodeScratch::RETAIN_BYTES); + } + /// A chunk whose pipeline decodes to fewer bytes than the chunk holds is /// an error naming the chunk, never a short buffer the reader pads. #[test] From f0db8176781192715a05658c191a00b60f924e47 Mon Sep 17 00:00:00 2001 From: osobh Date: Sat, 26 Sep 2026 10:12:38 -0500 Subject: [PATCH 2/7] format: decode full chunked reads straight into the output Full reads of a chunked dataset (the cached reader behind the facade's read_* and the uncached one behind mmap/lazy files and verify_provenance) now decode each chunk into this thread's reusable scratch buffers and copy it straight to its place in the output. Before, the cached reader decoded batches of 128 chunks into fresh Vecs and the uncached one decoded every chunk of the dataset into its own buffer before assembling any: a new 256 KiB allocation (and its page faults) per chunk and per filter stage. Chunks are still inserted into the file's chunk cache when the whole dataset fits in it. With the parallel feature the calling thread now decodes too, sharing the chunks with whichever rayon workers are free (run_with_helpers): a helper the busy pool only starts after the read is done returns at once. Before, the caller handed every chunk to the pool and slept, so readers outside a small pool (2-4 threads) queued behind its workers; with a one-thread pool the reads went sequential. Chunks are placed concurrently only when the index puts them on the chunk grid at distinct places (a corrupt index is read one chunk at a time), and the error returned is still the first failing chunk's. Fix: a chunk stored unfiltered in a filtered dataset (every filter-mask bit set) that is shorter than a chunk read as zeros where its data was missing through the cached reader (the facade's read_*); it is now an error naming the chunk, as the uncached reader already made it. Regression tests, both failing before this change: tests/busy_decode_pool.rs (both workers of a two-thread pool busy, four readers) and short_unfiltered_chunk_of_a_filtered_dataset_is_an_error. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/clawhdf5-format/src/chunked_read.rs | 1054 +++++++++++++------ crates/clawhdf5-format/src/parallel_read.rs | 174 +++ crates/clawhdf5/tests/busy_decode_pool.rs | 112 ++ 3 files changed, 1027 insertions(+), 313 deletions(-) create mode 100644 crates/clawhdf5/tests/busy_decode_pool.rs diff --git a/crates/clawhdf5-format/src/chunked_read.rs b/crates/clawhdf5-format/src/chunked_read.rs index 70176e2..7d3efd0 100644 --- a/crates/clawhdf5-format/src/chunked_read.rs +++ b/crates/clawhdf5-format/src/chunked_read.rs @@ -6,15 +6,16 @@ extern crate alloc; #[cfg(not(feature = "std"))] use alloc::{format, vec, vec::Vec}; -use crate::chunk_cache::CacheAlignedBuffer; #[cfg(feature = "std")] -use crate::chunk_cache::ChunkCache; +use crate::chunk_cache::{CacheAlignedBuffer, ChunkCache}; use crate::data_layout::DataLayout; use crate::dataspace::Dataspace; 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::{DecodeScratch, decompress_chunk_exact_with}; +#[cfg(feature = "std")] use crate::filters::{all_filters_skipped, decompress_chunk_exact}; use crate::fixed_array::{FixedArrayHeader, read_fixed_array_chunks}; #[cfg(feature = "std")] @@ -26,60 +27,329 @@ use crate::parallel_read; #[cfg(feature = "parallel")] use crate::lane_partition::PartitionStats; -/// Decompress all chunks into cache-line-aligned buffers, using lane-partitioned -/// parallel decompression when the `parallel` feature is enabled and the chunk -/// count exceeds the threshold. -fn decompress_all_chunks( - file_data: &[u8], - chunks: &[ChunkInfo], - pipeline: Option<&FilterPipeline>, - chunk_total_bytes: usize, - element_size: u32, -) -> Result, FormatError> { - #[cfg(feature = "parallel")] +/// Run `f` with this thread's chunk-decoding scratch buffers (see +/// [`DecodeScratch`]), kept between reads so decoding reuses memory instead +/// of faulting in fresh pages for every chunk. A re-entrant call (a +/// registered filter codec that itself reads a file) gets a fresh scratch. +fn with_scratch(f: impl FnOnce(&mut DecodeScratch) -> R) -> R { + #[cfg(feature = "std")] { - if let Some(pl) = pipeline - && parallel_read::should_use_parallel(chunks.len()) - && parallel_read::pool_can_parallelise() - { - // Seed from the first chunk's address and count for determinism. - let seed = chunks.first().map(|c| c.address).unwrap_or(0) ^ (chunks.len() as u64); - let (data, _stats) = parallel_read::decompress_chunks_lane_partitioned( - file_data, - chunks, - pl, - chunk_total_bytes, - element_size, - seed, - None, // auto-detect lane count - )?; - return Ok(data.into_iter().map(CacheAlignedBuffer::from_vec).collect()); + use std::cell::RefCell; + std::thread_local! { + static SCRATCH: RefCell = RefCell::new(DecodeScratch::new()); + } + let mut f = Some(f); + let kept = SCRATCH.try_with(|cell| { + let mut scratch = cell.try_borrow_mut().ok()?; + let f = f.take()?; + let r = f(&mut scratch); + scratch.trim(); + Some(r) + }); + if let Ok(Some(r)) = kept { + return r; + } + let f = f.expect("with_scratch: closure already run"); + f(&mut DecodeScratch::new()) + } + #[cfg(not(feature = "std"))] + f(&mut DecodeScratch::new()) +} + +/// A full read's output buffer, written through a raw pointer so that +/// several threads can place chunks into it at once. +struct OutBuf<'a> { + ptr: *mut u8, + len: usize, + _borrow: core::marker::PhantomData<&'a mut [u8]>, +} + +// SAFETY: `OutBuf` is a `&mut [u8]` that hands out writes; sharing it across +// threads is sound as long as concurrent writes do not overlap, which +// `write`'s contract requires. +unsafe impl Send for OutBuf<'_> {} +unsafe impl Sync for OutBuf<'_> {} + +impl<'a> OutBuf<'a> { + fn new(out: &'a mut [u8]) -> Self { + Self { + ptr: out.as_mut_ptr(), + len: out.len(), + _borrow: core::marker::PhantomData, } } - // Sequential fallback — allocate into aligned buffers - let mut result = Vec::with_capacity(chunks.len()); - for chunk_info in chunks { - let c_addr = chunk_info.address as usize; - let size = chunk_info.chunk_size as usize; - ensure_len(file_data, c_addr, size)?; - let raw_chunk = &file_data[c_addr..c_addr + size]; - - let decompressed = if let Some(pl) = pipeline { - decompress_chunk_exact( - raw_chunk, - pl, - chunk_total_bytes, - element_size, - chunk_info.filter_mask, - &chunk_info.offsets, - )? - } else { - raw_chunk.to_vec() - }; - result.push(CacheAlignedBuffer::from_vec(decompressed)); + fn len(&self) -> usize { + self.len } - Ok(result) + + /// Copy `src` to `[at, at + src.len())`. Out of range is a no-op (the + /// callers check first). + /// + /// # Safety + /// + /// No other thread may be writing an overlapping range at the same time. + unsafe fn write(&self, at: usize, src: &[u8]) { + if at.checked_add(src.len()).is_none_or(|end| end > self.len) { + debug_assert!(false, "OutBuf::write out of range"); + return; + } + // SAFETY: in range (checked above) of a live `&'a mut [u8]`; `src` + // cannot overlap it (the output is exclusively borrowed); no + // concurrent overlapping write (the caller's contract). + unsafe { core::ptr::copy_nonoverlapping(src.as_ptr(), self.ptr.add(at), src.len()) } + } +} + +/// How a chunked dataset's chunks map into its row-major output. +struct ChunkPlacer { + rank: usize, + chunk_dims: Vec, + ds_dims: Vec, + ds_strides: Vec, + chunk_strides: Vec, + elem_size: usize, +} + +impl ChunkPlacer { + /// `chunk_dims` and `ds_dims` have one entry per dimension; the dataset + /// must not be empty (so the stride products stay in range). + fn new(chunk_dims: &[usize], ds_dims: &[usize], elem_size: usize) -> Self { + let rank = chunk_dims.len(); + let mut ds_strides = vec![1usize; rank]; + let mut chunk_strides = vec![1usize; rank]; + for i in (0..rank.saturating_sub(1)).rev() { + ds_strides[i] = ds_strides[i + 1] * ds_dims[i + 1]; + chunk_strides[i] = chunk_strides[i + 1] * chunk_dims[i + 1]; + } + Self { + rank, + chunk_dims: chunk_dims.to_vec(), + ds_dims: ds_dims.to_vec(), + ds_strides, + chunk_strides, + elem_size, + } + } + + /// Copy the decoded chunk `data`, whose first element is at `offsets` + /// (one per dimension), to its place in `out`. + /// + /// # Safety + /// + /// No other thread may be placing a chunk whose region overlaps this + /// one's (see [`Self::regions_disjoint`]). + unsafe fn place(&self, data: &[u8], offsets: &[u64], out: &OutBuf<'_>) { + if self.rank == 0 { + let copy_len = data.len().min(out.len()); + // SAFETY: in range; exclusivity is the caller's contract. + unsafe { out.write(0, &data[..copy_len]) }; + return; + } + let mut short = [0usize; 8]; + let mut long = Vec::new(); + let chunk_offsets: &mut [usize] = if self.rank <= short.len() { + &mut short[..self.rank] + } else { + long.resize(self.rank, 0); + &mut long + }; + for (o, &off) in chunk_offsets.iter_mut().zip(offsets) { + *o = off as usize; + } + // SAFETY: the caller's contract. + unsafe { + copy_chunk_into( + data, + out, + chunk_offsets, + &self.chunk_dims, + &self.ds_dims, + &self.ds_strides, + &self.chunk_strides, + self.elem_size, + self.rank, + ) + }; + } + + /// Whether placing `chunks` writes pairwise disjoint regions of the + /// output, so they can be placed concurrently: every chunk starts on the + /// chunk grid and no two start at the same place (a chunk wholly outside + /// the dataset writes nothing and is ignored). A corrupt index that + /// breaks this is read one chunk at a time instead. + #[cfg(feature = "parallel")] + fn regions_disjoint(&self, chunks: &[ChunkInfo]) -> bool { + if self.rank == 0 { + return chunks.len() <= 1; + } + let mut keys = Vec::with_capacity(chunks.len()); + 'chunks: for c in chunks { + if c.offsets.len() < self.rank { + return false; + } + let mut key = 0u64; + for d in 0..self.rank { + let (off, cd, dd) = ( + c.offsets[d], + self.chunk_dims[d] as u64, + self.ds_dims[d] as u64, + ); + if off >= dd { + continue 'chunks; + } + if off % cd != 0 { + return false; + } + let Some(k) = key + .checked_mul(dd.div_ceil(cd)) + .and_then(|k| k.checked_add(off / cd)) + else { + return false; + }; + key = k; + } + keys.push(key); + } + keys.sort_unstable(); + keys.windows(2).all(|w| w[0] != w[1]) + } +} + +/// The per-file chunk cache a full read uses, if any: the cache, this +/// dataset's key in it (its chunk-index address), and whether the dataset +/// fits in it — only then are its chunks looked up and inserted. +#[cfg(feature = "std")] +type CacheUse<'a> = Option<(&'a ChunkCache, u64, bool)>; +#[cfg(not(feature = "std"))] +type CacheUse<'a> = Option<&'a core::convert::Infallible>; + +/// Decode every chunk in `chunks` and place it in `output` (the dataset's +/// whole row-major extent, zeroed). +/// +/// Each chunk goes straight from its decoder to its place in the output: +/// decoded into this thread's reusable scratch (or, when the cache is to +/// keep it, into a buffer the cache takes), then copied. With the +/// `parallel` feature and a filter pipeline, the chunks are shared out +/// between the calling thread and idle rayon workers +/// ([`parallel_read::run_with_helpers`]), so the caller never waits on a +/// busy pool. The error returned is the first failing chunk's, in `chunks` +/// order. +#[allow(clippy::too_many_arguments)] +fn fill_from_chunks( + file_data: &[u8], + chunks: &[ChunkInfo], + pipeline: Option<&FilterPipeline>, + placer: &ChunkPlacer, + chunk_total_bytes: usize, + cache: CacheUse<'_>, + output: &mut [u8], +) -> Result<(), FormatError> { + let rank = placer.rank; + let elem_size = placer.elem_size as u32; + let out = OutBuf::new(output); + #[cfg(not(feature = "std"))] + let _ = cache; + + // Decode chunk `i` and place it. Its callers below run it either on one + // thread, or on several for chunks whose regions are pairwise disjoint, + // each chunk once: no two threads ever write the same bytes. + let work = |i: usize, scratch: &mut DecodeScratch| -> Result<(), FormatError> { + let c = &chunks[i]; + if c.offsets.len() < rank { + return Err(FormatError::ChunkedReadError(format!( + "chunk index entry has {} offsets for a rank-{rank} dataset", + c.offsets.len() + ))); + } + let offsets = &c.offsets[..rank]; + let c_addr = c.address as usize; + let size = c.chunk_size as usize; + ensure_len(file_data, c_addr, size)?; + let raw = &file_data[c_addr..c_addr + size]; + let Some(pl) = pipeline else { + // SAFETY: see above. + unsafe { placer.place(raw, offsets, &out) }; + return Ok(()); + }; + // A chunk stored as-is (every filter skipped) is checked and placed + // straight from the file bytes, never cached. + #[cfg(feature = "std")] + if let Some((cache, key, true)) = cache + && !all_filters_skipped(pl, c.filter_mask) + { + let cached = match cache.get_decompressed_in(key, offsets) { + Some(hit) => hit, + None => { + let data = decompress_chunk_exact( + raw, + pl, + chunk_total_bytes, + elem_size, + c.filter_mask, + &c.offsets, + )?; + cache.put_decompressed_in(key, offsets.to_vec(), data) + } + }; + // SAFETY: see above. + unsafe { placer.place(&cached, offsets, &out) }; + return Ok(()); + } + let data = decompress_chunk_exact_with( + raw, + pl, + chunk_total_bytes, + elem_size, + c.filter_mask, + &c.offsets, + scratch, + )?; + // SAFETY: see above. + unsafe { placer.place(data, offsets, &out) }; + Ok(()) + }; + + #[cfg(feature = "parallel")] + if pipeline.is_some() + && parallel_read::should_use_parallel(chunks.len()) + && placer.regions_disjoint(chunks) + { + use core::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Mutex, PoisonError}; + + let n = chunks.len(); + let next = AtomicUsize::new(0); + let failed: Mutex> = Mutex::new(None); + let body = || { + with_scratch(|scratch| { + loop { + let i = next.fetch_add(1, Ordering::Relaxed); + if i >= n { + break; + } + if let Err(e) = work(i, scratch) { + let mut failed = failed.lock().unwrap_or_else(PoisonError::into_inner); + if failed.as_ref().is_none_or(|(at, _)| i < *at) { + *failed = Some((i, e)); + } + // Stop handing out chunks. Every chunk before `i` was + // claimed already and finishes, so the error kept is + // the first in order. + next.store(n, Ordering::Relaxed); + break; + } + } + }) + }; + parallel_read::run_with_helpers(parallel_read::helper_count(n), &body); + return match failed.into_inner().unwrap_or_else(PoisonError::into_inner) { + Some((_, e)) => Err(e), + None => Ok(()), + }; + } + + with_scratch(|scratch| (0..chunks.len()).try_for_each(|i| work(i, scratch))) } /// Decompress all chunks with lane-partitioned parallelism and return @@ -557,11 +827,6 @@ pub fn generate_implicit_chunks( chunks } -/// Read a chunked dataset, decompressing chunks as needed. -/// Chunks decompressed together before being copied out, bounding the extra -/// memory a parallel full read holds at once. -const DECODE_BATCH: usize = 128; - /// B-tree v2 record types used for chunk indexing. const BT2_CHUNK_UNFILTERED: u8 = 10; const BT2_CHUNK_FILTERED: u8 = 11; @@ -820,6 +1085,111 @@ pub fn list_chunks( Ok((chunks, chunk_dims)) } +/// The chunk cache a full read may use (`None` without `std`). +#[cfg(feature = "std")] +pub(crate) type CacheRef<'a> = Option<&'a ChunkCache>; +/// The chunk cache a full read may use (`None` without `std`). +#[cfg(not(feature = "std"))] +pub(crate) type CacheRef<'a> = Option<&'a core::convert::Infallible>; + +/// The body of every full chunked read: list the chunks (through the +/// cache's index for this dataset when there is a cache), allocate the +/// output with `alloc` (zeroed, `total_bytes` long, as bytes through +/// `bytes`), and decode every chunk straight into it. +#[allow(clippy::too_many_arguments)] +pub(crate) fn read_chunked_full( + file_data: &[u8], + layout: &DataLayout, + dataspace: &Dataspace, + datatype: &Datatype, + pipeline: Option<&FilterPipeline>, + offset_size: u8, + length_size: u8, + cache: CacheRef<'_>, + alloc: impl FnOnce(usize) -> Result, + bytes: impl FnOnce(&mut O) -> &mut [u8], +) -> Result { + check_chunk_element_size(layout, datatype, offset_size)?; + let elem_size = datatype.type_size() as usize; + let list = || { + list_chunks( + file_data, + layout, + dataspace, + elem_size, + offset_size, + length_size, + ) + }; + + #[cfg(feature = "std")] + let (chunks, chunk_dims, cache_key) = match cache { + Some(cache) => { + let (chunk_dimensions, version, addr_opt) = match layout { + DataLayout::Chunked { + chunk_dimensions, + version, + btree_address, + .. + } => (chunk_dimensions, *version, *btree_address), + _ => { + return Err(FormatError::ChunkedReadError( + "expected chunked layout".into(), + )); + } + }; + let addr = addr_opt.ok_or_else(|| { + FormatError::ChunkedReadError("no address for chunked layout".into()) + })?; + let (rank, chunk_dims) = + chunk_geometry(chunk_dimensions, version, dataspace, elem_size)?; + // 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().map(|(chunks, _)| chunks))?; + (chunks, chunk_dims, Some((cache, addr))) + } + None => { + let (chunks, chunk_dims) = list()?; + (chunks, chunk_dims, None) + } + }; + #[cfg(not(feature = "std"))] + let (chunks, chunk_dims) = { + let _ = cache; + list()? + }; + + let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?; + let mut output = alloc(total_bytes)?; + if total_bytes == 0 { + // Also keeps the stride products in range: with a zero-sized + // dimension the total is 0 even if other dimensions are huge. + return Ok(output); + } + let ds_dims: Vec = dataspace.dimensions.iter().map(|&d| d as usize).collect(); + let placer = ChunkPlacer::new(&chunk_dims, &ds_dims, elem_size); + let chunk_total_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?; + // Chunks are cached only when the whole dataset fits: pushing a larger + // dataset through the cache just evicts each chunk moments after + // inserting it. + #[cfg(feature = "std")] + let cache_use = cache_key.map(|(cache, key)| (cache, key, total_bytes <= cache.max_bytes())); + #[cfg(not(feature = "std"))] + let cache_use = None; + fill_from_chunks( + file_data, + &chunks, + pipeline, + &placer, + chunk_total_bytes, + cache_use, + bytes(&mut output), + )?; + Ok(output) +} + +/// Read a chunked dataset, decompressing chunks as needed. pub fn read_chunked_data( file_data: &[u8], layout: &DataLayout, @@ -829,118 +1199,26 @@ pub fn read_chunked_data( offset_size: u8, length_size: u8, ) -> Result, FormatError> { - check_chunk_element_size(layout, datatype, offset_size)?; - let elem_size = datatype.type_size() as usize; - let (chunks, chunk_dims) = list_chunks( + read_chunked_full( file_data, layout, dataspace, - elem_size, + datatype, + pipeline, offset_size, length_size, - )?; - let rank = chunk_dims.len(); - let ds_dims: Vec = dataspace.dimensions.iter().map(|&d| d as usize).collect(); - - // Assemble output - let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?; - if total_bytes == 0 { - // Also keeps the stride products below in range: with a zero-sized - // dimension the total is 0 even if other dimensions are huge. - return Ok(Vec::new()); - } - let mut output = alloc_output(total_bytes)?; - - let mut ds_strides = vec![1usize; rank]; - for i in (0..rank.saturating_sub(1)).rev() { - ds_strides[i] = ds_strides[i + 1] * ds_dims[i + 1]; - } - - let mut chunk_strides = vec![1usize; rank]; - for i in (0..rank.saturating_sub(1)).rev() { - chunk_strides[i] = chunk_strides[i + 1] * chunk_dims[i + 1]; - } - - let chunk_total_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?; - - // Fast path: no filters — copy directly from file_data without intermediate alloc - if pipeline.is_none() { - for chunk_info in &chunks { - let chunk_offsets: Vec = chunk_info - .offsets - .iter() - .take(rank) - .map(|&o| o as usize) - .collect(); - - let c_addr = chunk_info.address as usize; - let size = chunk_info.chunk_size as usize; - ensure_len(file_data, c_addr, size)?; - let chunk_data = &file_data[c_addr..c_addr + size]; - - if rank == 0 { - let copy_len = chunk_data.len().min(output.len()); - output[..copy_len].copy_from_slice(&chunk_data[..copy_len]); - } else { - copy_chunk_to_output( - chunk_data, - &mut output, - &chunk_offsets, - &chunk_dims, - &ds_dims, - &ds_strides, - &chunk_strides, - elem_size, - rank, - ); - } - } - return Ok(output); - } - - // Filtered path: decompress all chunks then assemble - let decompressed_chunks = decompress_all_chunks( - file_data, - &chunks, - pipeline, - chunk_total_bytes, - elem_size as u32, - )?; - - for (chunk_info, decompressed) in chunks.iter().zip(decompressed_chunks.iter()) { - let chunk_offsets: Vec = chunk_info - .offsets - .iter() - .take(rank) - .map(|&o| o as usize) - .collect(); - - if rank == 0 { - let copy_len = decompressed.len().min(output.len()); - output[..copy_len].copy_from_slice(&decompressed[..copy_len]); - } else { - copy_chunk_to_output( - decompressed, - &mut output, - &chunk_offsets, - &chunk_dims, - &ds_dims, - &ds_strides, - &chunk_strides, - elem_size, - rank, - ); - } - } - - Ok(output) + None, + alloc_output, + |out| out.as_mut_slice(), + ) } /// Read a chunked dataset with caching support. /// /// On the first call, scans the chunk index (B-tree / fixed array / etc.) once /// and populates the cache's hash index. Subsequent calls skip the index scan -/// entirely. Decompressed chunk data is also cached with LRU eviction. +/// entirely. Decompressed chunk data is also cached with LRU eviction, when +/// the whole dataset fits in the cache. #[cfg(feature = "std")] #[allow(clippy::too_many_arguments)] pub fn read_chunked_data_cached( @@ -953,159 +1231,18 @@ pub fn read_chunked_data_cached( length_size: u8, cache: &ChunkCache, ) -> Result, FormatError> { - let (chunk_dimensions, version, addr_opt) = match layout { - DataLayout::Chunked { - chunk_dimensions, - version, - btree_address, - .. - } => (chunk_dimensions, *version, *btree_address), - _ => { - return Err(FormatError::ChunkedReadError( - "expected chunked layout".into(), - )); - } - }; - - let addr = addr_opt - .ok_or_else(|| FormatError::ChunkedReadError("no address for chunked layout".into()))?; - - check_chunk_element_size(layout, datatype, offset_size)?; - let elem_size = datatype.type_size() as usize; - let (rank, chunk_dims) = chunk_geometry(chunk_dimensions, version, dataspace, elem_size)?; - let ds_dims: Vec = dataspace.dimensions.iter().map(|&d| d as usize).collect(); - - // 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, - ) - .map(|(chunks, _)| chunks) - })?; - - // Assemble output - let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?; - if total_bytes == 0 { - // Also keeps the stride products below in range: with a zero-sized - // dimension the total is 0 even if other dimensions are huge. - return Ok(Vec::new()); - } - let mut output = alloc_output(total_bytes)?; - - let mut ds_strides = vec![1usize; rank]; - for i in (0..rank.saturating_sub(1)).rev() { - ds_strides[i] = ds_strides[i + 1] * ds_dims[i + 1]; - } - - let mut chunk_strides = vec![1usize; rank]; - for i in (0..rank.saturating_sub(1)).rev() { - chunk_strides[i] = chunk_strides[i + 1] * chunk_dims[i + 1]; - } - - let chunk_total_bytes = checked_chunk_byte_len(&chunk_dims, elem_size)?; - - let mut place = |data: &[u8], chunk_info: &ChunkInfo| { - if rank == 0 { - let copy_len = data.len().min(output.len()); - output[..copy_len].copy_from_slice(&data[..copy_len]); - return; - } - let chunk_offsets: Vec = chunk_info - .offsets - .iter() - .take(rank) - .map(|&o| o as usize) - .collect(); - copy_chunk_to_output( - data, - &mut output, - &chunk_offsets, - &chunk_dims, - &ds_dims, - &ds_strides, - &chunk_strides, - elem_size, - rank, - ); - }; - let raw_bytes = |chunk_info: &ChunkInfo| -> Result<&[u8], FormatError> { - let c_addr = chunk_info.address as usize; - let size = chunk_info.chunk_size as usize; - ensure_len(file_data, c_addr, size)?; - Ok(&file_data[c_addr..c_addr + size]) - }; - - // Chunks stored as-is (no pipeline, or the filter mask says this chunk - // 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) { - place(raw_bytes(chunk_info)?, chunk_info); - continue; - } - let coord: Vec = chunk_info.offsets.iter().take(rank).copied().collect(); - match cache.get_decompressed_in(addr, &coord) { - Some(cached) => place(&cached, chunk_info), - None => misses.push(chunk_info), - } - } - - // Decompress what the cache didn't have, a bounded batch at a time — in - // parallel with the `parallel` feature (this path, the one the facade - // uses, was sequential; only the uncached reader was parallel), unless the - // pool has one thread: then every reading thread would queue behind that - // one worker, so each decodes its own chunks instead. Chunks are - // cached only when the whole dataset fits: pushing a larger dataset - // through the cache just evicts each chunk moments after inserting it. - let cache_them = total_bytes <= cache.max_bytes(); - if let Some(pl) = pipeline { - let decode = |c: &&ChunkInfo| -> Result, FormatError> { - decompress_chunk_exact( - raw_bytes(c)?, - pl, - chunk_total_bytes, - elem_size as u32, - c.filter_mask, - &c.offsets, - ) - }; - for batch in misses.chunks(DECODE_BATCH) { - #[cfg(feature = "parallel")] - let decoded: Vec, FormatError>> = - if batch.len() >= 4 && parallel_read::pool_can_parallelise() { - use rayon::prelude::*; - batch.par_iter().map(decode).collect() - } else { - batch.iter().map(decode).collect() - }; - #[cfg(not(feature = "parallel"))] - let decoded: Vec, FormatError>> = batch.iter().map(decode).collect(); - - for (chunk_info, data) in batch.iter().zip(decoded) { - let data = data?; - if cache_them { - let coord: Vec = chunk_info.offsets.iter().take(rank).copied().collect(); - let cached = cache.put_decompressed_in(addr, coord, data); - place(&cached, chunk_info); - } else { - place(&data, chunk_info); - } - } - } - } - - Ok(output) + read_chunked_full( + file_data, + layout, + dataspace, + datatype, + pipeline, + offset_size, + length_size, + Some(cache), + alloc_output, + |out| out.as_mut_slice(), + ) } /// Sweep context passed into `read_chunked_data_sweep` to enable adaptive @@ -1486,6 +1623,7 @@ pub fn read_chunked_data_indexed( } /// Copy chunk data into the output buffer at the correct N-D position. +#[cfg(any(feature = "std", test))] #[allow(clippy::too_many_arguments)] fn copy_chunk_to_output( chunk_data: &[u8], @@ -1497,6 +1635,41 @@ fn copy_chunk_to_output( chunk_strides: &[usize], elem_size: usize, rank: usize, +) { + // SAFETY: `output` is exclusively borrowed, so no other write can + // overlap this one. + unsafe { + copy_chunk_into( + chunk_data, + &OutBuf::new(output), + chunk_offsets, + chunk_dims, + ds_dims, + ds_strides, + chunk_strides, + elem_size, + rank, + ) + } +} + +/// [`copy_chunk_to_output`] through an [`OutBuf`]: only the bytes of the +/// chunk's region (clipped to the dataset) are written. +/// +/// # Safety +/// +/// No other thread may be writing an overlapping region of `output`. +#[allow(clippy::too_many_arguments)] +unsafe fn copy_chunk_into( + chunk_data: &[u8], + output: &OutBuf<'_>, + chunk_offsets: &[usize], + chunk_dims: &[usize], + ds_dims: &[usize], + ds_strides: &[usize], + chunk_strides: &[usize], + elem_size: usize, + rank: usize, ) { // Row-copy approach: iterate over outer dimensions, memcpy the innermost // dimension in bulk. For 1-D data this is a single memcpy per chunk. @@ -1518,7 +1691,8 @@ fn copy_chunk_to_output( .is_some_and(|end| end <= output.len()) && src_bytes <= chunk_data.len() { - output[dst_start..dst_start + src_bytes].copy_from_slice(&chunk_data[..src_bytes]); + // SAFETY: in range (checked); exclusivity is the caller's contract. + unsafe { output.write(dst_start, &chunk_data[..src_bytes]) }; } return; } @@ -1620,8 +1794,8 @@ fn copy_chunk_to_output( .checked_add(row_bytes) .is_some_and(|end| end <= output.len()); if fits { - output[dst_start..dst_start + row_bytes] - .copy_from_slice(&chunk_data[src_start..src_start + row_bytes]); + // SAFETY: in range (checked); exclusivity is the caller's contract. + unsafe { output.write(dst_start, &chunk_data[src_start..src_start + row_bytes]) }; } } } @@ -2071,6 +2245,198 @@ mod tests { ); } + /// A 2-D `f32` dataset of `rows x cols` in `cr x cc` chunks, deflated + /// behind shuffle: (file bytes, chunk list, pipeline, expected output). + #[cfg(feature = "deflate")] + fn deflated_grid( + rows: usize, + cols: usize, + cr: usize, + cc: usize, + ) -> (Vec, Vec, FilterPipeline, Vec) { + use crate::filter_pipeline::{FILTER_DEFLATE, FILTER_SHUFFLE, FilterDescription}; + let pipeline = FilterPipeline { + version: 2, + filters: vec![ + FilterDescription { + filter_id: FILTER_SHUFFLE, + name: None, + flags: 0, + client_data: vec![4], + }, + FilterDescription { + filter_id: FILTER_DEFLATE, + name: None, + flags: 0, + client_data: vec![1], + }, + ], + }; + let value = |r: usize, c: usize| (r * cols + c) as f32 * 0.5; + let expected: Vec = (0..rows * cols) + .flat_map(|i| value(i / cols, i % cols).to_le_bytes()) + .collect(); + let (mut file, mut chunks) = (vec![0u8; 8], Vec::new()); + for r0 in (0..rows).step_by(cr) { + for c0 in (0..cols).step_by(cc) { + // Full-size chunks; the edge padding is zeros. + let chunk: Vec = (0..cr * cc) + .flat_map(|i| { + let (r, c) = (r0 + i / cc, c0 + i % cc); + let v = if r < rows && c < cols { + value(r, c) + } else { + 0.0 + }; + v.to_le_bytes() + }) + .collect(); + let stored = crate::filters::compress_chunk(&chunk, &pipeline, 4).unwrap(); + chunks.push(ChunkInfo { + chunk_size: stored.len() as u32, + filter_mask: 0, + offsets: vec![r0 as u64, c0 as u64, 0], + address: file.len() as u64, + }); + file.extend(stored); + } + } + (file, chunks, pipeline, expected) + } + + #[cfg(feature = "deflate")] + fn fill( + file: &[u8], + chunks: &[ChunkInfo], + pipeline: &FilterPipeline, + dims: [usize; 2], + chunk: [usize; 2], + ) -> Result, FormatError> { + let placer = ChunkPlacer::new(&chunk, &dims, 4); + let mut out = vec![0u8; dims[0] * dims[1] * 4]; + fill_from_chunks( + file, + chunks, + Some(pipeline), + &placer, + chunk[0] * chunk[1] * 4, + None, + &mut out, + )?; + Ok(out) + } + + /// Chunks decoded straight into the output (in parallel with the + /// `parallel` feature, partial edge chunks included) land where the + /// row-by-row reference puts them. + #[test] + #[cfg(feature = "deflate")] + fn chunks_fill_the_output_in_place() { + for (dims, chunk) in [ + ([100, 70], [16, 32]), + ([64, 64], [16, 16]), + ([5, 3], [8, 8]), + ] { + let (file, chunks, pipeline, expected) = + deflated_grid(dims[0], dims[1], chunk[0], chunk[1]); + for _ in 0..4 { + assert_eq!( + fill(&file, &chunks, &pipeline, dims, chunk).unwrap(), + expected + ); + } + // Any chunk order gives the same output. + let mut reversed = chunks.clone(); + reversed.reverse(); + assert_eq!( + fill(&file, &reversed, &pipeline, dims, chunk).unwrap(), + expected + ); + } + } + + /// Of several corrupt chunks, the error names the first in chunk order, + /// however the chunks were shared out between threads. + #[test] + #[cfg(feature = "deflate")] + fn first_corrupt_chunk_is_the_error() { + let (dims, chunk) = ([128, 64], [8, 16]); + let (mut file, mut chunks, pipeline, _) = + deflated_grid(dims[0], dims[1], chunk[0], chunk[1]); + // Valid streams that decode short: an error naming the chunk. + let short = crate::filters::compress_chunk(&[1u8; 64], &pipeline, 4).unwrap(); + for bad in [5usize, 11, 40] { + chunks[bad].address = file.len() as u64; + chunks[bad].chunk_size = short.len() as u32; + file.extend_from_slice(&short); + } + for _ in 0..20 { + let err = fill(&file, &chunks, &pipeline, dims, chunk).unwrap_err(); + let want = format!("{:?}", chunks[5].offsets); + assert!(err.to_string().contains(&want), "{err} (want {want})"); + } + } + + /// Only chunks on the grid, each at a distinct place, may be placed + /// concurrently; anything else is read one chunk at a time. + #[test] + #[cfg(feature = "parallel")] + fn regions_disjoint_only_for_distinct_grid_chunks() { + let placer = ChunkPlacer::new(&[10, 10], &[25, 30], 4); + let chunk = |r: u64, c: u64| ChunkInfo { + chunk_size: 0, + filter_mask: 0, + offsets: vec![r, c, 0], + address: 0, + }; + let grid: Vec = (0..3) + .flat_map(|r| (0..3).map(move |c| chunk(r * 10, c * 10))) + .collect(); + assert!(placer.regions_disjoint(&grid)); + // Chunks wholly outside the dataset write nothing. + let mut outside = grid.clone(); + outside.push(chunk(30, 0)); + outside.push(chunk(30, 0)); + assert!(placer.regions_disjoint(&outside)); + let mut duplicate = grid.clone(); + duplicate.push(chunk(10, 20)); + assert!(!placer.regions_disjoint(&duplicate)); + let mut off_grid = grid.clone(); + off_grid[4] = chunk(15, 10); + assert!(!placer.regions_disjoint(&off_grid)); + let mut short = grid; + short[0].offsets.truncate(1); + assert!(!placer.regions_disjoint(&short)); + } + + /// A duplicated chunk in a corrupt index is not placed from two threads + /// at once: the read goes one chunk at a time, as before. + #[test] + #[cfg(feature = "deflate")] + fn duplicate_chunks_are_read_in_order() { + let (dims, chunk) = ([64, 64], [16, 16]); + let (file, mut chunks, pipeline, expected) = + deflated_grid(dims[0], dims[1], chunk[0], chunk[1]); + // A second entry at chunk 3's place, holding chunk 7's data: the + // later entry wins, as a sequential read has it. + let mut dup = chunks[7].clone(); + dup.offsets = chunks[3].offsets.clone(); + chunks.push(dup); + let got = fill(&file, &chunks, &pipeline, dims, chunk).unwrap(); + let chunk7_in_3: Vec = { + let mut e = expected.clone(); + let (r3, c3) = (chunks[3].offsets[0] as usize, chunks[3].offsets[1] as usize); + let (r7, c7) = (chunks[7].offsets[0] as usize, chunks[7].offsets[1] as usize); + for r in 0..16 { + let src = ((r7 + r) * 64 + c7) * 4; + let dst = ((r3 + r) * 64 + c3) * 4; + e[dst..dst + 64].copy_from_slice(&expected[src..src + 64]); + } + e + }; + assert_eq!(got, chunk7_in_3); + } + #[test] fn copy_chunk_to_output_1d_rejects_overflowing_offset_without_panicking() { // Found by fuzzing: `global_start * elem_size` overflowed for a @@ -2481,6 +2847,68 @@ mod tests { use crate::chunk_cache::ChunkCache; + /// A chunk stored unfiltered in a filtered dataset (every filter-mask + /// bit set) must still hold the whole chunk. The cached reader placed a + /// short one and read its missing rows as zeros; both readers refuse it. + #[test] + fn short_unfiltered_chunk_of_a_filtered_dataset_is_an_error() { + use crate::filter_pipeline::{FILTER_SHUFFLE, FilterDescription}; + let mut file_data = vec![0u8; 0x2000]; + let mut infos = Vec::new(); + for k in 0..3u64 { + let address = 0x1000 + k as usize * 80; + for i in 0..10u64 { + let at = address + i as usize * 8; + file_data[at..at + 8].copy_from_slice(&((k * 10 + i) as f64).to_le_bytes()); + } + infos.push(ChunkInfo { + chunk_size: if k == 1 { 40 } else { 80 }, + filter_mask: 1, + offsets: vec![k * 10, 0], + address: address as u64, + }); + } + let btree = build_chunk_btree_leaf(&infos, 2, 8); + file_data[0x100..0x100 + btree.len()].copy_from_slice(&btree); + let layout = DataLayout::Chunked { + chunk_dimensions: vec![10, 8], + btree_address: Some(0x100), + version: 3, + chunk_index_type: None, + single_chunk_filtered_size: None, + single_chunk_filter_mask: None, + dont_filter_partial_edge_chunks: false, + }; + let dataspace = simple_space(vec![30]); + let pipeline = FilterPipeline { + version: 2, + filters: vec![FilterDescription { + filter_id: FILTER_SHUFFLE, + name: None, + flags: 0, + client_data: vec![8], + }], + }; + let dt = make_f64_type(); + let cache = ChunkCache::new(); + for result in [ + read_chunked_data(&file_data, &layout, &dataspace, &dt, Some(&pipeline), 8, 8), + read_chunked_data_cached( + &file_data, + &layout, + &dataspace, + &dt, + Some(&pipeline), + 8, + 8, + &cache, + ), + ] { + let err = result.unwrap_err().to_string(); + assert!(err.contains("[10, 0]") && err.contains("40"), "{err}"); + } + } + #[test] fn cached_read_populates_index_and_returns_correct_data() { let values: Vec = (0..20).map(|i| i as f64).collect(); diff --git a/crates/clawhdf5-format/src/parallel_read.rs b/crates/clawhdf5-format/src/parallel_read.rs index 1b940c6..61d06cd 100644 --- a/crates/clawhdf5-format/src/parallel_read.rs +++ b/crates/clawhdf5-format/src/parallel_read.rs @@ -41,6 +41,126 @@ pub fn pool_can_parallelise() -> bool { rayon::current_num_threads() > 1 } +/// How many rayon workers [`run_with_helpers`] should ask to help with +/// `items` work items, given that the calling thread works too: the pool's +/// other threads (all of them when the caller is not one), at most one per +/// item beyond the caller's first. +pub(crate) fn helper_count(items: usize) -> usize { + let pool = rayon::current_num_threads(); + let others = if rayon::current_thread_index().is_some() { + pool.saturating_sub(1) + } else { + pool + }; + others.min(items.saturating_sub(1)) +} + +/// Run `body` on the calling thread and on up to `helpers` rayon workers at +/// once, returning when the caller's call has finished and every worker that +/// started one has too. `body` shares its work out itself (typically by +/// claiming items from an atomic counter until none are left). +/// +/// The caller never waits for a worker to *become* free: helpers are queued +/// on the pool, and one that only gets to run after the caller has finished +/// returns without calling `body`. So a busy or small pool can only fail to +/// speed a read up, never hold it back — with `par_iter`, the calling thread +/// (not a pool worker) handed all the work to the pool and slept, and N +/// threads reading through a 2-worker pool decoded on 2 cores. +/// +/// A panic in `body`, on any thread, is resumed on the caller once every +/// helper that started has stopped. +pub(crate) fn run_with_helpers(helpers: usize, body: &(dyn Fn() + Sync)) { + use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind}; + use std::sync::{Arc, Condvar, Mutex, PoisonError}; + + if helpers == 0 { + body(); + return; + } + + type Body = dyn Fn() + Sync + 'static; + struct Shared { + /// `body`, its lifetime erased. Only dereferenced by a helper that + /// registered in `state` while it was open (see below). + body: *const Body, + /// (closed, helpers inside `body`). + state: Mutex<(bool, usize)>, + idle: Condvar, + panic: Mutex>>, + } + // SAFETY: `body` points to a `Sync` closure, so calling it from other + // threads is allowed; the pointer is only used under the protocol below, + // which keeps it from outliving the closure. + unsafe impl Send for Shared {} + unsafe impl Sync for Shared {} + + fn help(shared: &Shared) { + { + let mut state = shared.state.lock().unwrap_or_else(PoisonError::into_inner); + if state.0 { + return; + } + state.1 += 1; + } + // SAFETY: registered while open, so the caller of `run_with_helpers` + // is still inside it (it closes, then waits until no helper is + // registered, before returning), and `body` is alive. + let body = unsafe { &*shared.body }; + if let Err(payload) = catch_unwind(AssertUnwindSafe(body)) { + shared + .panic + .lock() + .unwrap_or_else(PoisonError::into_inner) + .get_or_insert(payload); + } + let mut state = shared.state.lock().unwrap_or_else(PoisonError::into_inner); + state.1 -= 1; + if state.1 == 0 { + shared.idle.notify_all(); + } + } + + let body_ptr: *const (dyn Fn() + Sync + '_) = body; + // SAFETY: only the lifetime changes (same fat-pointer layout). The + // pointer is dereferenced only while this function is running: see + // `help` and the wait below. + let body_ptr: *const Body = unsafe { core::mem::transmute(body_ptr) }; + let shared = Arc::new(Shared { + body: body_ptr, + state: Mutex::new((false, 0)), + idle: Condvar::new(), + panic: Mutex::new(None), + }); + for _ in 0..helpers { + let shared = Arc::clone(&shared); + rayon::spawn(move || help(&shared)); + } + let caller = catch_unwind(AssertUnwindSafe(body)); + { + // Close, then wait for the helpers inside `body`; later ones return + // at once. This must happen even if `body` panicked on this thread. + let mut state = shared.state.lock().unwrap_or_else(PoisonError::into_inner); + state.0 = true; + while state.1 > 0 { + state = shared + .idle + .wait(state) + .unwrap_or_else(PoisonError::into_inner); + } + } + if let Err(payload) = caller { + resume_unwind(payload); + } + let helper_panic = shared + .panic + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take(); + if let Some(payload) = helper_panic { + resume_unwind(payload); + } +} + /// Decompress chunks in parallel using lane-partitioned assignment. /// /// Instead of naive `par_iter`, chunks are deterministically assigned to lanes @@ -258,6 +378,60 @@ mod tests { (file, infos) } + /// Every item is processed exactly once, whatever mix of caller and + /// helpers ends up doing it. + #[test] + fn run_with_helpers_shares_all_work() { + use core::sync::atomic::{AtomicUsize, Ordering}; + for helpers in [0, 1, 3, 16] { + let n = 1000; + let next = AtomicUsize::new(0); + let done: Vec = (0..n).map(|_| AtomicUsize::new(0)).collect(); + run_with_helpers(helpers, &|| { + loop { + let i = next.fetch_add(1, Ordering::Relaxed); + if i >= n { + break; + } + done[i].fetch_add(1, Ordering::Relaxed); + } + }); + assert!(done.iter().all(|d| d.load(Ordering::Relaxed) == 1)); + } + } + + /// A panic in the shared body reaches the caller whichever thread it + /// happened on, and only after the helpers inside the body have left it + /// (they borrow the caller's stack). + #[test] + fn run_with_helpers_propagates_panics() { + use core::sync::atomic::{AtomicUsize, Ordering}; + use std::panic::{AssertUnwindSafe, catch_unwind}; + let caller = std::thread::current().id(); + for panic_on_caller in [true, false] { + let inside = AtomicUsize::new(0); + let calls = AtomicUsize::new(0); + let result = catch_unwind(AssertUnwindSafe(|| { + run_with_helpers(4, &|| { + inside.fetch_add(1, Ordering::SeqCst); + calls.fetch_add(1, Ordering::SeqCst); + let on_caller = std::thread::current().id() == caller; + std::thread::sleep(std::time::Duration::from_millis(20)); + inside.fetch_sub(1, Ordering::SeqCst); + if on_caller == panic_on_caller { + panic!("boom"); + } + }); + })); + // A helper may never have run (the pool was slow to start it), + // in which case nothing panicked when `panic_on_caller` is false. + if panic_on_caller || calls.load(Ordering::SeqCst) > 1 { + assert!(result.is_err()); + } + assert_eq!(inside.load(Ordering::SeqCst), 0); + } + } + /// Every parallel decoder refuses a chunk that decodes short, naming it. #[test] fn short_decoded_chunk_is_an_error() { diff --git a/crates/clawhdf5/tests/busy_decode_pool.rs b/crates/clawhdf5/tests/busy_decode_pool.rs new file mode 100644 index 0000000..d53eeab --- /dev/null +++ b/crates/clawhdf5/tests/busy_decode_pool.rs @@ -0,0 +1,112 @@ +//! Full reads of chunked datasets must not wait for a busy rayon pool of +//! any size. +//! +//! A full read handed its chunks to the rayon pool (`par_iter`) and the +//! calling thread — not a pool worker — slept until the pool had decoded +//! them. With a small pool (2-4 threads) and more reading threads than +//! workers, every reader queued behind the same few workers +//! (`docs/known-issues.md`, "Concurrent and contiguous read performance"). +//! Now the calling thread decodes too, and pool workers only help when they +//! are free. The test keeps both workers of a two-thread pool busy and +//! requires reads to finish anyway, with the right values. +//! +//! One test in its own binary: it configures the process-wide rayon pool. + +#![cfg(feature = "parallel")] + +use std::sync::mpsc; +use std::time::Duration; + +use clawhdf5::{File, FileBuilder}; + +const N: usize = 4096; // 64 chunks of 64 elements + +fn values() -> Vec { + (0..N).map(|i| i as f64 * 0.25 - 7.0).collect() +} + +fn build() -> File { + let mut b = FileBuilder::new(); + b.create_dataset("data") + .with_f64_data(&values()) + .with_shape(&[N as u64]) + .with_chunks(&[64]) + .with_deflate(1) + .with_provenance("test-suite", "2026-09-26T00:00:00Z", None); + File::from_bytes(b.finish().unwrap()).unwrap() +} + +/// Run `f` on a fresh thread; `None` if it has not finished within `limit`. +fn finishes_within( + limit: Duration, + f: impl FnOnce() -> T + Send + 'static, +) -> Option { + let (tx, rx) = mpsc::channel(); + std::thread::spawn(move || { + let _ = tx.send(f()); + }); + rx.recv_timeout(limit).ok() +} + +#[test] +fn full_reads_do_not_wait_for_a_busy_small_pool() { + const WORKERS: usize = 2; + rayon::ThreadPoolBuilder::new() + .num_threads(WORKERS) + .build_global() + .expect("this test binary configures the global pool first"); + + // Built first: the writer compresses on the pool too. + let file = std::sync::Arc::new(build()); + + // Occupy every worker of the pool until the reads are done. + let (started_tx, started_rx) = mpsc::channel(); + let (release_tx, release_rx) = mpsc::channel::<()>(); + let release_rx = std::sync::Arc::new(std::sync::Mutex::new(release_rx)); + for _ in 0..WORKERS { + let (started_tx, release_rx) = (started_tx.clone(), release_rx.clone()); + rayon::spawn(move || { + started_tx.send(()).unwrap(); + let _ = release_rx.lock().unwrap().recv(); + }); + } + for _ in 0..WORKERS { + started_rx.recv().unwrap(); + } + + let limit = Duration::from_secs(20); + // Several readers at once, as in the `concurrent_read` benchmark: the + // cached full read (`read_*`), the typed one and the uncached reader + // behind `verify_provenance`. + let readers: Vec<_> = (0..4) + .map(|_| { + let file = std::sync::Arc::clone(&file); + std::thread::spawn(move || { + finishes_within(limit, move || { + let ds = file.dataset("data").unwrap(); + ( + ds.read_f64().unwrap(), + ds.read_f32().unwrap(), + ds.verify_provenance().unwrap(), + ) + }) + }) + }) + .collect(); + let results: Vec<_> = readers.into_iter().map(|h| h.join().unwrap()).collect(); + // Free the workers before asserting, so a failure does not hang the + // blocked reader threads forever. + for _ in 0..WORKERS { + release_tx.send(()).unwrap(); + } + + let want = values(); + let want_f32: Vec = want.iter().map(|&v| v as f32).collect(); + for result in results { + let (f64s, f32s, verified) = + result.expect("a full read waited for the busy two-thread rayon pool"); + assert_eq!(f64s, want); + assert_eq!(f32s, want_f32); + assert_eq!(verified, clawhdf5::provenance::VerifyResult::Ok); + } +} From 9e9b849dd7cf6bd9de4e7ca72a4abdf42c3f4c59 Mon Sep 17 00:00:00 2001 From: osobh Date: Sat, 26 Sep 2026 10:18:29 -0500 Subject: [PATCH 3/7] read chunked datasets straight into the typed output read_f32/read_f64/read_i32/read_i64/read_u64 of a chunked dataset that stores exactly that type in native byte order now decode every chunk straight into the Vec they return (data_read::read_chunked_native), on File (through its chunk cache), MmapFile and LazyFile. Before, the chunks went into a byte buffer that read_as_* then copied into a second, typed one: two dataset-sized allocations and a full extra copy per read. The output is zeroed pages from the allocator, backed by transparent huge pages when large, like the byte reader's. Other types and byte orders, and datasets with no storage or external data, keep converting through the byte readers; unallocated chunks read as the fill value as before. tests/chunked_read_paths_interop.rs checks every chunked read path (File twice, so cached; from_bytes; MmapFile; LazyFile; small, strided and point selections; with and without the parallel feature) against h5py for 1-8 byte integers and 2-8 byte floats in both byte orders, through deflate, shuffle, Fletcher32, LZF, SZIP and Blosc, with partial edge chunks, sparse datasets with default and non-default fill values, and datasets larger than the chunk cache. A filter this build lacks must be an error (or, when an optional filter declined every chunk, the right data). Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/clawhdf5-format/src/data_read.rs | 114 ++++++ crates/clawhdf5/src/lazy.rs | 42 ++ crates/clawhdf5/src/mmap_file.rs | 42 ++ crates/clawhdf5/src/reader.rs | 43 +++ .../tests/chunked_read_paths_interop.rs | 361 ++++++++++++++++++ 5 files changed, 602 insertions(+) create mode 100644 crates/clawhdf5/tests/chunked_read_paths_interop.rs diff --git a/crates/clawhdf5-format/src/data_read.rs b/crates/clawhdf5-format/src/data_read.rs index 8aec190..fd7054c 100644 --- a/crates/clawhdf5-format/src/data_read.rs +++ b/crates/clawhdf5-format/src/data_read.rs @@ -756,6 +756,120 @@ pub fn read_selection_native( crate::gather::gather::(raw, dims, elem_size, selection).map(Some) } +/// The bytes of a slice of [`NativeElement`]s. +#[cfg(feature = "std")] +fn bytes_of_mut(values: &mut [T]) -> &mut [u8] { + // SAFETY: `T: NativeElement` has no padding and every bit pattern is a + // valid value, so its storage may be viewed, and written, as bytes; the + // byte slice covers exactly the values' storage and borrows it + // exclusively for its lifetime. + unsafe { + core::slice::from_raw_parts_mut( + values.as_mut_ptr().cast::(), + core::mem::size_of_val(values), + ) + } +} + +/// `count` zeroed values of `T`, from zeroed pages where the allocator can +/// (see [`crate::chunked_read::alloc_output`]) and backed by huge pages when +/// large. A size taken from the file surfaces as an error, not an abort. +#[cfg(feature = "std")] +fn alloc_zeroed_values(count: usize) -> Result, FormatError> { + if count == 0 || core::mem::size_of::() == 0 { + return Ok(Vec::new()); + } + let failed = || { + FormatError::Overflow(format!( + "cannot allocate {count} values of {} bytes for dataset output", + core::mem::size_of::() + )) + }; + let layout = core::alloc::Layout::array::(count).map_err(|_| failed())?; + // SAFETY: `layout` has non-zero size (count > 0, T not zero-sized). + let ptr = unsafe { std::alloc::alloc_zeroed(layout) }; + if ptr.is_null() { + return Err(failed()); + } + crate::bulk_alloc::advise_huge_pages(ptr, layout.size()); + // SAFETY: allocated by the global allocator with the layout of + // `[T; count]`, which is what `Vec` with capacity `count` frees; all + // bytes are zero, a valid `T` (`NativeElement`: any bit pattern is). + Ok(unsafe { Vec::from_raw_parts(ptr.cast::(), count, count) }) +} + +/// Read a whole chunked dataset that stores `T` natively +/// ([`NativeElement::is_native`]) straight into a `Vec`: each chunk is +/// decoded and copied to its place in the typed output, with no byte buffer +/// to convert from afterwards. Unallocated chunks read as the dataset's fill +/// value, as [`crate::fill_value::read_full_with_fill`] makes them. +/// +/// `Ok(None)` when this does not apply — the datatype is not `T`'s native +/// representation (another type, another byte order: the caller converts +/// through the byte readers and the `read_as_*` functions), the layout is +/// not chunked, no storage is allocated, or the data lives in external +/// files. `cache` is the file's chunk cache, used as +/// [`crate::chunked_read::read_chunked_data_cached`] uses it. +#[cfg(feature = "std")] +#[allow(clippy::too_many_arguments)] +pub fn read_chunked_native( + messages: &[crate::object_header::HeaderMessage], + file_data: &[u8], + layout: &DataLayout, + dataspace: &Dataspace, + datatype: &Datatype, + pipeline: Option<&FilterPipeline>, + offset_size: u8, + length_size: u8, + cache: Option<&ChunkCache>, +) -> Result>, FormatError> { + use crate::fill_value; + use crate::message_type::MessageType; + + if !T::is_native(datatype) + || !matches!(layout, DataLayout::Chunked { .. }) + || !fill_value::has_storage(layout) + || messages + .iter() + .any(|m| m.msg_type == MessageType::ExternalDataFiles) + { + return Ok(None); + } + let size = core::mem::size_of::(); + let mut values = crate::chunked_read::read_chunked_full( + file_data, + layout, + dataspace, + datatype, + pipeline, + offset_size, + length_size, + cache, + |total_bytes| { + if !total_bytes.is_multiple_of(size) { + return Err(FormatError::DataSizeMismatch { + expected: total_bytes.next_multiple_of(size), + actual: total_bytes, + }); + } + alloc_zeroed_values::(total_bytes / size) + }, + |values| bytes_of_mut(values), + )?; + let fill = fill_value::dataset_fill_value_in(file_data, messages, offset_size, length_size)?; + fill_value::apply_to_unallocated_chunks( + bytes_of_mut(&mut values), + file_data, + layout, + dataspace, + size, + fill.as_deref(), + offset_size, + length_size, + )?; + Ok(Some(values)) +} + /// Convert raw bytes to `f64` values. pub fn read_as_f64(raw: &[u8], datatype: &Datatype) -> Result, FormatError> { // Array datatypes read as a flat sequence of their base elements, and diff --git a/crates/clawhdf5/src/lazy.rs b/crates/clawhdf5/src/lazy.rs index a33b9bb..c3b779e 100644 --- a/crates/clawhdf5/src/lazy.rs +++ b/crates/clawhdf5/src/lazy.rs @@ -349,6 +349,9 @@ impl<'f, R: HDF5Read> LazyDataset<'f, R> { /// Read all data as `f64` values. pub fn read_f64(&self) -> Result, Error> { + if let Some(values) = self.read_chunked_native::()? { + return Ok(values); + } let raw = self.read_raw()?; let dt = self.datatype()?; Ok(data_read::read_as_f64(&raw, &dt)?) @@ -396,6 +399,9 @@ impl<'f, R: HDF5Read> LazyDataset<'f, R> { /// Read all data as `f32` values. pub fn read_f32(&self) -> Result, Error> { + if let Some(values) = self.read_chunked_native::()? { + return Ok(values); + } let raw = self.read_raw()?; let dt = self.datatype()?; Ok(data_read::read_as_f32(&raw, &dt)?) @@ -403,6 +409,9 @@ impl<'f, R: HDF5Read> LazyDataset<'f, R> { /// Read all data as `i32` values. pub fn read_i32(&self) -> Result, Error> { + if let Some(values) = self.read_chunked_native::()? { + return Ok(values); + } let raw = self.read_raw()?; let dt = self.datatype()?; Ok(data_read::read_as_i32(&raw, &dt)?) @@ -410,6 +419,9 @@ impl<'f, R: HDF5Read> LazyDataset<'f, R> { /// Read all data as `i64` values. pub fn read_i64(&self) -> Result, Error> { + if let Some(values) = self.read_chunked_native::()? { + return Ok(values); + } let raw = self.read_raw()?; let dt = self.datatype()?; Ok(data_read::read_as_i64(&raw, &dt)?) @@ -417,6 +429,9 @@ impl<'f, R: HDF5Read> LazyDataset<'f, R> { /// Read all data as `u64` values. pub fn read_u64(&self) -> Result, Error> { + if let Some(values) = self.read_chunked_native::()? { + return Ok(values); + } let raw = self.read_raw()?; let dt = self.datatype()?; Ok(data_read::read_as_u64(&raw, &dt)?) @@ -550,6 +565,33 @@ impl<'f, R: HDF5Read> LazyDataset<'f, R> { .transpose() } + /// A chunked dataset that stores `T` natively, decoded straight into a + /// `Vec` (no byte buffer to convert); `None` for any other dataset + /// (see [`data_read::read_chunked_native`]). + fn read_chunked_native(&self) -> Result>, Error> { + let dl = self.data_layout()?; + if !matches!(dl, DataLayout::Chunked { .. }) { + return Ok(None); + } + let dt = self.datatype()?; + if !T::is_native(&dt) { + return Ok(None); + } + let ds = self.dataspace()?; + let pipeline = self.filter_pipeline()?; + Ok(data_read::read_chunked_native::( + &self.header.messages, + self.file.hdf5_bytes(), + &dl, + &ds, + &dt, + pipeline.as_ref(), + self.file.offset_size(), + self.file.length_size(), + None, + )?) + } + fn read_raw(&self) -> Result, Error> { let dt = self.datatype()?; let ds = self.dataspace()?; diff --git a/crates/clawhdf5/src/mmap_file.rs b/crates/clawhdf5/src/mmap_file.rs index 76d9ea1..7813b9c 100644 --- a/crates/clawhdf5/src/mmap_file.rs +++ b/crates/clawhdf5/src/mmap_file.rs @@ -276,6 +276,9 @@ impl<'f> MmapDataset<'f> { /// Read all data as `f64` values. pub fn read_f64(&self) -> Result, Error> { + if let Some(values) = self.read_chunked_native::()? { + return Ok(values); + } let raw = self.read_raw()?; let dt = self.datatype()?; Ok(data_read::read_as_f64(&raw, &dt)?) @@ -310,6 +313,9 @@ impl<'f> MmapDataset<'f> { /// Read all data as `f32` values. pub fn read_f32(&self) -> Result, Error> { + if let Some(values) = self.read_chunked_native::()? { + return Ok(values); + } let raw = self.read_raw()?; let dt = self.datatype()?; Ok(data_read::read_as_f32(&raw, &dt)?) @@ -317,6 +323,9 @@ impl<'f> MmapDataset<'f> { /// Read all data as `i32` values. pub fn read_i32(&self) -> Result, Error> { + if let Some(values) = self.read_chunked_native::()? { + return Ok(values); + } let raw = self.read_raw()?; let dt = self.datatype()?; Ok(data_read::read_as_i32(&raw, &dt)?) @@ -324,6 +333,9 @@ impl<'f> MmapDataset<'f> { /// Read all data as `i64` values. pub fn read_i64(&self) -> Result, Error> { + if let Some(values) = self.read_chunked_native::()? { + return Ok(values); + } let raw = self.read_raw()?; let dt = self.datatype()?; Ok(data_read::read_as_i64(&raw, &dt)?) @@ -331,6 +343,9 @@ impl<'f> MmapDataset<'f> { /// Read all data as `u64` values. pub fn read_u64(&self) -> Result, Error> { + if let Some(values) = self.read_chunked_native::()? { + return Ok(values); + } let raw = self.read_raw()?; let dt = self.datatype()?; Ok(data_read::read_as_u64(&raw, &dt)?) @@ -496,6 +511,33 @@ impl<'f> MmapDataset<'f> { .transpose() } + /// A chunked dataset that stores `T` natively, decoded straight into a + /// `Vec` (no byte buffer to convert); `None` for any other dataset + /// (see [`data_read::read_chunked_native`]). + fn read_chunked_native(&self) -> Result>, Error> { + let dl = self.data_layout()?; + if !matches!(dl, DataLayout::Chunked { .. }) { + return Ok(None); + } + let dt = self.datatype()?; + if !T::is_native(&dt) { + return Ok(None); + } + let ds = self.dataspace()?; + let pipeline = self.filter_pipeline()?; + Ok(data_read::read_chunked_native::( + &self.header.messages, + self.file.hdf5_bytes(), + &dl, + &ds, + &dt, + pipeline.as_ref(), + self.file.offset_size(), + self.file.length_size(), + None, + )?) + } + fn read_raw(&self) -> Result, Error> { let dt = self.datatype()?; let ds = self.dataspace()?; diff --git a/crates/clawhdf5/src/reader.rs b/crates/clawhdf5/src/reader.rs index 90faf11..446a5a7 100644 --- a/crates/clawhdf5/src/reader.rs +++ b/crates/clawhdf5/src/reader.rs @@ -507,6 +507,9 @@ impl<'f> Dataset<'f> { if let Ok(Some(bytes)) = self.read_raw_ref() { return Ok(data_read::read_as_f64(bytes, &dt)?); } + if let Some(values) = self.read_chunked_native::()? { + return Ok(values); + } let raw = self.read_raw()?; Ok(data_read::read_as_f64(&raw, &dt)?) } @@ -524,6 +527,9 @@ impl<'f> Dataset<'f> { if let Ok(Some(bytes)) = self.read_raw_ref() { return Ok(data_read::read_as_f32(bytes, &dt)?); } + if let Some(values) = self.read_chunked_native::()? { + return Ok(values); + } let raw = self.read_raw()?; Ok(data_read::read_as_f32(&raw, &dt)?) } @@ -536,6 +542,9 @@ impl<'f> Dataset<'f> { if let Ok(Some(bytes)) = self.read_raw_ref() { return Ok(data_read::read_as_i32(bytes, &dt)?); } + if let Some(values) = self.read_chunked_native::()? { + return Ok(values); + } let raw = self.read_raw()?; Ok(data_read::read_as_i32(&raw, &dt)?) } @@ -548,6 +557,9 @@ impl<'f> Dataset<'f> { if let Ok(Some(bytes)) = self.read_raw_ref() { return Ok(data_read::read_as_i64(bytes, &dt)?); } + if let Some(values) = self.read_chunked_native::()? { + return Ok(values); + } let raw = self.read_raw()?; Ok(data_read::read_as_i64(&raw, &dt)?) } @@ -560,6 +572,9 @@ impl<'f> Dataset<'f> { if let Ok(Some(bytes)) = self.read_raw_ref() { return Ok(data_read::read_as_u64(bytes, &dt)?); } + if let Some(values) = self.read_chunked_native::()? { + return Ok(values); + } let raw = self.read_raw()?; Ok(data_read::read_as_u64(&raw, &dt)?) } @@ -1050,6 +1065,34 @@ impl<'f> Dataset<'f> { .transpose() } + /// A chunked dataset that stores `T` natively, decoded straight into a + /// `Vec` through the file's chunk cache (no byte buffer to convert); + /// `None` for any other dataset (see + /// [`data_read::read_chunked_native`]). + fn read_chunked_native(&self) -> Result>, Error> { + let dl = self.data_layout()?; + if !matches!(dl, DataLayout::Chunked { .. }) { + return Ok(None); + } + let dt = self.datatype()?; + if !T::is_native(&dt) { + return Ok(None); + } + let ds = self.dataspace()?; + let pipeline = self.filter_pipeline()?; + Ok(data_read::read_chunked_native::( + &self.header.messages, + self.file.data.as_bytes(), + &dl, + &ds, + &dt, + pipeline.as_ref(), + self.file.offset_size(), + self.file.length_size(), + Some(&self.file.chunk_cache), + )?) + } + fn read_raw(&self) -> Result, Error> { let dt = self.datatype()?; let ds = self.dataspace()?; diff --git a/crates/clawhdf5/tests/chunked_read_paths_interop.rs b/crates/clawhdf5/tests/chunked_read_paths_interop.rs new file mode 100644 index 0000000..2865a22 --- /dev/null +++ b/crates/clawhdf5/tests/chunked_read_paths_interop.rs @@ -0,0 +1,361 @@ +//! Every full and selection read path of chunked datasets against h5py. +//! +//! h5py (libhdf5) writes chunked datasets of every numeric type the typed +//! readers cover, in both byte orders, through deflate, shuffle, +//! Fletcher32, LZF, SZIP and Blosc, in 1-3 dimensional shapes whose chunks +//! do not divide them (partial edge chunks), plus sparse datasets whose +//! unwritten chunks read as a fill value. Next to each it stores the values +//! as contiguous `f64`, and checks that h5py reads the chunked dataset back +//! as those values. +//! +//! clawhdf5 must read every chunked dataset as those values through every +//! reader: `File` (the chunk-cached reader, twice so the second read can hit +//! the cache, and the typed readers that decode straight into their output), +//! `File::from_bytes`, `MmapFile`, `LazyFile`, and the selection readers +//! (a small hyperslab, a strided one covering most of the dataset, points). +//! With `--features parallel` the same reads decode chunks on several +//! threads. A filter this build does not include must be an error, never +//! data. +//! +//! Skipped when python3 with h5py is unavailable, unless +//! `CLAWHDF5_REQUIRE_INTEROP=1`. + +use std::process::Command; + +use clawhdf5::{File, LazyFile, MmapFile, Selection}; + +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") +} + +/// Whether the interop test can run; panics instead of skipping when +/// `CLAWHDF5_REQUIRE_INTEROP=1`. +fn have_python() -> bool { + let ok = Command::new(python()) + .args(["-c", "import h5py, numpy"]) + .output() + .map(|o| o.status.success()) + .unwrap_or(false); + if ok { + return true; + } + assert!( + !interop_required(), + "CLAWHDF5_REQUIRE_INTEROP=1 but python3 with h5py is not available" + ); + eprintln!("SKIP: python3 with h5py not available"); + false +} + +/// Writes `c/` (chunked) and `e/` (the expected values, `f64`, +/// contiguous) for every case; prints one line per case: +/// `name filter_ids(comma-separated or -)`. +const GENERATE: &str = r#" +import sys +import numpy as np, h5py +try: + import hdf5plugin +except ImportError: + hdf5plugin = None +path = sys.argv[1] + +dtypes = [] +for code in ['i1', 'u1', 'i2', 'u2', 'i4', 'u4', 'i8', 'u8', 'f2', 'f4', 'f8']: + orders = ['|'] if code[1] == '1' else ['<', '>'] + dtypes += [np.dtype(o + code) for o in orders] + +filters = { + 'none': {}, + 'gzip': dict(compression='gzip', compression_opts=4), + 'shuffle_gzip': dict(shuffle=True, compression='gzip', compression_opts=1), + # h5py puts the checksum last (applied after deflate). + 'shuffle_gzip_fletcher': dict(shuffle=True, compression='gzip', fletcher32=True), + 'fletcher': dict(fletcher32=True), + 'shuffle_lzf': dict(shuffle=True, compression='lzf'), +} +if h5py.h5z.filter_avail(h5py.h5z.FILTER_SZIP): + filters['szip'] = dict(compression='szip', compression_opts=('nn', 8)) +if hdf5plugin is not None: + filters['blosc'] = dict(**hdf5plugin.Blosc(cname='lz4', clevel=5, + shuffle=hdf5plugin.Blosc.SHUFFLE)) + +shapes = [((37, 23), (8, 5)), ((101,), (16,)), ((9, 10, 11), (4, 3, 5))] + +def values(n, dt): + i = np.arange(n, dtype=np.int64) + if dt.kind == 'f': + v = ((i * 7 + 3) % 1000) / 8.0 - 60.0 + elif dt.kind == 'i': + v = (i * 7 + 3) % 200 - 100 + else: + v = (i * 7 + 3) % 250 + return v.astype(dt) + +def filter_ids(dset): + plist = dset.id.get_create_plist() + ids = [str(plist.get_filter(k)[0]) for k in range(plist.get_nfilters())] + return ','.join(ids) or '-' + +with h5py.File(path, 'w') as f: + n = 0 + for dt in dtypes: + for fname, fopts in filters.items(): + for s, (shape, chunks) in enumerate(shapes): + name = f"{dt.str.replace('|', 'x').replace('<', 'le').replace('>', 'be')}_{fname}_{s}" + data = values(int(np.prod(shape)), dt).reshape(shape) + try: + d = f.create_dataset('c/' + name, data=data, chunks=chunks, **fopts) + except (ValueError, TypeError) as e: + continue # a filter that refuses this type + f.create_dataset('e/' + name, data=data.astype('i4'): + name = f"{dt.str.replace('<', 'le').replace('>', 'be')}_large" + shape = (2600, 2048) + data = values(shape[0] * shape[1], dt).reshape(shape) + d = f.create_dataset('c/' + name, data=data, chunks=(256, 256), + shuffle=True, compression='gzip', compression_opts=1) + f.create_dataset('e/' + name, data=data.astype('', 'be')}_sparse_{fname}_{fill}" + shape, chunks = (37, 23), (8, 5) + kw = dict(filters[fname]) + if fill is not None: + kw['fillvalue'] = np.array(fill, dtype=dt) + d = f.create_dataset('c/' + name, shape=shape, dtype=dt, chunks=chunks, **kw) + full = values(37 * 23, dt).reshape(shape) + d[3:17, 4:12] = full[3:17, 4:12] + d[30:, 20:] = full[30:, 20:] + expect = d[()] + want = np.full(shape, 0 if fill is None else fill, dtype=dt) + want[3:17, 4:12] = full[3:17, 4:12] + want[30:, 20:] = full[30:, 20:] + assert np.array_equal(expect, want), name + f.create_dataset('e/' + name, data=want.astype(' bool { + ids == "-" + || ids.split(',').all(|id| { + clawhdf5_format::filter_registry::is_filter_available(id.parse().expect("filter id")) + }) +} + +/// The expected values converted as libhdf5 converts them for each typed +/// reader (every case's values are exact in `f32` and within `i32`). +struct Expected { + f64s: Vec, +} + +impl Expected { + fn f32s(&self) -> Vec { + self.f64s.iter().map(|&v| v as f32).collect() + } + /// Truncation toward zero; negative values read as unsigned are 0. + fn i32s(&self) -> Vec { + self.f64s.iter().map(|&v| v as i32).collect() + } + fn i64s(&self) -> Vec { + self.f64s.iter().map(|&v| v as i64).collect() + } + fn u64s(&self) -> Vec { + self.f64s.iter().map(|&v| v as u64).collect() + } + fn select(&self, idx: &[usize]) -> Expected { + Expected { + f64s: idx.iter().map(|&i| self.f64s[i]).collect(), + } + } +} + +/// The typed readers' results for one dataset, compared with `want`. +macro_rules! check_typed { + ($ds:expr, $want:expr, $name:expr, $path:expr) => {{ + let (ds, want, name, path) = (&$ds, &$want, $name, $path); + assert_eq!(ds.read_f64().unwrap(), want.f64s, "{name} {path} f64"); + assert_eq!(ds.read_f32().unwrap(), want.f32s(), "{name} {path} f32"); + assert_eq!(ds.read_i32().unwrap(), want.i32s(), "{name} {path} i32"); + assert_eq!(ds.read_i64().unwrap(), want.i64s(), "{name} {path} i64"); + assert_eq!(ds.read_u64().unwrap(), want.u64s(), "{name} {path} u64"); + }}; +} + +/// Row-major indices of the elements `sel` picks from a dataset of `dims`. +fn selected(sel: &Selection, dims: &[u64]) -> Vec { + let strides: Vec = (0..dims.len()) + .map(|d| dims[d + 1..].iter().product::() as usize) + .collect(); + match sel { + Selection::Hyperslab { + start, + stride, + count, + block, + } => { + // Per dimension, the selected coordinates in order. + let axes: Vec> = (0..dims.len()) + .map(|d| { + (0..count[d]) + .flat_map(|c| (0..block[d]).map(move |b| start[d] + c * stride[d] + b)) + .collect() + }) + .collect(); + let mut out = vec![0usize]; + for (d, axis) in axes.iter().enumerate() { + let stride = strides[d]; + out = out + .iter() + .flat_map(|&base| axis.iter().map(move |&x| base + x as usize * stride)) + .collect(); + } + out + } + Selection::Points(points) => points + .iter() + .map(|p| p.iter().zip(&strides).map(|(&x, &s)| x as usize * s).sum()) + .collect(), + _ => unreachable!(), + } +} + +/// A small box (read through the partial-read path), a strided selection +/// covering most of the dataset (read in full, then selected), and points. +fn selections(dims: &[u64]) -> Vec { + let rank = dims.len(); + let small = Selection::Hyperslab { + start: dims.iter().map(|&d| d / 3).collect(), + stride: vec![1; rank], + count: dims.iter().map(|&d| (d / 4).max(1)).collect(), + block: vec![1; rank], + }; + let strided = Selection::Hyperslab { + start: vec![0; rank], + stride: vec![2; rank], + count: dims.iter().map(|&d| d.div_ceil(2)).collect(), + block: vec![1; rank], + }; + let points = Selection::Points(vec![ + vec![0; rank], + dims.iter().map(|&d| d - 1).collect(), + dims.iter().map(|&d| d / 2).collect(), + dims.iter().map(|&d| (d * 2) / 3).collect(), + ]); + vec![small, strided, points] +} + +#[test] +fn chunked_reads_match_h5py_on_every_path() { + if !have_python() { + return; + } + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("chunked_paths.h5"); + let out = Command::new(python()) + .arg("-c") + .arg(GENERATE) + .arg(&path) + .output() + .expect("run python"); + assert!( + out.status.success(), + "generator failed: {}", + String::from_utf8_lossy(&out.stderr) + ); + let cases: Vec<(String, String)> = String::from_utf8(out.stdout) + .unwrap() + .lines() + .map(|l| { + let (name, ids) = l.split_once(' ').unwrap(); + (name.to_string(), ids.to_string()) + }) + .collect(); + assert!(cases.len() > 300, "only {} cases", cases.len()); + if interop_required() { + // The generator must have covered the plugin filters too. + for f in ["_szip_", "_blosc_", "_shuffle_lzf_", "_sparse_"] { + assert!(cases.iter().any(|(n, _)| n.contains(f)), "no {f} case"); + } + } + + let file = File::open(&path).unwrap(); + let owned = File::from_bytes(std::fs::read(&path).unwrap()).unwrap(); + let mmap = MmapFile::open(&path).unwrap(); + let lazy = LazyFile::open_mmap(&path).unwrap(); + let (mut checked, mut refused) = (0, 0); + for (name, ids) in &cases { + let chunked = format!("c/{name}"); + let want = Expected { + f64s: file + .dataset(&format!("e/{name}")) + .unwrap() + .read_f64() + .unwrap(), + }; + let ds = file.dataset(&chunked).unwrap(); + if !filters_available(ids) { + // Unsupported filter: an error from every reader, never wrong + // data. (An optional filter, as Blosc is, may have declined every + // chunk — they are then stored as-is and read fine.) + let results = [ + ds.read_f64(), + owned.dataset(&chunked).unwrap().read_f64(), + mmap.dataset(&chunked).unwrap().read_f64(), + lazy.dataset(&chunked).unwrap().read_f64(), + ]; + let errors = results.iter().filter(|r| r.is_err()).count(); + assert!(errors == 0 || errors == results.len(), "{name}"); + for values in results.into_iter().flatten() { + assert_eq!(values, want.f64s, "{name}"); + } + if errors > 0 { + assert!(ds.read_f32().is_err(), "{name}"); + refused += 1; + continue; + } + } + // Twice: the second read of a small dataset comes from the cache. + check_typed!(ds, want, name, "File"); + check_typed!(ds, want, name, "File (cached)"); + check_typed!(owned.dataset(&chunked).unwrap(), want, name, "from_bytes"); + check_typed!(mmap.dataset(&chunked).unwrap(), want, name, "MmapFile"); + check_typed!(lazy.dataset(&chunked).unwrap(), want, name, "LazyFile"); + + let dims = ds.shape().unwrap(); + for sel in selections(&dims) { + let want = want.select(&selected(&sel, &dims)); + assert_eq!( + ds.read_f64_selection(&sel).unwrap(), + want.f64s, + "{name} {sel:?}" + ); + assert_eq!( + ds.read_f32_selection(&sel).unwrap(), + want.f32s(), + "{name} {sel:?}" + ); + assert_eq!( + ds.read_i64_selection(&sel).unwrap(), + want.i64s(), + "{name} {sel:?}" + ); + } + checked += 1; + } + eprintln!("{checked} datasets read on every path, {refused} refused (filter not built in)"); + assert!(checked > 300); +} From 9e608b975c7deb219dcdb8f34307fa29683088b4 Mon Sep 17 00:00:00 2001 From: osobh Date: Sat, 26 Sep 2026 10:20:07 -0500 Subject: [PATCH 4/7] format: decode selected chunks into reusable buffers A selection read (a hyperslab or points covering at most half the dataset) decoded each chunk it overlaps into fresh buffers, one per filter stage; it now uses the thread's chunk-decoding scratch like the full readers. Covered by tests/chunked_read_paths_interop.rs (small hyperslabs and points over every filter and type) and the partial-read equivalence tests. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/clawhdf5-format/src/chunked_read.rs | 2 +- crates/clawhdf5-format/src/partial_read.rs | 99 +++++++++++----------- 2 files changed, 52 insertions(+), 49 deletions(-) diff --git a/crates/clawhdf5-format/src/chunked_read.rs b/crates/clawhdf5-format/src/chunked_read.rs index 7d3efd0..712aab8 100644 --- a/crates/clawhdf5-format/src/chunked_read.rs +++ b/crates/clawhdf5-format/src/chunked_read.rs @@ -31,7 +31,7 @@ use crate::lane_partition::PartitionStats; /// [`DecodeScratch`]), kept between reads so decoding reuses memory instead /// of faulting in fresh pages for every chunk. A re-entrant call (a /// registered filter codec that itself reads a file) gets a fresh scratch. -fn with_scratch(f: impl FnOnce(&mut DecodeScratch) -> R) -> R { +pub(crate) fn with_scratch(f: impl FnOnce(&mut DecodeScratch) -> R) -> R { #[cfg(feature = "std")] { use std::cell::RefCell; diff --git a/crates/clawhdf5-format/src/partial_read.rs b/crates/clawhdf5-format/src/partial_read.rs index 9865c46..0151255 100644 --- a/crates/clawhdf5-format/src/partial_read.rs +++ b/crates/clawhdf5-format/src/partial_read.rs @@ -24,7 +24,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::{all_filters_skipped, decompress_chunk_exact}; +use crate::filters::{all_filters_skipped, decompress_chunk_exact_with}; use crate::selection::Selection; /// The smallest axis-aligned box containing every selected element, as @@ -305,54 +305,57 @@ pub fn read_selection( let rank = dims.len(); let chunk_shape: Vec = chunk_dims.iter().map(|&d| d as u64).collect(); let chunk_bytes = crate::chunked_read::checked_chunk_byte_len(&chunk_dims, elem_size)?; - for chunk in &chunks { - if chunk.offsets.len() < rank || chunk.address == u64::MAX { - continue; - } - let origin = &chunk.offsets[..rank]; - let overlaps = (0..rank).all(|d| { - origin[d] < box_start[d] + box_extent[d] - && origin[d].saturating_add(chunk_shape[d]) > box_start[d] - }); - if !overlaps { - continue; - } - let at = usize::try_from(chunk.address) - .map_err(|_| FormatError::Overflow("chunk address exceeds usize".into()))?; - let raw = at - .checked_add(chunk.chunk_size as usize) - .and_then(|end| file_data.get(at..end)) - .ok_or(FormatError::UnexpectedEof { - expected: at.saturating_add(chunk.chunk_size as usize), - available: file_data.len(), - })?; - // 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 !all_filters_skipped(pl, chunk.filter_mask) => { - decoded = decompress_chunk_exact( - raw, - pl, - chunk_bytes, - elem_size as u32, - chunk.filter_mask, - &chunk.offsets[..rank], - )?; - &decoded + // Chunks are decoded into this thread's reusable buffers. + crate::chunked_read::with_scratch(|scratch| -> Result<(), FormatError> { + for chunk in &chunks { + if chunk.offsets.len() < rank || chunk.address == u64::MAX { + continue; } - _ => raw, - }; - copy_overlap( - data, - origin, - &chunk_shape, - &mut boxed, - &box_start, - &box_extent, - elem_size, - ); - } + let origin = &chunk.offsets[..rank]; + let overlaps = (0..rank).all(|d| { + origin[d] < box_start[d] + box_extent[d] + && origin[d].saturating_add(chunk_shape[d]) > box_start[d] + }); + if !overlaps { + continue; + } + let at = usize::try_from(chunk.address) + .map_err(|_| FormatError::Overflow("chunk address exceeds usize".into()))?; + let raw = at + .checked_add(chunk.chunk_size as usize) + .and_then(|end| file_data.get(at..end)) + .ok_or(FormatError::UnexpectedEof { + expected: at.saturating_add(chunk.chunk_size as usize), + available: file_data.len(), + })?; + // Mirrors the full-read path: filter-mask bit i set means + // filter i was not applied to this chunk. + let data: &[u8] = match pipeline { + Some(pl) if !all_filters_skipped(pl, chunk.filter_mask) => { + decompress_chunk_exact_with( + raw, + pl, + chunk_bytes, + elem_size as u32, + chunk.filter_mask, + &chunk.offsets[..rank], + scratch, + )? + } + _ => raw, + }; + copy_overlap( + data, + origin, + &chunk_shape, + &mut boxed, + &box_start, + &box_extent, + elem_size, + ); + } + Ok(()) + })?; } _ => return Ok(None), } From 94b6df986c60dfba77069da516188cceec90abe4 Mon Sep 17 00:00:00 2001 From: osobh Date: Sat, 26 Sep 2026 10:36:44 -0500 Subject: [PATCH 5/7] docs: chunked full reads decode in place; small pools no longer block CHANGELOG entry for the chunked read changes, and the known-issues entry on concurrent chunked reads updated: both causes it names (per-read page faults, readers waiting on a small pool) are fixed; the 16-thread comparison with h5py stays open until re-measured on an idle machine. Co-Authored-By: Claude Opus 5.5 (1M context) --- CHANGELOG.md | 46 ++++++++++++++++++++++++++++++++++++++++++++ docs/known-issues.md | 21 +++++++++++++++----- 2 files changed, 62 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 29804fc..a877c30 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,52 @@ ## Unreleased +### Chunked full reads (2026-09-26) +- **Chunks are decoded straight into the output, into reused buffers.** A + full read of a chunked dataset faulted in about three times its size in + fresh pages: every chunk was decoded into a new buffer per filter stage + (the cached reader behind `read_*` decoded 128 chunks at a time before + placing any; the uncached one behind `MmapFile`, `LazyFile` and + `verify_provenance` decoded the whole dataset first), then assembled into + a byte buffer, which the typed readers copied once more. Now each chunk + is decoded into buffers the thread keeps between chunks and reads + (`clawhdf5_format::filters::DecodeScratch`, + `decompress_chunk_exact_with`: deflate inflates into a kept buffer with a + reset inflater, shuffle into the other one, Fletcher32 is checked in + place; other filters go through the registry as before) and copied + directly to its place in the output. Chunks still go into the file's + chunk cache when the whole dataset fits. Selection reads decode the + chunks they touch the same way. +- **Typed full reads of chunked data skip the byte buffer.** `read_f32`, + `read_f64`, `read_i32`, `read_i64` and `read_u64` (on `File`, `MmapFile` + and `LazyFile`) of a chunked dataset stored as that type in native byte + order decode every chunk into the returned `Vec` (huge-page backed when + large, like the byte readers' output); other types and byte orders + convert as before. New public + `clawhdf5_format::data_read::read_chunked_native`. +- **Reading threads no longer wait for a busy rayon pool.** A full read + handed its chunks to rayon and the calling thread slept until the pool had + decoded them, so with a small pool (2-4 threads) readers outside it queued + behind its workers. The calling thread now decodes too, and pool workers + join in only when free; a helper the pool starts after the read has + finished returns at once. A single read still spreads over the default + pool. This replaces the one-thread-pool special case below for full + reads. Chunks are placed from several threads only when the chunk index + puts them on the chunk grid at distinct places; a corrupt index is read + one chunk at a time, and the error reported is still the first failing + chunk's. New test `crates/clawhdf5/tests/busy_decode_pool.rs`. +- **Fixed:** in a filtered dataset, a chunk stored with every filter skipped + (filter mask) and shorter than a chunk read with zeros in place of its + missing part through `File`'s `read_*`; it is now an error naming the + chunk, as `MmapFile`/`LazyFile` already made it. +- New h5py comparison `crates/clawhdf5/tests/chunked_read_paths_interop.rs`: + every chunked read path (cached and uncached full reads, `MmapFile`, + `LazyFile`, small, strided and point selections, with and without the + `parallel` feature) for 1-8-byte integers and 2-8-byte floats in both + byte orders, through deflate, shuffle, Fletcher32, LZF, SZIP and Blosc, + with partial edge chunks, sparse datasets and fill values, and datasets + larger than the chunk cache. + ### Concurrent reads (2026-09-26) - **Full reads of chunked datasets scale with threads again when rayon's pool has one thread.** Each full read handed its chunks to rayon to diff --git a/docs/known-issues.md b/docs/known-issues.md index 2657cdf..0d2556b 100644 --- a/docs/known-issues.md +++ b/docs/known-issues.md @@ -26,8 +26,9 @@ performance" below. ## Concurrent and contiguous read performance (measured 2026-09-26) -**Status:** open for chunked full reads (one cause fixed 2026-09-26); the -contiguous item is fixed (2026-09-26). Measured on +**Status:** open for chunked full reads at 16 threads until re-measured +(the causes identified below are fixed as of 2026-09-26); the contiguous +item is fixed (2026-09-26). Measured on tank with `concurrent_read` against h5py 3.16 / HDF5 2.0 (`BENCHMARKS.md`, "Concurrent reads"): - **Partly fixed 2026-09-26.** Full reads of chunked datasets from several threads @@ -48,9 +49,19 @@ tank with `concurrent_read` against h5py 3.16 / HDF5 2.0 (`BENCHMARKS.md`, readers outside it still wait on its workers. Datasets larger than the cache's budget were already read without inserting into it, and skipping its lookups entirely gained only a few percent at 16 threads. Remaining - per-read overhead, not yet addressed: each full `read_f32` of a chunked - dataset faults in about three times its size in fresh pages (the output, - the `f32` copy of it, and a new buffer per decoded chunk). + per-read overhead: each full `read_f32` of a chunked dataset faults in + about three times its size in fresh pages (the output, the `f32` copy of + it, and a new buffer per decoded chunk). + **Fixed 2026-09-26** (both causes; see `CHANGELOG.md`, "Chunked full + reads"): chunks are decoded into buffers each thread reuses and copied + straight into the output, and the typed readers decode into the `Vec` + they return, so a full read no longer faults in a buffer per chunk or a + second copy of its output; and the reading + thread decodes its own chunks with pool workers helping when free, so no + reader waits on a small or busy pool + (`crates/clawhdf5/tests/busy_decode_pool.rs`). The 16-thread comparison + with h5py processes has not been re-measured yet (tank was busy with + other work); this item stays open until it is. - Contiguous datasets read 4x slower than h5py on one thread (2.5 vs 9.8 GB/s full, 0.12x for 256 x 256 hyperslabs). **Fixed 2026-09-26** (re-measured on tank at `408f69e`: 13665 MB/s From 8295d016141ea491bfe892b1f7462cab8f0fb4db Mon Sep 17 00:00:00 2001 From: osobh Date: Sat, 26 Sep 2026 10:37:40 -0500 Subject: [PATCH 6/7] format: a chunk offset past usize writes nothing Placing a chunk cast its u64 offsets to usize; on a 32-bit target an offset past the address space wrapped into the output (and could then overlap another chunk's region when chunks are placed concurrently). Such an offset is past the dataset, so it now saturates and the chunk writes nothing, as the concurrent-placement check already assumed. No change on 64-bit targets, where the cast cannot wrap (so no test here). Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/clawhdf5-format/src/chunked_read.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/crates/clawhdf5-format/src/chunked_read.rs b/crates/clawhdf5-format/src/chunked_read.rs index 712aab8..ddadda7 100644 --- a/crates/clawhdf5-format/src/chunked_read.rs +++ b/crates/clawhdf5-format/src/chunked_read.rs @@ -155,7 +155,9 @@ impl ChunkPlacer { &mut long }; for (o, &off) in chunk_offsets.iter_mut().zip(offsets) { - *o = off as usize; + // An offset past `usize` is past the dataset: the chunk writes + // nothing (a plain cast would wrap it into the output). + *o = usize::try_from(off).unwrap_or(usize::MAX); } // SAFETY: the caller's contract. unsafe { From bf4aefcd00abd5629701a207fc42f0fb9d494973 Mon Sep 17 00:00:00 2001 From: osobh Date: Sat, 26 Sep 2026 11:22:07 -0500 Subject: [PATCH 7/7] format: a one-thread pool decodes on the caller only; keep 1 MiB scratch Review follow-ups. With run_with_helpers a one-thread rayon pool gave each read a second core (the caller plus the worker), so --decode-threads 1 no longer matched h5py's one core per call; such a pool now adds no helper. Per-thread decode scratch is kept up to 1 MiB per buffer (was 4 MiB), bounding what never-exiting pool workers hold. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/clawhdf5-format/src/filters.rs | 7 +++++-- crates/clawhdf5-format/src/parallel_read.rs | 6 ++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/crates/clawhdf5-format/src/filters.rs b/crates/clawhdf5-format/src/filters.rs index 872b8d8..266e530 100644 --- a/crates/clawhdf5-format/src/filters.rs +++ b/crates/clawhdf5-format/src/filters.rs @@ -183,8 +183,11 @@ enum Stage { } impl DecodeScratch { - /// Largest buffer [`trim`](Self::trim) keeps (4 MiB). - pub const RETAIN_BYTES: usize = 4 << 20; + /// Largest buffer [`trim`](Self::trim) keeps (1 MiB): enough for common + /// chunk sizes (a 256 x 256 `f32` chunk is 256 KiB) while bounding what + /// every long-lived thread (rayon's workers never exit) holds on to, at + /// two buffers each. + pub const RETAIN_BYTES: usize = 1 << 20; /// An empty scratch; buffers are allocated on first use. pub fn new() -> Self { diff --git a/crates/clawhdf5-format/src/parallel_read.rs b/crates/clawhdf5-format/src/parallel_read.rs index 61d06cd..d9927c5 100644 --- a/crates/clawhdf5-format/src/parallel_read.rs +++ b/crates/clawhdf5-format/src/parallel_read.rs @@ -47,6 +47,12 @@ pub fn pool_can_parallelise() -> bool { /// item beyond the caller's first. pub(crate) fn helper_count(items: usize) -> usize { let pool = rayon::current_num_threads(); + // A one-thread pool means "decode on the calling thread" (the setting + // benchmarks use to compare with h5py, where each call decodes on its + // caller): no helper, so one read never uses two cores. + if pool <= 1 { + return 0; + } let others = if rayon::current_thread_index().is_some() { pool.saturating_sub(1) } else {