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) <[email protected]>
This commit is contained in:
@@ -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;
|
||||
|
||||
@@ -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,6 +305,8 @@ 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)?;
|
||||
// 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;
|
||||
@@ -328,18 +330,17 @@ pub fn read_selection(
|
||||
})?;
|
||||
// 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(
|
||||
decompress_chunk_exact_with(
|
||||
raw,
|
||||
pl,
|
||||
chunk_bytes,
|
||||
elem_size as u32,
|
||||
chunk.filter_mask,
|
||||
&chunk.offsets[..rank],
|
||||
)?;
|
||||
&decoded
|
||||
scratch,
|
||||
)?
|
||||
}
|
||||
_ => raw,
|
||||
};
|
||||
@@ -353,6 +354,8 @@ pub fn read_selection(
|
||||
elem_size,
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
})?;
|
||||
}
|
||||
_ => return Ok(None),
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user