osobhandClaude Opus 5.5
7d629f49e3
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]>
433 lines
16 KiB
Rust
433 lines
16 KiB
Rust
//! Selection reads that cost what the selection costs, not what the dataset
|
|
//! costs.
|
|
//!
|
|
//! [`crate::data_read::read_raw_data_selection`] used to decode the *entire*
|
|
//! dataset and then pick elements out of it, so reading a 64x64 window of a
|
|
//! large dataset took about as long as reading all of it. A contiguous
|
|
//! dataset's selection is now copied straight out of the file, one `memcpy`
|
|
//! per contiguous run of selected elements (`crate::gather`). For chunked
|
|
//! data the selection's bounding box is materialised — only the chunks that
|
|
//! overlap it — and the extractor runs over that small buffer with the
|
|
//! selection translated to the box's origin. Extraction semantics are
|
|
//! therefore exactly the full-read ones.
|
|
|
|
#[cfg(not(feature = "std"))]
|
|
use alloc::string as alloc_or_std;
|
|
#[cfg(not(feature = "std"))]
|
|
use alloc::{format, vec, vec::Vec};
|
|
#[cfg(feature = "std")]
|
|
use std::string as alloc_or_std;
|
|
|
|
use crate::chunked_read::{alloc_output, checked_byte_len, list_chunks_for_read_in};
|
|
use crate::data_layout::DataLayout;
|
|
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_with};
|
|
use crate::selection::Selection;
|
|
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
|
|
/// the selection is not valid for `dims` (the caller's full path then reports
|
|
/// the error exactly as before).
|
|
fn bounding_box(selection: &Selection, dims: &[u64]) -> Option<(Vec<u64>, Vec<u64>)> {
|
|
match selection {
|
|
Selection::Hyperslab {
|
|
start,
|
|
stride,
|
|
count,
|
|
block,
|
|
} => {
|
|
let rank = dims.len();
|
|
if [start.len(), stride.len(), count.len(), block.len()] != [rank; 4] {
|
|
return None;
|
|
}
|
|
let mut extent = Vec::with_capacity(rank);
|
|
for d in 0..rank {
|
|
if count[d] == 0 || block[d] == 0 {
|
|
return None;
|
|
}
|
|
// Last selected index + 1, relative to start.
|
|
let span = (count[d] - 1)
|
|
.checked_mul(stride[d])?
|
|
.checked_add(block[d])?;
|
|
if start[d].checked_add(span)? > dims[d] {
|
|
return None;
|
|
}
|
|
extent.push(span);
|
|
}
|
|
Some((start.clone(), extent))
|
|
}
|
|
Selection::Points(points) => {
|
|
let rank = dims.len();
|
|
let first = points.first()?;
|
|
if first.len() != rank {
|
|
return None;
|
|
}
|
|
let (mut lo, mut hi) = (first.clone(), first.clone());
|
|
for p in points {
|
|
if p.len() != rank {
|
|
return None;
|
|
}
|
|
for d in 0..rank {
|
|
if p[d] >= dims[d] {
|
|
return None;
|
|
}
|
|
lo[d] = lo[d].min(p[d]);
|
|
hi[d] = hi[d].max(p[d]);
|
|
}
|
|
}
|
|
let extent = lo.iter().zip(&hi).map(|(l, h)| h - l + 1).collect();
|
|
Some((lo, extent))
|
|
}
|
|
Selection::All | Selection::None => None,
|
|
}
|
|
}
|
|
|
|
/// Check that `selection` addresses only elements that exist in a dataset of
|
|
/// shape `dims`. Without this an out-of-range selection read *something*: a
|
|
/// hyperslab past the edge came back padded with zeros, and a point whose
|
|
/// column was out of range wrapped into the next row.
|
|
pub fn validate(selection: &Selection, dims: &[u64]) -> Result<(), FormatError> {
|
|
let rank = dims.len();
|
|
let bad = |msg: alloc_or_std::String| Err(FormatError::SelectionOutOfBounds(msg));
|
|
match selection {
|
|
Selection::All | Selection::None => Ok(()),
|
|
Selection::Hyperslab {
|
|
start,
|
|
stride,
|
|
count,
|
|
block,
|
|
} => {
|
|
if [start.len(), stride.len(), count.len(), block.len()] != [rank; 4] {
|
|
return bad(format!("hyperslab rank does not match dataset rank {rank}"));
|
|
}
|
|
for d in 0..rank {
|
|
if count[d] == 0 || block[d] == 0 {
|
|
continue; // selects nothing along this dimension
|
|
}
|
|
let end = (count[d] - 1)
|
|
.checked_mul(stride[d])
|
|
.and_then(|v| v.checked_add(block[d]))
|
|
.and_then(|v| v.checked_add(start[d]));
|
|
if !end.is_some_and(|end| end <= dims[d]) {
|
|
return bad(format!(
|
|
"dimension {d}: start {} stride {} count {} block {} exceeds extent {}",
|
|
start[d], stride[d], count[d], block[d], dims[d]
|
|
));
|
|
}
|
|
if block[d] > stride[d] && count[d] > 1 {
|
|
return bad(format!(
|
|
"dimension {d}: block {} larger than stride {} (overlapping blocks)",
|
|
block[d], stride[d]
|
|
));
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
Selection::Points(points) => {
|
|
for p in points {
|
|
if p.len() != rank {
|
|
return bad(format!("point {p:?} does not match dataset rank {rank}"));
|
|
}
|
|
if let Some(d) = (0..rank).find(|&d| p[d] >= dims[d]) {
|
|
return bad(format!(
|
|
"point {p:?}: coordinate {} exceeds extent {} of dimension {d}",
|
|
p[d], dims[d]
|
|
));
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|
|
|
|
/// The same selection expressed relative to `origin`.
|
|
fn translate(selection: &Selection, origin: &[u64]) -> Selection {
|
|
match selection {
|
|
Selection::Hyperslab {
|
|
start,
|
|
stride,
|
|
count,
|
|
block,
|
|
} => Selection::Hyperslab {
|
|
start: start.iter().zip(origin).map(|(s, o)| s - o).collect(),
|
|
stride: stride.clone(),
|
|
count: count.clone(),
|
|
block: block.clone(),
|
|
},
|
|
Selection::Points(points) => Selection::Points(
|
|
points
|
|
.iter()
|
|
.map(|p| p.iter().zip(origin).map(|(c, o)| c - o).collect())
|
|
.collect(),
|
|
),
|
|
other => other.clone(),
|
|
}
|
|
}
|
|
|
|
/// Copy the part of a source region that overlaps the box into `out` (which
|
|
/// is the box, row-major).
|
|
///
|
|
/// The source region starts at `src_origin` in dataset coordinates, has shape
|
|
/// `src_shape`, and its elements are in `src` row-major. One `memcpy` per
|
|
/// overlapping row of the last dimension.
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn copy_overlap(
|
|
src: &[u8],
|
|
src_origin: &[u64],
|
|
src_shape: &[u64],
|
|
out: &mut [u8],
|
|
box_start: &[u64],
|
|
box_extent: &[u64],
|
|
elem_size: usize,
|
|
) {
|
|
let rank = box_start.len();
|
|
// Overlap in dataset coordinates.
|
|
let mut lo = vec![0u64; rank];
|
|
let mut hi = vec![0u64; rank];
|
|
for d in 0..rank {
|
|
lo[d] = src_origin[d].max(box_start[d]);
|
|
hi[d] = (src_origin[d] + src_shape[d]).min(box_start[d] + box_extent[d]);
|
|
if lo[d] >= hi[d] {
|
|
return;
|
|
}
|
|
}
|
|
let strides = |shape: &[u64]| {
|
|
let mut s = vec![1u64; rank];
|
|
for d in (0..rank.saturating_sub(1)).rev() {
|
|
s[d] = s[d + 1] * shape[d + 1];
|
|
}
|
|
s
|
|
};
|
|
let (src_strides, out_strides) = (strides(src_shape), strides(box_extent));
|
|
let last = rank - 1;
|
|
// Byte offsets into the in-memory buffers; one that does not fit `usize`
|
|
// (a 32-bit target) is out of both buffers, like one past their ends.
|
|
let bytes = |elements: u64| usize::try_from(elements).ok()?.checked_mul(elem_size);
|
|
let Some(run) = bytes(hi[last] - lo[last]) else {
|
|
return;
|
|
};
|
|
|
|
let mut idx = lo.clone();
|
|
loop {
|
|
let src_at: u64 = (0..rank)
|
|
.map(|d| (idx[d] - src_origin[d]) * src_strides[d])
|
|
.sum();
|
|
let out_at: u64 = (0..rank)
|
|
.map(|d| (idx[d] - box_start[d]) * out_strides[d])
|
|
.sum();
|
|
if let (Some(s), Some(o)) = (bytes(src_at), bytes(out_at))
|
|
&& let (Some(from), Some(to)) = (
|
|
src.get(s..s.saturating_add(run)),
|
|
out.get_mut(o..o.saturating_add(run)),
|
|
)
|
|
{
|
|
to.copy_from_slice(from);
|
|
}
|
|
// Advance over every dimension but the last.
|
|
let mut d = last;
|
|
loop {
|
|
if d == 0 {
|
|
return;
|
|
}
|
|
d -= 1;
|
|
idx[d] += 1;
|
|
if idx[d] < hi[d] {
|
|
break;
|
|
}
|
|
idx[d] = lo[d];
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Read `selection` without materialising the whole dataset, when that is
|
|
/// possible and worthwhile. `Ok(None)` means "use the full-read path": an
|
|
/// `All`/`None`/invalid selection, a layout this doesn't handle (compact,
|
|
/// virtual, storage-less), or a bounding box covering most of the dataset.
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn read_selection(
|
|
file_data: &[u8],
|
|
layout: &DataLayout,
|
|
dataspace: &Dataspace,
|
|
elem_size: usize,
|
|
pipeline: Option<&FilterPipeline>,
|
|
offset_size: u8,
|
|
length_size: u8,
|
|
selection: &Selection,
|
|
) -> Result<Option<Vec<u8>>, FormatError> {
|
|
read_selection_in(
|
|
file_data,
|
|
layout,
|
|
dataspace,
|
|
elem_size,
|
|
pipeline,
|
|
offset_size,
|
|
length_size,
|
|
selection,
|
|
)
|
|
}
|
|
|
|
/// [`read_selection`] over any [`Storage`].
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub fn read_selection_in<S: Storage + ?Sized>(
|
|
file_data: &S,
|
|
layout: &DataLayout,
|
|
dataspace: &Dataspace,
|
|
elem_size: usize,
|
|
pipeline: Option<&FilterPipeline>,
|
|
offset_size: u8,
|
|
length_size: u8,
|
|
selection: &Selection,
|
|
) -> Result<Option<Vec<u8>>, FormatError> {
|
|
let dims = &dataspace.dimensions;
|
|
if dims.is_empty() || elem_size == 0 {
|
|
return Ok(None);
|
|
}
|
|
let total = dataspace.checked_num_elements()?;
|
|
// Contiguous data is addressable in place: copy the selection's runs
|
|
// straight out of it, whatever fraction of the dataset it covers, with no
|
|
// intermediate box (and no full copy for a large selection).
|
|
if let (
|
|
DataLayout::Contiguous {
|
|
address: Some(address),
|
|
..
|
|
},
|
|
Selection::Hyperslab { .. } | Selection::Points(_),
|
|
) = (layout, selection)
|
|
{
|
|
validate(selection, dims)?;
|
|
let base = usize::try_from(*address)
|
|
.map_err(|_| FormatError::Overflow("data address exceeds usize".into()))?;
|
|
let file_len = crate::storage::len_usize(file_data);
|
|
let eof = FormatError::UnexpectedEof {
|
|
expected: base,
|
|
available: file_len,
|
|
};
|
|
if let Some(all) = file_data.as_contiguous() {
|
|
let data = all
|
|
.get(base..)
|
|
.and_then(|d| d.get(..checked_byte_len(total, elem_size).ok()?))
|
|
.ok_or(eof)?;
|
|
return crate::gather::gather::<u8>(data, dims, elem_size, selection).map(Some);
|
|
}
|
|
// Not in memory: the same bounds check, then only the selected runs
|
|
// are read.
|
|
let len = checked_byte_len(total, elem_size)
|
|
.ok()
|
|
.filter(|&len| base <= file_len && len <= file_len - base)
|
|
.ok_or(eof)?;
|
|
return crate::gather::gather_storage(
|
|
file_data,
|
|
base as u64,
|
|
len,
|
|
dims,
|
|
elem_size,
|
|
selection,
|
|
)
|
|
.map(Some);
|
|
}
|
|
let Some((box_start, box_extent)) = bounding_box(selection, dims) else {
|
|
return Ok(None);
|
|
};
|
|
let box_elements = box_extent
|
|
.iter()
|
|
.try_fold(1u64, |acc, &e| acc.checked_mul(e))
|
|
.ok_or_else(|| FormatError::Overflow("selection bounding box overflows".into()))?;
|
|
// A box covering most of the dataset gains nothing over the full path.
|
|
if box_elements.saturating_mul(2) > total {
|
|
return Ok(None);
|
|
}
|
|
let mut boxed = alloc_output(checked_byte_len(box_elements, elem_size)?)?;
|
|
|
|
match layout {
|
|
DataLayout::Chunked {
|
|
btree_address: Some(_),
|
|
..
|
|
} => {
|
|
let (chunks, chunk_dims) = list_chunks_for_read_in(
|
|
file_data,
|
|
layout,
|
|
dataspace,
|
|
elem_size,
|
|
pipeline,
|
|
offset_size,
|
|
length_size,
|
|
)?;
|
|
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)?;
|
|
// The chunks overlapping the box, in index order.
|
|
let wanted: Vec<&crate::chunked_read::ChunkInfo> = chunks
|
|
.iter()
|
|
.filter(|chunk| {
|
|
if chunk.offsets.len() < rank || chunk.address == u64::MAX {
|
|
return false;
|
|
}
|
|
let origin = &chunk.offsets[..rank];
|
|
(0..rank).all(|d| {
|
|
origin[d] < box_start[d] + box_extent[d]
|
|
&& origin[d].saturating_add(chunk_shape[d]) > box_start[d]
|
|
})
|
|
})
|
|
.collect();
|
|
// 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),
|
|
}
|
|
|
|
extract_selection_from_buffer(
|
|
&boxed,
|
|
&box_extent,
|
|
elem_size,
|
|
&translate(selection, &box_start),
|
|
)
|
|
.map(Some)
|
|
}
|