format: bound and batch every chunk fetch over Storage
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
+12
-3
@@ -47,9 +47,18 @@
|
||||
`read_vl_strings_in`, `read_vl_bytes_in`),
|
||||
`AttributeMessage::read_vl_strings_in`, `provenance::verify_dataset_in`.
|
||||
- A chunked read first lists the chunks it needs, then fetches all their
|
||||
stored bytes with **one `read_ranges` call** (per 64 MiB of stored
|
||||
data), so a remote backend can coalesce and parallelise them, then
|
||||
decodes as before (in parallel with the `parallel` feature). Chunks the
|
||||
stored bytes with **one `read_ranges` call** per batch of at most 64 MiB
|
||||
(`storage::RAW_BATCH_BYTES`), so a remote backend can coalesce and
|
||||
parallelise them, then decodes each batch as before (in parallel with
|
||||
the `parallel` feature) before fetching the next. Every path that reads
|
||||
chunks — full, cached, indexed, sweep, selection and the
|
||||
`parallel_read` decoders — goes through the same batching, and no chunk
|
||||
fetches more of its stored bytes than its decoded size can need (the
|
||||
chunk size if unfiltered; else each applied filter's worst-case growth,
|
||||
generously: `n + n/4 + 4096` per codec, unbounded only for a codec the
|
||||
application registered). A crafted chunk index that points every chunk
|
||||
at one huge extent therefore costs a bounded fetch, not
|
||||
`chunks x extent` bytes (`tests/raw_fetch_bounds.rs`). Chunks the
|
||||
file's chunk cache already holds are not fetched. A selection fetches
|
||||
only the chunks its bounding box overlaps; a contiguous selection only
|
||||
its runs (adjacent ones merged). A global-heap collection is read once
|
||||
|
||||
@@ -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
|
||||
#[cfg(not(feature = "std"))]
|
||||
let hits: Vec<CacheHit> = Vec::new();
|
||||
let reqs: Vec<ExtentReq> = 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)
|
||||
chunk_req(c, pipeline, chunk_total_bytes, wanted)
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
let raw_bytes = ExtentBytes::fetch(file_data, &extents)?;
|
||||
.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 {
|
||||
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);
|
||||
} else {
|
||||
let raw_chunk = raw_bytes.get(i, *file_offset, *file_size as usize)?;
|
||||
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];
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
pipeline: Option<&FilterPipeline>,
|
||||
chunk_total_bytes: usize,
|
||||
) -> Vec<ExtentReq> {
|
||||
chunks
|
||||
.iter()
|
||||
.map(|c| (c.address, c.chunk_size as usize, true))
|
||||
.map(|c| crate::chunked_read::chunk_req(c, pipeline, chunk_total_bytes, true))
|
||||
.collect()
|
||||
};
|
||||
ExtentBytes::fetch(file_data, &extents)
|
||||
}
|
||||
|
||||
/// Threshold: only use parallel decompression when chunk count exceeds this.
|
||||
@@ -238,21 +235,28 @@ 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();
|
||||
|
||||
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();
|
||||
// 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
|
||||
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 {
|
||||
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, chunk_info.address, size)?;
|
||||
let raw_chunk = raw_bytes.get(index, &reqs[index])?;
|
||||
|
||||
let decompressed = decompress_chunk_exact(
|
||||
raw_chunk,
|
||||
@@ -276,14 +280,19 @@ pub fn decompress_chunks_lane_partitioned_in<S: Storage + ?Sized>(
|
||||
Ok((results, stats))
|
||||
})
|
||||
.collect();
|
||||
|
||||
let lane_results = lane_results?;
|
||||
|
||||
// Aggregate stats
|
||||
let mut partition_stats = PartitionStats::new(num_lanes);
|
||||
partition_stats.total_chunks = chunks.len();
|
||||
for (lane_idx, (_, stats)) in lane_results.iter().enumerate() {
|
||||
partition_stats.per_lane[lane_idx] = stats.clone();
|
||||
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
|
||||
@@ -292,8 +301,9 @@ pub fn decompress_chunks_lane_partitioned_in<S: Storage + ?Sized>(
|
||||
.flat_map(|(chunks, _)| chunks)
|
||||
.collect();
|
||||
all_chunks.sort_by_key(|dc| dc.index);
|
||||
|
||||
let ordered = all_chunks.into_iter().map(|dc| dc.data).collect();
|
||||
ordered.extend(all_chunks.into_iter().map(|dc| dc.data));
|
||||
Ok(())
|
||||
})?;
|
||||
Ok((ordered, partition_stats))
|
||||
}
|
||||
|
||||
@@ -325,13 +335,15 @@ 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,
|
||||
@@ -351,7 +363,10 @@ pub fn decompress_chunks_parallel_in<S: Storage + ?Sized>(
|
||||
|
||||
let mut result_vec = results?;
|
||||
result_vec.sort_by_key(|dc| dc.index);
|
||||
Ok(result_vec.into_iter().map(|dc| dc.data).collect())
|
||||
ordered.extend(result_vec.into_iter().map(|dc| dc.data));
|
||||
Ok(())
|
||||
})?;
|
||||
Ok(ordered)
|
||||
}
|
||||
|
||||
/// Decompress chunks sequentially (fallback when parallel is not warranted).
|
||||
@@ -373,11 +388,12 @@ 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(
|
||||
@@ -393,6 +409,8 @@ pub fn decompress_chunks_sequential_in<S: Storage + ?Sized>(
|
||||
};
|
||||
result.push(decompressed);
|
||||
}
|
||||
Ok(())
|
||||
})?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
|
||||
@@ -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,25 +373,24 @@ 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
|
||||
// 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| (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.
|
||||
.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, chunk) in wanted.iter().enumerate() {
|
||||
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, 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.
|
||||
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(
|
||||
@@ -417,6 +416,7 @@ pub fn read_selection_in<S: Storage + ?Sized>(
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
})?;
|
||||
}
|
||||
_ => return Ok(None),
|
||||
|
||||
@@ -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) {
|
||||
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) {
|
||||
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`)
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
@@ -349,9 +349,11 @@ impl File {
|
||||
/// selections, variable-length data, virtual datasets).
|
||||
///
|
||||
/// Every read goes through the storage's [`Storage::read_at`] and
|
||||
/// [`Storage::read_ranges`] (a chunked read fetches all the chunks it
|
||||
/// needs with one `read_ranges` call per 64 MiB), so nothing is read that
|
||||
/// the operation does not need. A storage that has the whole file in
|
||||
/// [`Storage::read_ranges`] (a chunked read fetches the chunks it needs
|
||||
/// with one `read_ranges` call per batch of at most 64 MiB, decoding
|
||||
/// each batch before the next, and never more of a chunk than its
|
||||
/// decoded size can need), so nothing is read that the operation does
|
||||
/// not need. A storage that has the whole file in
|
||||
/// memory ([`Storage::as_contiguous`]) is read as [`File::from_bytes`]
|
||||
/// reads its buffer. The storage holds the whole file: a user block is
|
||||
/// found and skipped, and bytes past the end of file the superblock
|
||||
|
||||
@@ -435,8 +435,12 @@ fast path within benchmark noise.
|
||||
(checked: identical conformance results; the bench gate below still
|
||||
has to be run on an idle machine).
|
||||
- Chunked reads fetch in batches of at most 64 MiB of stored bytes, one
|
||||
`read_ranges` call each, so a remote read never holds more than that
|
||||
undecoded; chunks already in the chunk cache are not fetched.
|
||||
`read_ranges` call each, decoding each batch before fetching the next,
|
||||
on every chunk-reading path (one helper,
|
||||
`storage::for_each_extent_batch`); and no chunk fetches more than its
|
||||
decoded size can need (`filters::stored_chunk_limit`), so a remote read
|
||||
holds at most one batch undecoded even over a crafted chunk index;
|
||||
chunks already in the chunk cache are not fetched.
|
||||
- The typed readers' zero-copy fast path became "read the contiguous
|
||||
bytes once": over a range storage a contiguous `read_f64` is one read,
|
||||
and a native contiguous selection reads only its runs.
|
||||
|
||||
Reference in New Issue
Block a user