osobhandClaude Opus 5.5
585e14d5e2
A chunk's filter mask has one bit per pipeline filter; bit i set means filter i was not applied to that chunk (an optional filter that declined, or a direct chunk write). Every read path treated any nonzero mask as "no filters applied" and returned the stored bytes, so a chunk that skipped only gzip in a shuffle+gzip pipeline came back still shuffled (h5py write_direct_chunk with filter_mask=0b10: 8 of 32 values wrong). decompress_chunk_masked undoes the filters the mask leaves set and skips the rest; an unsupported filter is no longer an error when the chunk skipped it. The full, cached, sweep, indexed, parallel and selection (partial_read) paths all use it, and a chunk is copied straight from the file only when every filter was skipped. decompress_chunk is the mask-0 case. Regression: h5py_partial_filter_mask_skips_only_masked_filters (1-D shuffle+gzip with masks 0, 0b10 and 0b11; 2-D with 0b01; full and hyperslab reads) and filter_mask_skips_only_the_masked_filters. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
366 lines
13 KiB
Rust
366 lines
13 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. Here the selection's
|
|
//! bounding box is materialised instead — only the rows of a contiguous
|
|
//! dataset, or only the chunks, that overlap it — and the existing 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};
|
|
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_masked};
|
|
use crate::selection::Selection;
|
|
|
|
/// 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;
|
|
let run = ((hi[last] - lo[last]) as usize) * elem_size;
|
|
|
|
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();
|
|
let (s, o) = (src_at as usize * elem_size, out_at as usize * elem_size);
|
|
if let (Some(from), Some(to)) = (src.get(s..s + run), out.get_mut(o..o + 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> {
|
|
let dims = &dataspace.dimensions;
|
|
if dims.is_empty() || elem_size == 0 {
|
|
return Ok(None);
|
|
}
|
|
let Some((box_start, box_extent)) = bounding_box(selection, dims) else {
|
|
return Ok(None);
|
|
};
|
|
let total = dataspace.checked_num_elements()?;
|
|
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::Contiguous {
|
|
address: Some(address),
|
|
..
|
|
} => {
|
|
let base = usize::try_from(*address)
|
|
.map_err(|_| FormatError::Overflow("data address exceeds usize".into()))?;
|
|
let data = file_data
|
|
.get(base..)
|
|
.and_then(|d| d.get(..checked_byte_len(total, elem_size).ok()?))
|
|
.ok_or(FormatError::UnexpectedEof {
|
|
expected: base,
|
|
available: file_data.len(),
|
|
})?;
|
|
let origin = vec![0u64; dims.len()];
|
|
copy_overlap(
|
|
data,
|
|
&origin,
|
|
dims,
|
|
&mut boxed,
|
|
&box_start,
|
|
&box_extent,
|
|
elem_size,
|
|
);
|
|
}
|
|
DataLayout::Chunked {
|
|
btree_address: Some(_),
|
|
..
|
|
} => {
|
|
let (chunks, chunk_dims) = list_chunks(
|
|
file_data,
|
|
layout,
|
|
dataspace,
|
|
elem_size,
|
|
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)?;
|
|
for chunk in &chunks {
|
|
if chunk.offsets.len() < rank || chunk.address == u64::MAX {
|
|
continue;
|
|
}
|
|
let origin = &chunk.offsets[..rank];
|
|
let overlaps = (0..rank).all(|d| {
|
|
origin[d] < box_start[d] + box_extent[d]
|
|
&& origin[d].saturating_add(chunk_shape[d]) > box_start[d]
|
|
});
|
|
if !overlaps {
|
|
continue;
|
|
}
|
|
let at = usize::try_from(chunk.address)
|
|
.map_err(|_| FormatError::Overflow("chunk address exceeds usize".into()))?;
|
|
let raw = at
|
|
.checked_add(chunk.chunk_size as usize)
|
|
.and_then(|end| file_data.get(at..end))
|
|
.ok_or(FormatError::UnexpectedEof {
|
|
expected: at.saturating_add(chunk.chunk_size as usize),
|
|
available: file_data.len(),
|
|
})?;
|
|
// Mirrors the full-read path: filter-mask bit i set means
|
|
// filter i was not applied to this chunk.
|
|
let decoded;
|
|
let data: &[u8] = match pipeline {
|
|
Some(pl) if !all_filters_skipped(pl, chunk.filter_mask) => {
|
|
decoded = decompress_chunk_masked(
|
|
raw,
|
|
pl,
|
|
chunk_bytes,
|
|
elem_size as u32,
|
|
chunk.filter_mask,
|
|
)?;
|
|
&decoded
|
|
}
|
|
_ => raw,
|
|
};
|
|
copy_overlap(
|
|
data,
|
|
origin,
|
|
&chunk_shape,
|
|
&mut boxed,
|
|
&box_start,
|
|
&box_extent,
|
|
elem_size,
|
|
);
|
|
}
|
|
}
|
|
_ => return Ok(None),
|
|
}
|
|
|
|
extract_selection_from_buffer(
|
|
&boxed,
|
|
&box_extent,
|
|
elem_size,
|
|
&translate(selection, &box_start),
|
|
)
|
|
.map(Some)
|
|
}
|