quantumclaw/clawhdf5
3
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

format: bound and batch every chunk fetch over Storage

Browse Source 
Only the full read split its chunk fetches into 64 MiB batches. The
selection path, the indexed read and the parallel_read decoders fetched
every chunk's stored bytes in one read_ranges call, each extent bounded only
by the file length, so a crafted chunk index pointing many chunks at one
large extent made File::open_storage hold chunks x extent bytes (3.3 GB from
a 16.8 MB file) before the first decode error.

- storage::for_each_extent_batch is now the one way raw-data reads fetch
  chunk bytes: batches of at most RAW_BATCH_BYTES (now pub), each decoded
  before the next is fetched. Used by the full, cached, indexed, selection
  and parallel_read paths; the sweep read uses read_extent per chunk.
- ExtentReq carries each chunk's claimed extent (bounds-checked as before,
  same errors) and the prefix actually fetched:
  filters::stored_chunk_limit — the chunk size if unfiltered, else each
  applied filter's worst-case growth (n + n/4 + 4096 per codec; unbounded
  only for an application-registered codec). The in-memory path cuts the
  slice it decodes the same way, so both paths still agree.
- tests/raw_fetch_bounds.rs: a crafted chunked_large.h5 (ten chunks all
  claiming 20 MiB at one padding blob) read through every path over a
  storage that records the largest single fetch; and 160 MiB of legitimate
  unfiltered chunks fetched batch by batch. Before: one 80 MiB fetch
  (selection) and one 160 MiB fetch; after: within the budget.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
osobh
2026-09-26 18:30:30 -05:00
co-authored by Claude Opus 5.5
parent f191dc09d5
commit 7d629f49e3
10 changed files with 825 additions and 285 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/clawhdf5-format/src/chunked_read.rs
+123 -80
View File
@@ -19,7 +19,7 @@ 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_in};
use crate::storage::{ExtentBytes, Storage, Window, raw_batches, read_exact_at};
use crate::storage::{ExtentBytes, ExtentReq, Storage, Window, for_each_extent_batch, read_extent};
#[cfg(feature = "std")]
use std::sync::Arc;
@@ -256,52 +256,17 @@ fn fill_from_chunks<S: Storage + ?Sized>(
cache: CacheUse<'_>,
output: &mut [u8],
) -> Result<(), FormatError> {
let contiguous = file_data.as_contiguous().is_some();
let out = OutBuf::new(output);
let batches = raw_batches(chunks.len(), contiguous, |i| chunks[i].chunk_size as usize);
for batch in batches {
fill_batch(
file_data,
&chunks[batch],
pipeline,
placer,
chunk_total_bytes,
cache,
&out,
)?;
}
Ok(())
}
/// Whether a full read looks chunk `c` up in the cache (see
/// [`fill_from_chunks`]): a filtered chunk of a dataset the cache keeps.
#[cfg(feature = "std")]
fn uses_cache(cache: CacheUse<'_>, pipeline: Option<&FilterPipeline>, c: &ChunkInfo) -> bool {
matches!(cache, Some((_, _, true)))
&& pipeline.is_some_and(|pl| !all_filters_skipped(pl, c.filter_mask))
}
/// [`fill_from_chunks`] for one batch of chunks.
#[allow(clippy::too_many_arguments)]
fn fill_batch<S: Storage + ?Sized>(
file_data: &S,
chunks: &[ChunkInfo],
pipeline: Option<&FilterPipeline>,
placer: &ChunkPlacer,
chunk_total_bytes: usize,
cache: CacheUse<'_>,
out: &OutBuf<'_>,
) -> Result<(), FormatError> {
let rank = placer.rank;
let elem_size = placer.elem_size as u32;
#[cfg(not(feature = "std"))]
let _ = cache;
#[cfg(feature = "std")]
let rank = placer.rank;
// Over a backend without the whole file in memory, the decoded chunks
// the cache holds are taken now (so a chunk evicted before it is placed
// is not left without bytes), and only the others are fetched.
#[cfg(feature = "std")]
let hits: Vec<Option<Arc<CacheAlignedBuffer>>> = match cache {
let hits: Vec<CacheHit> = match cache {
Some((cache, key, true)) if file_data.as_contiguous().is_none() => chunks
.iter()
.map(|c| {
@@ -314,25 +279,90 @@ fn fill_batch<S: Storage + ?Sized>(
.collect(),
_ => Vec::new(),
};
let extents: Vec<(u64, usize, bool)> = if file_data.as_contiguous().is_some() {
Vec::new()
} else {
chunks
.iter()
.enumerate()
.map(|(i, c)| {
#[cfg(feature = "std")]
let wanted = hits.get(i).is_none_or(Option::is_none);
#[cfg(not(feature = "std"))]
let wanted = {
let _ = i;
true
};
(c.address, c.chunk_size as usize, wanted)
})
.collect()
};
let raw_bytes = ExtentBytes::fetch(file_data, &extents)?;
#[cfg(not(feature = "std"))]
let hits: Vec<CacheHit> = Vec::new();
let reqs: Vec<ExtentReq> = chunks
.iter()
.enumerate()
.map(|(i, c)| {
let wanted = hits.get(i).is_none_or(Option::is_none);
chunk_req(c, pipeline, chunk_total_bytes, wanted)
})
.collect();
for_each_extent_batch(file_data, &reqs, |batch, raw_bytes| {
fill_batch(
&chunks[batch.clone()],
&reqs[batch.clone()],
hits.get(batch.clone()).unwrap_or_default(),
batch.start,
raw_bytes,
pipeline,
placer,
chunk_total_bytes,
cache,
&out,
)
})
}
/// The extent of chunk `c`'s stored bytes to fetch (`wanted`) or only
/// bounds-check: its whole stored size is checked against the file, and at
/// most [`crate::filters::stored_chunk_limit`] of it is read and decoded.
pub(crate) fn chunk_req(
c: &ChunkInfo,
pipeline: Option<&FilterPipeline>,
chunk_bytes: usize,
wanted: bool,
) -> ExtentReq {
let len = c.chunk_size as usize;
ExtentReq {
addr: c.address,
len,
fetch: wanted.then(|| {
len.min(crate::filters::stored_chunk_limit(
pipeline,
c.filter_mask,
chunk_bytes,
))
}),
}
}
/// A decoded chunk a full read took from the cache.
#[cfg(feature = "std")]
type CacheHit = Option<Arc<CacheAlignedBuffer>>;
#[cfg(not(feature = "std"))]
type CacheHit = Option<core::convert::Infallible>;
/// Whether a full read looks chunk `c` up in the cache (see
/// [`fill_from_chunks`]): a filtered chunk of a dataset the cache keeps.
#[cfg(feature = "std")]
fn uses_cache(cache: CacheUse<'_>, pipeline: Option<&FilterPipeline>, c: &ChunkInfo) -> bool {
matches!(cache, Some((_, _, true)))
&& pipeline.is_some_and(|pl| !all_filters_skipped(pl, c.filter_mask))
}
/// [`fill_from_chunks`] for one batch of chunks: `chunks` (with their
/// extents `reqs` and, over a backend without the file in memory, the cache
/// hits taken for them) are chunks `first..` of the read, whose stored bytes
/// `raw_bytes` holds.
#[allow(clippy::too_many_arguments)]
fn fill_batch(
chunks: &[ChunkInfo],
reqs: &[ExtentReq],
hits: &[CacheHit],
first: usize,
raw_bytes: &ExtentBytes<'_>,
pipeline: Option<&FilterPipeline>,
placer: &ChunkPlacer,
chunk_total_bytes: usize,
cache: CacheUse<'_>,
out: &OutBuf<'_>,
) -> Result<(), FormatError> {
let rank = placer.rank;
let elem_size = placer.elem_size as u32;
#[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,
@@ -346,15 +376,14 @@ fn fill_batch<S: Storage + ?Sized>(
)));
}
let offsets = &c.offsets[..rank];
let size = c.chunk_size as usize;
#[cfg(feature = "std")]
if let Some(Some(hit)) = hits.get(i) {
raw_bytes.check(i, c.address, size)?;
raw_bytes.check(first + i, &reqs[i])?;
// SAFETY: see above.
unsafe { placer.place(hit, offsets, out) };
return Ok(());
}
let raw = raw_bytes.get(i, c.address, size)?;
let raw = raw_bytes.get(first + i, &reqs[i])?;
let Some(pl) = pipeline else {
// SAFETY: see above.
unsafe { placer.place(raw, offsets, out) };
@@ -1913,9 +1942,9 @@ pub fn read_chunked_data_sweep_in<S: Storage + ?Sized>(
cached
} else {
// Decompress from file
let c_addr = to_usize(chunk_info.address)?;
let size = chunk_info.chunk_size as usize;
let raw_chunk = read_exact_at(file_data, c_addr as u64, size)?;
to_usize(chunk_info.address)?;
let req = chunk_req(chunk_info, pipeline, chunk_total_bytes, true);
let raw_chunk = read_extent(file_data, &req)?;
let raw_chunk = &*raw_chunk;
let dec = if let Some(pl) = pipeline {
decompress_chunk_exact(
@@ -2052,46 +2081,60 @@ pub fn read_chunked_data_indexed_in<S: Storage + ?Sized>(
let chunk_total_bytes = plan.chunk_total_bytes;
// The decoded chunks the cache holds, and the stored bytes of the
// others: fetched in one batch when the file is not in memory.
// others: fetched batch by batch when the file is not in memory.
let hits: Vec<Option<Arc<CacheAlignedBuffer>>> = plan
.mappings
.iter()
.map(|m| cache.get_decompressed_in(addr, &m.coord))
.collect();
let extents: Vec<(u64, usize, bool)> = plan
let reqs: Vec<ExtentReq> = plan
.mappings
.iter()
.zip(&hits)
.map(|(m, hit)| (m.file_offset, m.file_size as usize, hit.is_none()))
.map(|(m, hit)| {
let len = m.file_size as usize;
ExtentReq {
addr: m.file_offset,
len,
fetch: hit.is_none().then(|| {
len.min(crate::filters::stored_chunk_limit(
pipeline,
m.filter_mask,
chunk_total_bytes,
))
}),
}
})
.collect();
let raw_bytes = ExtentBytes::fetch(file_data, &extents)?;
// Decompress chunks (using LRU cache where possible)
let mut chunk_buffers: Vec<Arc<CacheAlignedBuffer>> = Vec::with_capacity(plan.mappings.len());
for (i, (m, hit)) in plan.mappings.iter().zip(hits).enumerate() {
let (coord, file_offset, file_size, filter_mask) =
(&m.coord, &m.file_offset, &m.file_size, &m.filter_mask);
if let Some(cached) = hit {
chunk_buffers.push(cached);
} else {
let raw_chunk = raw_bytes.get(i, *file_offset, *file_size as usize)?;
let mut hits = hits.into_iter();
for_each_extent_batch(file_data, &reqs, |batch, raw_bytes| {
for i in batch {
let m = &plan.mappings[i];
if let Some(cached) = hits.next().flatten() {
chunk_buffers.push(cached);
continue;
}
let raw_chunk = raw_bytes.get(i, &reqs[i])?;
let decompressed = if let Some(pl) = pipeline {
decompress_chunk_exact(
raw_chunk,
pl,
chunk_total_bytes,
elem_size as u32,
*filter_mask,
coord,
m.filter_mask,
&m.coord,
)?
} else {
raw_chunk.to_vec()
};
let aligned = CacheAlignedBuffer::from_vec(decompressed);
let arc = cache.put_decompressed_aligned_in(addr, coord.clone(), aligned);
chunk_buffers.push(arc);
chunk_buffers.push(cache.put_decompressed_aligned_in(addr, m.coord.clone(), aligned));
}
}
Ok(())
})?;
// Assemble using pre-computed layout
let mut output = vec![0u8; plan.output_bytes];
crates/clawhdf5-format/src/filter_registry.rs
+16
View File
@@ -188,6 +188,22 @@ pub fn is_filter_available(id: u16) -> bool {
}
}
/// Whether chunks filtered with `id` may be decoded by a codec the
/// application registered (whose stored sizes this crate cannot bound).
pub(crate) fn may_be_registered(id: u16) -> bool {
if builtin_filter(id).is_some_and(|b| !b.is_shared()) {
return false;
}
#[cfg(feature = "std")]
{
registered(id).is_some()
}
#[cfg(not(feature = "std"))]
{
false
}
}
#[cfg(feature = "std")]
mod custom {
use super::FilterCodec;
crates/clawhdf5-format/src/filters.rs
+38
View File
@@ -63,6 +63,44 @@ fn filter_output_bound(filter_id: u16, input: usize) -> usize {
}
}
/// Most stored bytes a chunk whose decoded size is `chunk_bytes` can need:
/// what a raw-data read fetches of (and decodes from) a chunk, whatever size
/// its index entry claims. A crafted index that points many chunks at huge
/// extents then costs no more than legitimate chunks would.
///
/// An unfiltered chunk (no pipeline, or every filter skipped by
/// `filter_mask`) is the chunk itself: `chunk_bytes`, and bytes past them
/// were never used. A filtered chunk is bounded by each applied filter's
/// worst-case growth in the write direction: shuffle keeps the size,
/// Fletcher32 adds 4 bytes, and any other codec gets `n + n/4 + 4096` — a
/// deliberately generous bound (bzip2 grows 1000 random bytes by 252, more
/// than the decoders' own `n + n/8 + 64` output bound), since a legitimate
/// chunk cut short here would fail to read. A filter handled by a codec the
/// application registered is not ours to bound: such a chunk is limited
/// only by the file.
pub(crate) fn stored_chunk_limit(
pipeline: Option<&FilterPipeline>,
filter_mask: u32,
chunk_bytes: usize,
) -> usize {
let Some(pipeline) = pipeline else {
return chunk_bytes;
};
let mut size = chunk_bytes;
for (i, filter) in pipeline.filters.iter().enumerate() {
if filter_skipped(filter_mask, i) {
continue;
}
size = match filter.filter_id {
FILTER_SHUFFLE => size,
FILTER_FLETCHER32 => size.saturating_add(4),
id if filter_registry::may_be_registered(id) => return usize::MAX,
_ => size.saturating_add(size / 4).saturating_add(4096),
};
}
size
}
/// Whether bit `index` of a chunk's filter mask says filter `index` was
/// skipped when the chunk was written.
fn filter_skipped(filter_mask: u32, index: usize) -> bool {
crates/clawhdf5-format/src/parallel_read.rs
+128 -110
View File
@@ -12,24 +12,21 @@ use crate::error::FormatError;
use crate::filter_pipeline::FilterPipeline;
use crate::filters::decompress_chunk_exact;
use crate::lane_partition::{self, LaneStats, PartitionStats};
use crate::storage::{ExtentBytes, Storage};
use crate::storage::{ExtentReq, Storage, for_each_extent_batch};
/// The stored bytes of every chunk in `chunks`, fetched in one
/// [`Storage::read_ranges`] call when the file is not in memory (each
/// chunk's bounds error is reported when that chunk is decoded, as before).
fn fetch_all<'a, S: Storage + ?Sized>(
file_data: &'a S,
/// The extents of `chunks`' stored bytes (see
/// [`crate::chunked_read::chunk_req`]), fetched batch by batch with
/// [`for_each_extent_batch`] when the file is not in memory (each chunk's
/// bounds error is reported when that chunk is decoded, as before).
fn chunk_reqs(
chunks: &[ChunkInfo],
) -> Result<ExtentBytes<'a>, FormatError> {
let extents: Vec<(u64, usize, bool)> = if file_data.as_contiguous().is_some() {
Vec::new()
} else {
chunks
.iter()
.map(|c| (c.address, c.chunk_size as usize, true))
.collect()
};
ExtentBytes::fetch(file_data, &extents)
pipeline: Option<&FilterPipeline>,
chunk_total_bytes: usize,
) -> Vec<ExtentReq> {
chunks
.iter()
.map(|c| crate::chunked_read::chunk_req(c, pipeline, chunk_total_bytes, true))
.collect()
}
/// Threshold: only use parallel decompression when chunk count exceeds this.
@@ -238,62 +235,75 @@ pub fn decompress_chunks_lane_partitioned_in<S: Storage + ?Sized>(
.unwrap_or(1)
});
let raw_bytes = fetch_all(file_data, chunks)?;
let assignments = lane_partition::partition_chunks(chunks.len(), lanes, seed);
let num_lanes = assignments.len();
// Each lane processes its assigned chunks and returns results + stats.
let lane_results: Result<Vec<(Vec<DecompressedChunk>, LaneStats)>, FormatError> = assignments
.into_par_iter()
.map(|indices| {
let mut results = Vec::with_capacity(indices.len());
let mut stats = LaneStats::default();
for &index in &indices {
let chunk_info = &chunks[index];
let size = chunk_info.chunk_size as usize;
let raw_chunk = raw_bytes.get(index, chunk_info.address, size)?;
let decompressed = decompress_chunk_exact(
raw_chunk,
pipeline,
chunk_total_bytes,
element_size,
chunk_info.filter_mask,
&chunk_info.offsets,
)?;
stats.chunks_processed += 1;
stats.compressed_bytes += size as u64;
stats.decompressed_bytes += decompressed.len() as u64;
results.push(DecompressedChunk {
index,
data: decompressed,
});
}
Ok((results, stats))
})
.collect();
let lane_results = lane_results?;
// Aggregate stats
let mut partition_stats = PartitionStats::new(num_lanes);
let reqs = chunk_reqs(chunks, Some(pipeline), chunk_total_bytes);
let mut ordered: Vec<Vec<u8>> = Vec::with_capacity(chunks.len());
let mut partition_stats = PartitionStats::new(0);
partition_stats.total_chunks = chunks.len();
for (lane_idx, (_, stats)) in lane_results.iter().enumerate() {
partition_stats.per_lane[lane_idx] = stats.clone();
}
// Each batch of fetched chunks is partitioned into lanes and decoded
// before the next batch is fetched (with the file in memory there is
// one batch: all the chunks).
for_each_extent_batch(file_data, &reqs, |batch, raw_bytes| {
let assignments = lane_partition::partition_chunks(batch.len(), lanes, seed);
// Each lane processes its assigned chunks and returns results + stats.
let lane_results: Result<Vec<(Vec<DecompressedChunk>, LaneStats)>, FormatError> =
assignments
.into_par_iter()
.map(|indices| {
let mut results = Vec::with_capacity(indices.len());
let mut stats = LaneStats::default();
// Flatten and sort by original index to restore order
let mut all_chunks: Vec<DecompressedChunk> = lane_results
.into_iter()
.flat_map(|(chunks, _)| chunks)
.collect();
all_chunks.sort_by_key(|dc| dc.index);
for &local in &indices {
let index = batch.start + local;
let chunk_info = &chunks[index];
let size = chunk_info.chunk_size as usize;
let raw_chunk = raw_bytes.get(index, &reqs[index])?;
let ordered = all_chunks.into_iter().map(|dc| dc.data).collect();
let decompressed = decompress_chunk_exact(
raw_chunk,
pipeline,
chunk_total_bytes,
element_size,
chunk_info.filter_mask,
&chunk_info.offsets,
)?;
stats.chunks_processed += 1;
stats.compressed_bytes += size as u64;
stats.decompressed_bytes += decompressed.len() as u64;
results.push(DecompressedChunk {
index,
data: decompressed,
});
}
Ok((results, stats))
})
.collect();
let lane_results = lane_results?;
// Aggregate stats
if partition_stats.per_lane.len() < lane_results.len() {
partition_stats
.per_lane
.resize_with(lane_results.len(), LaneStats::default);
partition_stats.num_lanes = lane_results.len();
}
for (lane, (_, stats)) in partition_stats.per_lane.iter_mut().zip(&lane_results) {
lane.chunks_processed += stats.chunks_processed;
lane.compressed_bytes += stats.compressed_bytes;
lane.decompressed_bytes += stats.decompressed_bytes;
}
// Flatten and sort by original index to restore order
let mut all_chunks: Vec<DecompressedChunk> = lane_results
.into_iter()
.flat_map(|(chunks, _)| chunks)
.collect();
all_chunks.sort_by_key(|dc| dc.index);
ordered.extend(all_chunks.into_iter().map(|dc| dc.data));
Ok(())
})?;
Ok((ordered, partition_stats))
}
@@ -325,33 +335,38 @@ pub fn decompress_chunks_parallel_in<S: Storage + ?Sized>(
) -> Result<Vec<Vec<u8>>, FormatError> {
use rayon::prelude::*;
let raw_bytes = fetch_all(file_data, chunks)?;
let results: Result<Vec<DecompressedChunk>, FormatError> = chunks
.par_iter()
.enumerate()
.map(|(index, chunk_info)| {
let size = chunk_info.chunk_size as usize;
let raw_chunk = raw_bytes.get(index, chunk_info.address, size)?;
let reqs = chunk_reqs(chunks, Some(pipeline), chunk_total_bytes);
let mut ordered: Vec<Vec<u8>> = Vec::with_capacity(chunks.len());
for_each_extent_batch(file_data, &reqs, |batch, raw_bytes| {
let results: Result<Vec<DecompressedChunk>, FormatError> = batch
.clone()
.into_par_iter()
.map(|index| {
let chunk_info = &chunks[index];
let raw_chunk = raw_bytes.get(index, &reqs[index])?;
let decompressed = decompress_chunk_exact(
raw_chunk,
pipeline,
chunk_total_bytes,
element_size,
chunk_info.filter_mask,
&chunk_info.offsets,
)?;
let decompressed = decompress_chunk_exact(
raw_chunk,
pipeline,
chunk_total_bytes,
element_size,
chunk_info.filter_mask,
&chunk_info.offsets,
)?;
Ok(DecompressedChunk {
index,
data: decompressed,
Ok(DecompressedChunk {
index,
data: decompressed,
})
})
})
.collect();
.collect();
let mut result_vec = results?;
result_vec.sort_by_key(|dc| dc.index);
Ok(result_vec.into_iter().map(|dc| dc.data).collect())
let mut result_vec = results?;
result_vec.sort_by_key(|dc| dc.index);
ordered.extend(result_vec.into_iter().map(|dc| dc.data));
Ok(())
})?;
Ok(ordered)
}
/// Decompress chunks sequentially (fallback when parallel is not warranted).
@@ -373,26 +388,29 @@ pub fn decompress_chunks_sequential_in<S: Storage + ?Sized>(
chunk_total_bytes: usize,
element_size: u32,
) -> Result<Vec<Vec<u8>>, FormatError> {
let raw_bytes = fetch_all(file_data, chunks)?;
let reqs = chunk_reqs(chunks, pipeline, chunk_total_bytes);
let mut result = Vec::with_capacity(chunks.len());
for (i, chunk_info) in chunks.iter().enumerate() {
let size = chunk_info.chunk_size as usize;
let raw_chunk = raw_bytes.get(i, chunk_info.address, size)?;
for_each_extent_batch(file_data, &reqs, |batch, raw_bytes| {
for i in batch {
let chunk_info = &chunks[i];
let raw_chunk = raw_bytes.get(i, &reqs[i])?;
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(decompressed);
}
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(decompressed);
}
Ok(())
})?;
Ok(result)
}
crates/clawhdf5-format/src/partial_read.rs
+45 -45
View File
@@ -26,7 +26,7 @@ use crate::error::FormatError;
use crate::filter_pipeline::FilterPipeline;
use crate::filters::{all_filters_skipped, decompress_chunk_exact_with};
use crate::selection::Selection;
use crate::storage::{ExtentBytes, Storage};
use crate::storage::{ExtentReq, Storage, for_each_extent_batch};
/// The smallest axis-aligned box containing every selected element, as
/// `(start, extent)` per dimension. `None` when there is nothing to gain or
@@ -373,50 +373,50 @@ pub fn read_selection_in<S: Storage + ?Sized>(
})
})
.collect();
// Their stored bytes, in one batch when the file is not in memory.
let extents: Vec<(u64, usize, bool)> = if file_data.as_contiguous().is_some() {
Vec::new()
} else {
wanted
.iter()
.map(|c| (c.address, c.chunk_size as usize, true))
.collect()
};
let raw_bytes = ExtentBytes::fetch(file_data, &extents)?;
// Chunks are decoded into this thread's reusable buffers.
crate::chunked_read::with_scratch(|scratch| -> Result<(), FormatError> {
for (i, chunk) in wanted.iter().enumerate() {
let origin = &chunk.offsets[..rank];
usize::try_from(chunk.address)
.map_err(|_| FormatError::Overflow("chunk address exceeds usize".into()))?;
let raw = raw_bytes.get(i, chunk.address, chunk.chunk_size as usize)?;
// 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(())
// Their stored bytes, batch by batch when the file is not in
// memory; each batch's chunks are decoded into this thread's
// reusable buffers before the next batch is fetched.
let reqs: Vec<ExtentReq> = wanted
.iter()
.map(|c| crate::chunked_read::chunk_req(c, pipeline, chunk_bytes, true))
.collect();
for_each_extent_batch(file_data, &reqs, |batch, raw_bytes| {
crate::chunked_read::with_scratch(|scratch| -> Result<(), FormatError> {
for i in batch {
let chunk = wanted[i];
let origin = &chunk.offsets[..rank];
usize::try_from(chunk.address).map_err(|_| {
FormatError::Overflow("chunk address exceeds usize".into())
})?;
let raw = raw_bytes.get(i, &reqs[i])?;
// 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),
crates/clawhdf5-format/src/storage.rs
+120 -42
View File
@@ -337,14 +337,57 @@ pub fn read_upto<S: Storage + ?Sized>(
}
/// Most stored bytes fetched by one [`Storage::read_ranges`] call when a
/// read gathers many extents (a chunked dataset's chunks): a larger read is
/// fetched and decoded batch by batch, so a remote backend never holds more
/// than this much undecoded data per read.
pub(crate) const RAW_BATCH_BYTES: usize = 64 << 20;
/// read gathers many extents (a chunked dataset's chunks, a selection's
/// runs): a larger read is fetched and decoded batch by batch, so a backend
/// without the file in memory never holds more than this much undecoded
/// data per read (or one extent, when a single one is larger — and every
/// chunk's extent is bounded by what the chunk can need, see
/// [`crate::filters::stored_chunk_limit`]).
pub const RAW_BATCH_BYTES: usize = 64 << 20;
/// The stored bytes of a list of extents (chunks, contiguous runs), fetched
/// together: [`Storage::read_ranges`] is called once for all of them, so a
/// remote backend can coalesce and parallelise the requests.
/// One extent of a raw-data read: `len` bytes stored at `addr`, whose
/// bounds are checked against the file, of which the first `fetch` bytes
/// are read (`None`: only checked, not read — its bytes are not needed).
///
/// `fetch` below `len` bounds what a crafted size field can make a read
/// fetch: a chunk never needs more of its stored bytes than its decoded
/// size allows, however large its index entry says it is.
#[derive(Debug, Clone, Copy)]
pub(crate) struct ExtentReq {
pub addr: u64,
pub len: usize,
pub fetch: Option<usize>,
}
impl ExtentReq {
/// How many bytes are read for this extent.
#[inline]
fn fetch_len(&self) -> usize {
self.fetch.map_or(0, |f| f.min(self.len))
}
}
/// One extent's bytes on their own (see [`ExtentReq`]): the whole extent's
/// bounds checked as [`read_exact_at`] checks them, and its first
/// `req.fetch` bytes read (none when `fetch` is `None`).
pub(crate) fn read_extent<'a, S: Storage + ?Sized>(
file: &'a S,
req: &ExtentReq,
) -> Result<Cow<'a, [u8]>, FormatError> {
let start = usize::try_from(req.addr).unwrap_or(usize::MAX);
match start.checked_add(req.len) {
Some(end) if end <= len_usize(file) => read_exact_at(file, req.addr, req.fetch_len()),
_ => Err(FormatError::UnexpectedEof {
expected: start.saturating_add(req.len),
available: len_usize(file),
}),
}
}
/// The stored bytes of one batch of extents (chunks, contiguous runs),
/// fetched together: [`Storage::read_ranges`] is called once for the batch,
/// so a remote backend can coalesce and parallelise the requests. See
/// [`for_each_extent_batch`], which is how every raw-data read gets them.
///
/// With the whole file in memory nothing is fetched: [`Self::get`] slices
/// it, as the slice readers did. Either way an extent that does not lie in
@@ -356,8 +399,12 @@ pub(crate) const RAW_BATCH_BYTES: usize = 64 << 20;
pub(crate) enum ExtentBytes<'a> {
/// The whole file.
Contiguous(&'a [u8]),
/// Each extent's bytes, or its bounds error.
Fetched(Vec<Extent<'a>>),
/// Each extent's bytes, or its bounds error; the first is extent
/// `base` of the read.
Fetched {
base: usize,
extents: Vec<Extent<'a>>,
},
}
/// One extent of [`ExtentBytes::Fetched`].
@@ -371,33 +418,35 @@ pub(crate) enum Extent<'a> {
}
impl<'a> ExtentBytes<'a> {
/// Fetch `extents` (`(address, length, wanted)`): the bytes of those
/// `wanted`, and the bounds check of all of them.
pub(crate) fn fetch<S: Storage + ?Sized>(
/// Fetch `reqs`, extents `base..base + reqs.len()` of the read: the
/// bytes of those wanted, and the bounds check of all of them.
fn fetch<S: Storage + ?Sized>(
file: &'a S,
extents: &[(u64, usize, bool)],
reqs: &[ExtentReq],
base: usize,
) -> Result<Self, FormatError> {
if let Some(all) = file.as_contiguous() {
return Ok(ExtentBytes::Contiguous(all));
}
let file_len = len_usize(file);
let mut ranges = Vec::new();
let mut out = Vec::with_capacity(extents.len());
let mut out = Vec::with_capacity(reqs.len());
// Positions in `out` of the extents being read, in `ranges` order.
let mut slots = Vec::new();
for &(addr, len, wanted) in extents {
let checked =
crate::addr::to_usize(addr).and_then(|start| match start.checked_add(len) {
for req in reqs {
let checked = crate::addr::to_usize(req.addr).and_then(|start| {
match start.checked_add(req.len) {
Some(end) if end <= file_len => Ok(()),
_ => Err(FormatError::UnexpectedEof {
expected: start.saturating_add(len),
expected: start.saturating_add(req.len),
available: file_len,
}),
});
}
});
match checked {
Ok(()) if wanted => {
Ok(()) if req.fetch.is_some() => {
slots.push(out.len());
ranges.push(addr..addr + len as u64);
ranges.push(req.addr..req.addr + req.fetch_len() as u64);
out.push(Extent::NotFetched);
}
Ok(()) => out.push(Extent::NotFetched),
@@ -418,41 +467,46 @@ impl<'a> ExtentBytes<'a> {
out[slot] = Extent::Bytes(bytes);
}
}
Ok(ExtentBytes::Fetched(out))
Ok(ExtentBytes::Fetched { base, extents: out })
}
/// Whether extent `i` (at `addr`, `len` bytes, as passed to
/// [`Self::fetch`]) lies in the file: its bounds error if not.
pub(crate) fn check(&self, i: usize, addr: u64, len: usize) -> Result<(), FormatError> {
/// Whether extent `i` of the read (`req`) lies in the file: its bounds
/// error if not.
pub(crate) fn check(&self, i: usize, req: &ExtentReq) -> Result<(), FormatError> {
match self {
ExtentBytes::Contiguous(_) => self.get(i, addr, len).map(|_| ()),
ExtentBytes::Fetched(v) => match v.get(i) {
Some(Extent::Err(e)) => Err(e.clone()),
Some(_) => Ok(()),
None => Err(not_fetched()),
},
ExtentBytes::Contiguous(_) => self.get(i, req).map(|_| ()),
ExtentBytes::Fetched { base, extents } => {
match i.checked_sub(*base).and_then(|j| extents.get(j)) {
Some(Extent::Err(e)) => Err(e.clone()),
Some(_) => Ok(()),
None => Err(not_fetched()),
}
}
}
}
/// Extent `i`'s bytes (at `addr`, `len` bytes, as passed to
/// [`Self::fetch`]).
pub(crate) fn get(&self, i: usize, addr: u64, len: usize) -> Result<&[u8], FormatError> {
/// Extent `i` of the read (`req`): its first `req.fetch` bytes, the
/// same whether the file is in memory or not.
pub(crate) fn get(&self, i: usize, req: &ExtentReq) -> Result<&[u8], FormatError> {
match self {
ExtentBytes::Contiguous(all) => {
let start = crate::addr::to_usize(addr)?;
let start = crate::addr::to_usize(req.addr)?;
start
.checked_add(len)
.checked_add(req.len)
.and_then(|end| all.get(start..end))
.map(|b| &b[..req.fetch_len()])
.ok_or(FormatError::UnexpectedEof {
expected: start.saturating_add(len),
expected: start.saturating_add(req.len),
available: <[u8]>::len(all),
})
}
ExtentBytes::Fetched(v) => match v.get(i) {
Some(Extent::Bytes(b)) => Ok(b),
Some(Extent::Err(e)) => Err(e.clone()),
_ => Err(not_fetched()),
},
ExtentBytes::Fetched { base, extents } => {
match i.checked_sub(*base).and_then(|j| extents.get(j)) {
Some(Extent::Bytes(b)) => Ok(b),
Some(Extent::Err(e)) => Err(e.clone()),
_ => Err(not_fetched()),
}
}
}
}
}
@@ -462,6 +516,30 @@ fn not_fetched() -> FormatError {
FormatError::Storage("an extent that was not fetched was asked for".into())
}
/// The one way raw-data reads fetch stored bytes: `reqs` are split into
/// consecutive batches of at most [`RAW_BATCH_BYTES`] of fetched bytes (at
/// least one extent each — and no extent fetches more than its
/// [`ExtentReq::fetch`]), and for each batch in turn its bytes are fetched
/// with one [`Storage::read_ranges`] call and `f(batch, &bytes)` is called,
/// with `bytes` indexed by the extent's position in `reqs`. A batch's bytes
/// are dropped before the next batch is fetched, and an error from `f`
/// stops the read before anything more is fetched.
///
/// With the whole file in memory there is nothing to fetch: one call, over
/// all of `reqs`, that slices the file.
pub(crate) fn for_each_extent_batch<'a, S: Storage + ?Sized>(
file: &'a S,
reqs: &[ExtentReq],
mut f: impl FnMut(Range<usize>, &ExtentBytes<'a>) -> Result<(), FormatError>,
) -> Result<(), FormatError> {
let contiguous = file.as_contiguous().is_some();
for batch in raw_batches(reqs.len(), contiguous, |i| reqs[i].fetch_len()) {
let bytes = ExtentBytes::fetch(file, &reqs[batch.clone()], batch.start)?;
f(batch, &bytes)?;
}
Ok(())
}
/// Split `n` extents, whose sizes `size(i)` gives, into consecutive batches
/// of at most [`RAW_BATCH_BYTES`] (at least one extent each): the ranges of
/// `0..n` to fetch together. With the whole file in memory (`contiguous`)
crates/clawhdf5-format/tests/raw_fetch_bounds.rs
+332
View File
@@ -0,0 +1,332 @@
//! What a raw-data read fetches from a [`Storage`] without the file in
//! memory is bounded: by batch (at most `RAW_BATCH_BYTES` per
//! `read_ranges` call) and by chunk (never more of a chunk's stored bytes
//! than its decoded size can need), on every path that reads chunks — full,
//! cached, indexed, sweep, selection and the `parallel_read` decoders.
//!
//! A crafted chunk index can point every chunk at one huge extent. Slicing
//! an in-memory file costs nothing there, but a backend that fetches would
//! hold `chunks x extent` bytes before the first chunk failed to decode.
use std::borrow::Cow;
use std::ops::Range;
use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
use clawhdf5_format::chunk_cache::ChunkCache;
use clawhdf5_format::chunked_read::{
ChunkInfo, SweepContext, list_chunks, read_chunked_data_sweep_in,
};
use clawhdf5_format::data_layout::DataLayout;
use clawhdf5_format::data_read::{
read_raw_data_cached_in, read_raw_data_full_in, read_raw_data_indexed_in,
read_raw_data_selection_in,
};
use clawhdf5_format::dataspace::Dataspace;
use clawhdf5_format::datatype::Datatype;
use clawhdf5_format::error::FormatError;
use clawhdf5_format::filter_pipeline::FilterPipeline;
use clawhdf5_format::group_v2;
use clawhdf5_format::message_type::MessageType;
use clawhdf5_format::object_header::ObjectHeader;
use clawhdf5_format::selection::Selection;
use clawhdf5_format::storage::{RAW_BATCH_BYTES, Storage};
use clawhdf5_format::superblock::Superblock;
/// A read_at-only storage that records the most bytes one call fetched
/// (a `read_ranges` call counts all its ranges together) and the total.
struct PeakStorage {
data: Vec<u8>,
peak: AtomicU64,
total: AtomicU64,
}
impl PeakStorage {
fn new(data: Vec<u8>) -> Self {
PeakStorage {
data,
peak: AtomicU64::new(0),
total: AtomicU64::new(0),
}
}
fn reset(&self) {
self.peak.store(0, Relaxed);
self.total.store(0, Relaxed);
}
fn served(&self, offset: u64, len: usize) -> Vec<u8> {
self.data
.as_slice()
.read_at(offset, len)
.unwrap()
.into_owned()
}
}
impl Storage for PeakStorage {
fn read_at(&self, offset: u64, len: usize) -> Result<Cow<'_, [u8]>, FormatError> {
let got = self.served(offset, len);
self.peak.fetch_max(got.len() as u64, Relaxed);
self.total.fetch_add(got.len() as u64, Relaxed);
Ok(Cow::Owned(got))
}
fn len(&self) -> u64 {
self.data.len() as u64
}
fn read_ranges(&self, ranges: &[Range<u64>]) -> Result<Vec<Cow<'_, [u8]>>, FormatError> {
let got: Vec<Vec<u8>> = ranges
.iter()
.map(|r| self.served(r.start, (r.end - r.start) as usize))
.collect();
let bytes: u64 = got.iter().map(|g| g.len() as u64).sum();
self.peak.fetch_max(bytes, Relaxed);
self.total.fetch_add(bytes, Relaxed);
Ok(got.into_iter().map(Cow::Owned).collect())
}
}
struct Chunked {
layout: DataLayout,
dataspace: Dataspace,
datatype: Datatype,
pipeline: Option<FilterPipeline>,
os: u8,
ls: u8,
}
/// The one chunked dataset of fixture `name`.
fn chunked(bytes: &[u8]) -> Chunked {
let sb = Superblock::parse(bytes, 0).unwrap();
let (os, ls) = (sb.offset_size, sb.length_size);
for child in group_v2::resolve_group_children(bytes, &sb, sb.root_group_address).unwrap() {
let header =
ObjectHeader::parse(bytes, child.object_header_address as usize, os, ls).unwrap();
let msg = |t: MessageType| {
header
.messages
.iter()
.find(|m| m.msg_type == t)
.map(|m| m.data.clone())
};
let Some(dl) = msg(MessageType::DataLayout) else {
continue;
};
let layout = DataLayout::parse(&dl, os, ls).unwrap();
if !matches!(layout, DataLayout::Chunked { .. }) {
continue;
}
return Chunked {
layout,
datatype: Datatype::parse(&msg(MessageType::Datatype).unwrap())
.unwrap()
.0,
dataspace: Dataspace::parse(&msg(MessageType::Dataspace).unwrap(), ls).unwrap(),
pipeline: msg(MessageType::FilterPipeline).map(|p| FilterPipeline::parse(&p).unwrap()),
os,
ls,
};
}
panic!("no chunked dataset");
}
/// Claimed stored size of every chunk in the crafted file.
const HUGE: u32 = 20 << 20;
/// `chunked_large.h5` (1000 `i32` in ten gzip chunks, a v1 B-tree index)
/// with `HUGE` bytes of padding appended and every chunk's index entry
/// rewritten to claim `HUGE` stored bytes at the padding: ten chunks, 200
/// MiB of extents, in a 20 MiB file.
fn crafted() -> (Vec<u8>, Chunked, Vec<ChunkInfo>) {
let dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures");
let mut bytes = std::fs::read(dir.join("chunked_large.h5")).unwrap();
let ds = chunked(&bytes);
let es = ds.datatype.type_size() as usize;
let (chunks, _) = list_chunks(&bytes, &ds.layout, &ds.dataspace, es, ds.os, ds.ls).unwrap();
assert_eq!(chunks.len(), 10);
let blob = bytes.len() as u64;
for c in &chunks {
// v1 B-tree key (size, filter mask, offsets + 0) then the child
// address.
let mut pat = Vec::new();
pat.extend_from_slice(&c.chunk_size.to_le_bytes());
pat.extend_from_slice(&c.filter_mask.to_le_bytes());
// The key holds one offset per dimension plus the element offset
// (0); `offsets` may or may not list that last one.
for d in 0..=ds.dataspace.dimensions.len() {
pat.extend_from_slice(&c.offsets.get(d).copied().unwrap_or(0).to_le_bytes());
}
pat.extend_from_slice(&c.address.to_le_bytes());
let at = bytes
.windows(pat.len())
.position(|w| w == pat.as_slice())
.expect("chunk key");
bytes[at..at + 4].copy_from_slice(&HUGE.to_le_bytes());
let a = at + pat.len() - 8;
bytes[a..a + 8].copy_from_slice(&blob.to_le_bytes());
}
bytes.resize(bytes.len() + HUGE as usize, 0x5a);
let ds = chunked(&bytes);
let (chunks, _) = list_chunks(&bytes, &ds.layout, &ds.dataspace, es, ds.os, ds.ls).unwrap();
assert!(
chunks
.iter()
.all(|c| c.chunk_size == HUGE && c.address == blob)
);
(bytes, ds, chunks)
}
/// Most a crafted chunk of the fixture may fetch: its decoded size (400
/// bytes) grown by one codec, generously.
const CHUNK_LIMIT: u64 = 400 + 100 + 4096;
#[test]
fn crafted_chunk_index_cannot_amplify_fetches() {
let (bytes, ds, chunks) = crafted();
let st = PeakStorage::new(bytes.clone());
let pl = ds.pipeline.as_ref();
let (dl, sp, dt, os, ls) = (&ds.layout, &ds.dataspace, &ds.datatype, ds.os, ds.ls);
let check =
|what: &str, got: Result<Vec<u8>, FormatError>, want: Result<Vec<u8>, FormatError>| {
// Same outcome as slicing the whole file.
assert_eq!(got, want, "{what}");
let (peak, total) = (st.peak.load(Relaxed), st.total.load(Relaxed));
assert!(
peak <= RAW_BATCH_BYTES as u64,
"{what}: one fetch of {peak} bytes"
);
// Every chunk's fetch is bounded by what it can need, whatever its
// index entry claims (plus the index and header reads).
assert!(
total <= chunks.len() as u64 * CHUNK_LIMIT + 64 * 1024,
"{what}: fetched {total} bytes"
);
st.reset();
};
let slice: &[u8] = &bytes;
let sel = Selection::Hyperslab {
start: vec![100],
stride: vec![1],
count: vec![400],
block: vec![1],
};
check(
"selection",
read_raw_data_selection_in(&st, dl, sp, dt, pl, os, ls, &sel),
read_raw_data_selection_in(slice, dl, sp, dt, pl, os, ls, &sel),
);
check(
"full",
read_raw_data_full_in(&st, dl, sp, dt, pl, os, ls),
read_raw_data_full_in(slice, dl, sp, dt, pl, os, ls),
);
check(
"cached",
read_raw_data_cached_in(&st, dl, sp, dt, pl, os, ls, &ChunkCache::new()),
read_raw_data_cached_in(slice, dl, sp, dt, pl, os, ls, &ChunkCache::new()),
);
check(
"indexed",
read_raw_data_indexed_in(&st, dl, sp, dt, pl, os, ls, &ChunkCache::new()),
read_raw_data_indexed_in(slice, dl, sp, dt, pl, os, ls, &ChunkCache::new()),
);
check(
"sweep",
read_chunked_data_sweep_in(
&st,
dl,
sp,
dt,
pl,
os,
ls,
&ChunkCache::new(),
&mut SweepContext::new(4, 2),
),
read_chunked_data_sweep_in(
slice,
dl,
sp,
dt,
pl,
os,
ls,
&ChunkCache::new(),
&mut SweepContext::new(4, 2),
),
);
#[cfg(feature = "parallel")]
{
use clawhdf5_format::parallel_read::{
decompress_chunks_lane_partitioned_in, decompress_chunks_parallel_in,
decompress_chunks_sequential_in,
};
let pl = pl.unwrap();
let flat = |r: Result<Vec<Vec<u8>>, FormatError>| r.map(|v| v.concat());
check(
"parallel",
flat(decompress_chunks_parallel_in(&st, &chunks, pl, 400, 4)),
flat(decompress_chunks_parallel_in(slice, &chunks, pl, 400, 4)),
);
check(
"sequential",
flat(decompress_chunks_sequential_in(
&st,
&chunks,
Some(pl),
400,
4,
)),
flat(decompress_chunks_sequential_in(
slice,
&chunks,
Some(pl),
400,
4,
)),
);
check(
"lane partitioned",
flat(
decompress_chunks_lane_partitioned_in(&st, &chunks, pl, 400, 4, 7, Some(3))
.map(|(v, _)| v),
),
flat(
decompress_chunks_lane_partitioned_in(slice, &chunks, pl, 400, 4, 7, Some(3))
.map(|(v, _)| v),
),
);
}
}
/// Legitimately large chunks (unfiltered, 4 MiB each, 160 MiB in all) are
/// fetched batch by batch: no call holds more than the batch budget, and
/// the data is right.
#[cfg(feature = "parallel")]
#[test]
fn large_reads_are_fetched_in_batches() {
use clawhdf5_format::parallel_read::decompress_chunks_sequential_in;
const CHUNK: usize = 4 << 20;
let data: Vec<u8> = (0..2 * CHUNK).map(|i| (i % 251) as u8).collect();
let chunks: Vec<ChunkInfo> = (0..40u64)
.map(|i| ChunkInfo {
chunk_size: CHUNK as u32,
filter_mask: 0,
offsets: vec![i * CHUNK as u64],
address: (i % 2) * CHUNK as u64,
})
.collect();
let st = PeakStorage::new(data.clone());
let got = decompress_chunks_sequential_in(&st, &chunks, None, CHUNK, 1).unwrap();
assert_eq!(got.len(), 40);
for (i, c) in got.iter().enumerate() {
let at = (i % 2) * CHUNK;
assert!(c == &data[at..at + CHUNK], "chunk {i}");
}
let peak = st.peak.load(Relaxed);
assert!(peak <= RAW_BATCH_BYTES as u64, "one fetch of {peak} bytes");
assert_eq!(st.total.load(Relaxed), 40 * CHUNK as u64);
}