osobhandClaude Fable 5.1
c6a7bbfc67
read_raw_data_selection computed which chunks a selection intersects, threw the answer away, decoded the entire dataset and picked elements out of it — for contiguous layouts too. A 64x64 window of a 64 MB deflate dataset cost 105 ms, about half a full read; every selection cost the same whatever its size. New partial_read module: materialise only the selection's bounding box — the overlapping rows of a contiguous dataset (straight from the file bytes) or the overlapping chunks (only those are decompressed) — then run the existing extractor over that buffer with the selection translated to the box origin, so extraction semantics are exactly the full-read ones. It declines (falling back to the old path) for All/None, compact/virtual/storage-less layouts, and boxes covering more than half the dataset. That window now takes 0.39 ms, one row 2.7 ms, one column 5.2 ms. Selections are validated against the dataset shape first. They were not: a hyperslab past an edge came back padded with zeros and a point with an out-of-range column wrapped into the next row, returning the wrong element with no error. Now FormatError::SelectionOutOfBounds (also rank mismatch and overlapping blocks); the facade's fill-aware path validates too. Tests: equivalence against a reference extraction from a full read over 60 random hyperslabs/point lists per layout (contiguous, chunked, deflate) for ranks 1-3. New read_harness bench binary with before/after in BENCHMARKS.md. Co-Authored-By: Claude Fable 5.1 <[email protected]>
360 lines
13 KiB
Rust
360 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::decompress_chunk;
|
|
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: a non-zero filter mask means the
|
|
// chunk was stored unfiltered.
|
|
let decoded;
|
|
let data: &[u8] = match pipeline {
|
|
Some(pl) if chunk.filter_mask == 0 => {
|
|
decoded = decompress_chunk(raw, pl, chunk_bytes, elem_size as u32)?;
|
|
&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)
|
|
}
|