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), }