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<R>(f: impl FnOnce(&mut DecodeScratch) -> R) -> R {
+pub(crate) fn with_scratch<R>(f: impl FnOnce(&mut DecodeScratch) -> R) -> R {
     #[cfg(feature = "std")]
     {
         use std::cell::RefCell;

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