perf(py): read an index list one group of chunks at a time
Each run of consecutive indices was its own uncached hyperslab read, so a list over a compressed chunked dataset decoded the same chunk once per run (d[range(0, 200000, 40)] over 20 gzip chunks: 8 s, h5py 0.014 s). Plan::reads now groups the indices — a group ends only where a whole chunk holds no selected index, or, unchunked, at a gap over 64 KiB — and the selected rows are gathered from each group's block in Rust. Now 3.8 ms (h5py 4.1 ms, release, tank). The new test (1-D, 2-D and contiguous, compared with h5py, 2 s bound) took 5.8 s before. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
osobh
+13
-2
@@ -26,8 +26,8 @@
|
||||
read the whole dataset and slice it in numpy, and knew six dtypes. Now
|
||||
integers (negative from the end), slices with positive steps, `...`, one
|
||||
increasing list of integers per key and compound field names map onto the
|
||||
facade's hyperslab selection (a list becomes one hyperslab per run of
|
||||
consecutive indices), with h5py's results (numpy scalar for an all-integer
|
||||
facade's hyperslab selection (a list is read one group of neighbouring
|
||||
chunks at a time and picked from in memory), with h5py's results (numpy scalar for an all-integer
|
||||
key, 0-d array for `scalar[...]`) and h5py's errors for everything else
|
||||
(negative steps, `None`, boolean masks, out-of-range indices).
|
||||
`Dataset.dtype` is the numpy dtype h5py reports, for every integer and
|
||||
@@ -69,6 +69,17 @@
|
||||
from the file (h5py's, zero for files it wrote) and stays zero-copy.
|
||||
The h5py comparisons now also compare every byte of structured values
|
||||
(`test_compound_padding_bytes_match_h5py` and `assert_same`).
|
||||
- **Index lists no longer decode the same chunks once per run.** A list
|
||||
index was one uncached hyperslab read per run of consecutive indices, so
|
||||
on a chunked, compressed dataset every run decoded its chunk again:
|
||||
`d[list(range(0, 200000, 40))]` over 20 gzip chunks took 8 s (h5py:
|
||||
0.014 s). The list is now read in groups — for a chunked dataset a group
|
||||
ends only where a whole chunk holds no selected index, so each chunk is
|
||||
decoded once; otherwise at a gap of more than 64 KiB — and the selected
|
||||
rows are picked from each group in Rust. The same read now takes 3.8 ms
|
||||
(h5py 4.1 ms; release build on tank, best of 5).
|
||||
`test_a_long_index_list_decodes_each_chunk_once` compares 1-D, 2-D and
|
||||
contiguous cases with h5py under a 2 s bound (5.8 s before, debug build).
|
||||
- **CI builds and tests the Python package.** It was excluded from CI.
|
||||
`scripts/ci-test.sh` now lints `clawhdf5-py`, builds the wheel with
|
||||
maturin, unpacks it under `target/` and runs the pytest suite; skipped
|
||||
|
||||
@@ -31,6 +31,8 @@ pub struct PyDataset {
|
||||
path: String,
|
||||
/// `None` for a dataset with a null dataspace (h5py's `Empty`).
|
||||
shape: Option<Vec<u64>>,
|
||||
/// The chunk shape, for a chunked dataset.
|
||||
chunks: Option<Vec<u64>>,
|
||||
datatype: Datatype,
|
||||
/// Why the datatype cannot be read into numpy, if it cannot.
|
||||
conv: Result<Converter, String>,
|
||||
@@ -56,10 +58,14 @@ impl PyDataset {
|
||||
};
|
||||
let conv = Converter::new(py, &datatype, file.superblock().offset_size)
|
||||
.map_err(|e| e.value(py).to_string());
|
||||
let chunks = shape
|
||||
.as_ref()
|
||||
.and_then(|s| node::chunk_shape(&file, &hdr, s.len()));
|
||||
Ok(Self {
|
||||
file,
|
||||
path,
|
||||
shape,
|
||||
chunks,
|
||||
datatype,
|
||||
conv,
|
||||
})
|
||||
@@ -81,17 +87,23 @@ impl PyDataset {
|
||||
let arr = if plan.is_empty() {
|
||||
conv.empty(py, &out_shape)?
|
||||
} else {
|
||||
let (reads, list_axis) = plan.reads(dims);
|
||||
let (vl, elem_size, unit) = (conv.is_vl(), conv.elem_size, conv.vl_unit);
|
||||
let list_axis = plan.list_axis();
|
||||
let chunk_len = match (&self.chunks, list_axis) {
|
||||
(Some(c), Some(a)) => c.get(a).copied(),
|
||||
_ => None,
|
||||
};
|
||||
let (reads, list_axis) = plan.reads(dims, chunk_len, elem_size);
|
||||
let read_shape = plan.read_shape();
|
||||
let file = &*self.file;
|
||||
let path = self.path.as_str();
|
||||
let (vl, elem_size, unit) = (conv.is_vl(), conv.elem_size, conv.vl_unit);
|
||||
// Everything below touches only Rust data: release the GIL.
|
||||
let read = || -> Result<Elements, ReadError> {
|
||||
let ds = file.dataset(path)?;
|
||||
let mut blocks = Vec::with_capacity(reads.len());
|
||||
for (sel, shape) in reads {
|
||||
let raw = ds.read_selection(&sel)?;
|
||||
for read in reads {
|
||||
let raw = ds.read_selection(&read.sel)?;
|
||||
let mut shape = read.shape;
|
||||
let want = shape.iter().product::<usize>() * elem_size;
|
||||
if raw.len() != want {
|
||||
return Err(ReadError::Other(format!(
|
||||
@@ -99,6 +111,14 @@ impl PyDataset {
|
||||
raw.len()
|
||||
)));
|
||||
}
|
||||
let raw = match (&read.pick, list_axis) {
|
||||
(Some(pick), Some(axis)) => {
|
||||
let kept = select::gather_along(&raw, &shape, axis, pick, elem_size);
|
||||
shape[axis] = pick.len();
|
||||
kept
|
||||
}
|
||||
_ => raw,
|
||||
};
|
||||
blocks.push((raw, shape));
|
||||
}
|
||||
// Several blocks only for a list index: join their bytes
|
||||
|
||||
@@ -135,6 +135,33 @@ pub(crate) fn dataspace(file: &clawhdf5_rs::File, hdr: &ObjectHeader) -> PyResul
|
||||
})
|
||||
}
|
||||
|
||||
/// The chunk shape of a chunked dataset (one entry per dataset dimension),
|
||||
/// or `None` for other layouts or a layout message that does not parse.
|
||||
pub(crate) fn chunk_shape(
|
||||
file: &clawhdf5_rs::File,
|
||||
hdr: &ObjectHeader,
|
||||
rank: usize,
|
||||
) -> Option<Vec<u64>> {
|
||||
let sb = file.superblock();
|
||||
let msg = hdr
|
||||
.messages
|
||||
.iter()
|
||||
.find(|m| m.msg_type == MessageType::DataLayout)?;
|
||||
match clawhdf5_format::data_layout::DataLayout::parse(&msg.data, sb.offset_size, sb.length_size)
|
||||
.ok()?
|
||||
{
|
||||
clawhdf5_format::data_layout::DataLayout::Chunked {
|
||||
chunk_dimensions, ..
|
||||
} if chunk_dimensions.len() >= rank => Some(
|
||||
chunk_dimensions[..rank]
|
||||
.iter()
|
||||
.map(|&d| u64::from(d))
|
||||
.collect(),
|
||||
),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_null(space: &Dataspace) -> bool {
|
||||
space.space_type == DataspaceType::Null
|
||||
}
|
||||
|
||||
@@ -62,6 +62,11 @@ impl Plan {
|
||||
self.axes.iter().map(|a| a.len() as usize).collect()
|
||||
}
|
||||
|
||||
/// The axis indexed by a list, if any.
|
||||
pub fn list_axis(&self) -> Option<usize> {
|
||||
self.axes.iter().position(|a| matches!(a, Axis::List(_)))
|
||||
}
|
||||
|
||||
/// Whether the selection is empty.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.axes.iter().any(|a| a.len() == 0)
|
||||
@@ -69,16 +74,42 @@ impl Plan {
|
||||
|
||||
/// The hyperslab reads that make up this selection, each with the shape
|
||||
/// of its block (index axes kept at length 1). More than one only when an
|
||||
/// axis is indexed by a list: one read per run of consecutive indices,
|
||||
/// joined along `list_axis` afterwards (`join_along`).
|
||||
pub fn reads(&self, dims: &[u64]) -> (Vec<(Selection, Vec<usize>)>, Option<usize>) {
|
||||
let list_axis = self.axes.iter().position(|a| matches!(a, Axis::List(_)));
|
||||
let runs: Vec<(u64, u64)> = match list_axis.map(|i| &self.axes[i]) {
|
||||
Some(Axis::List(idx)) => consecutive_runs(idx),
|
||||
_ => vec![(0, 0)],
|
||||
/// axis is indexed by a list; those are joined along `list_axis`
|
||||
/// afterwards (`join_along`), after keeping each read's `pick` rows.
|
||||
///
|
||||
/// A list is read in groups, each one hyperslab over a stretch of the
|
||||
/// axis, not once per index: every read decodes the chunks it touches
|
||||
/// (and lists the dataset's chunks), so a read per run of indices decoded
|
||||
/// the same chunk again and again. For a chunked dataset (`chunk_len` is
|
||||
/// the chunk's length along the list axis) a group ends only where a
|
||||
/// whole chunk holds no selected index, so no chunk is decoded twice or
|
||||
/// without need. Otherwise a group ends at a gap of more than
|
||||
/// [`MAX_GAP_BYTES`] of unselected data.
|
||||
pub fn reads(
|
||||
&self,
|
||||
dims: &[u64],
|
||||
chunk_len: Option<u64>,
|
||||
elem_size: usize,
|
||||
) -> (Vec<Read>, Option<usize>) {
|
||||
let list_axis = self.list_axis();
|
||||
let groups: Vec<&[u64]> = match list_axis.map(|i| &self.axes[i]) {
|
||||
Some(Axis::List(idx)) => {
|
||||
let row_bytes = self.row_bytes(elem_size);
|
||||
group_indices(idx, |last, next| match chunk_len {
|
||||
Some(c) if c > 0 => next / c <= last / c + 1,
|
||||
_ => (next - last - 1).saturating_mul(row_bytes) <= MAX_GAP_BYTES,
|
||||
})
|
||||
}
|
||||
_ => vec![&[]],
|
||||
};
|
||||
let mut out = Vec::with_capacity(runs.len());
|
||||
for (run_start, run_len) in runs {
|
||||
let mut out = Vec::with_capacity(groups.len());
|
||||
for group in groups {
|
||||
let (first, span) = match (group.first(), group.last()) {
|
||||
(Some(&f), Some(&l)) => (f, l - f + 1),
|
||||
_ => (0, 0),
|
||||
};
|
||||
let pick = (span != group.len() as u64)
|
||||
.then(|| group.iter().map(|&i| (i - first) as usize).collect());
|
||||
let mut start = Vec::with_capacity(dims.len());
|
||||
let mut stride = Vec::with_capacity(dims.len());
|
||||
let mut count = Vec::with_capacity(dims.len());
|
||||
@@ -86,14 +117,14 @@ impl Plan {
|
||||
let (s, st, c) = match axis {
|
||||
Axis::Index(i) => (*i, 1, 1),
|
||||
Axis::Slice { start, step, count } => (*start, *step, *count),
|
||||
Axis::List(_) => (run_start, 1, run_len),
|
||||
Axis::List(_) => (first, 1, span),
|
||||
};
|
||||
start.push(s);
|
||||
// A stride only matters between blocks; keep it >= 1.
|
||||
stride.push(if c <= 1 { 1 } else { st });
|
||||
count.push(c);
|
||||
}
|
||||
let block_shape: Vec<usize> = count.iter().map(|&c| c as usize).collect();
|
||||
let shape: Vec<usize> = count.iter().map(|&c| c as usize).collect();
|
||||
let whole = start.iter().all(|&s| s == 0)
|
||||
&& stride.iter().all(|&s| s == 1)
|
||||
&& count.as_slice() == dims;
|
||||
@@ -108,21 +139,74 @@ impl Plan {
|
||||
block,
|
||||
}
|
||||
};
|
||||
out.push((sel, block_shape));
|
||||
out.push(Read { sel, shape, pick });
|
||||
}
|
||||
(out, list_axis)
|
||||
}
|
||||
|
||||
/// Bytes of one step along the list axis within a read's bounding box.
|
||||
fn row_bytes(&self, elem_size: usize) -> u64 {
|
||||
self.axes
|
||||
.iter()
|
||||
.map(|a| match a {
|
||||
Axis::Slice { step, count, .. } if *count > 0 => (count - 1) * step + 1,
|
||||
_ => 1,
|
||||
})
|
||||
.fold(elem_size as u64, u64::saturating_mul)
|
||||
}
|
||||
}
|
||||
|
||||
fn consecutive_runs(idx: &[u64]) -> Vec<(u64, u64)> {
|
||||
let mut runs: Vec<(u64, u64)> = Vec::new();
|
||||
for &i in idx {
|
||||
match runs.last_mut() {
|
||||
Some((s, n)) if *s + *n == i => *n += 1,
|
||||
_ => runs.push((i, 1)),
|
||||
/// Unselected data a read of a non-chunked dataset copies through rather
|
||||
/// than start another read.
|
||||
pub(crate) const MAX_GAP_BYTES: u64 = 64 * 1024;
|
||||
|
||||
/// One hyperslab read of a selection.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub(crate) struct Read {
|
||||
pub sel: Selection,
|
||||
/// The block's shape (index axes at length 1).
|
||||
pub shape: Vec<usize>,
|
||||
/// For a list: the positions along the list axis, within the block, to
|
||||
/// keep (`None`: all of them).
|
||||
pub pick: Option<Vec<usize>>,
|
||||
}
|
||||
|
||||
/// Split increasing indices into groups; `joins(last, next)` says whether
|
||||
/// `next` extends the group whose last index is `last`.
|
||||
fn group_indices(idx: &[u64], joins: impl Fn(u64, u64) -> bool) -> Vec<&[u64]> {
|
||||
let mut groups = Vec::new();
|
||||
let mut from = 0;
|
||||
for k in 1..idx.len() {
|
||||
if !joins(idx[k - 1], idx[k]) {
|
||||
groups.push(&idx[from..k]);
|
||||
from = k;
|
||||
}
|
||||
}
|
||||
runs
|
||||
if from < idx.len() {
|
||||
groups.push(&idx[from..]);
|
||||
}
|
||||
groups
|
||||
}
|
||||
|
||||
/// Keep the elements at positions `pick` along `axis` of a row-major block.
|
||||
pub(crate) fn gather_along(
|
||||
bytes: &[u8],
|
||||
shape: &[usize],
|
||||
axis: usize,
|
||||
pick: &[usize],
|
||||
elem_size: usize,
|
||||
) -> Vec<u8> {
|
||||
let outer: usize = shape[..axis].iter().product();
|
||||
let inner: usize = shape[axis + 1..].iter().product::<usize>() * elem_size;
|
||||
let len = shape[axis];
|
||||
let mut out = Vec::with_capacity(outer * pick.len() * inner);
|
||||
for o in 0..outer {
|
||||
for &p in pick {
|
||||
let at = (o * len + p) * inner;
|
||||
out.extend_from_slice(&bytes[at..at + inner]);
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Join row-major blocks of `elem_size`-byte elements whose shapes differ
|
||||
@@ -336,12 +420,53 @@ mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn runs_group_consecutive_indices() {
|
||||
fn indices_group_by_chunk() {
|
||||
let chunked = |c: u64| move |last: u64, next: u64| next / c <= last / c + 1;
|
||||
// Chunks of 10: 3, 5 and 15 are in neighbouring chunks; 42 skips two.
|
||||
let idx = [3, 5, 15, 42, 43, 99];
|
||||
assert_eq!(
|
||||
consecutive_runs(&[1, 2, 3, 7, 9, 10]),
|
||||
vec![(1, 3), (7, 1), (9, 2)]
|
||||
group_indices(&idx, chunked(10)),
|
||||
vec![&[3, 5, 15][..], &[42, 43], &[99]]
|
||||
);
|
||||
assert_eq!(group_indices(&[], chunked(10)), Vec::<&[u64]>::new());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_list_reads_once_per_group() {
|
||||
let plan = Plan {
|
||||
axes: vec![
|
||||
Axis::List(vec![0, 2, 3, 40]),
|
||||
Axis::Slice {
|
||||
start: 0,
|
||||
step: 1,
|
||||
count: 5,
|
||||
},
|
||||
],
|
||||
fields: vec![],
|
||||
scalar: false,
|
||||
};
|
||||
// Chunks of 8 rows: rows 0-3 are one read, row 40 another.
|
||||
let (reads, axis) = plan.reads(&[50, 5], Some(8), 4);
|
||||
assert_eq!(axis, Some(0));
|
||||
assert_eq!(reads.len(), 2);
|
||||
assert_eq!(reads[0].shape, vec![4, 5]);
|
||||
assert_eq!(reads[0].pick, Some(vec![0, 2, 3]));
|
||||
assert_eq!(reads[1].shape, vec![1, 5]);
|
||||
assert_eq!(reads[1].pick, None);
|
||||
// Not chunked: a gap under MAX_GAP_BYTES is read through.
|
||||
let (reads, _) = plan.reads(&[50, 5], None, 4);
|
||||
assert_eq!(reads.len(), 1);
|
||||
assert_eq!(reads[0].pick, Some(vec![0, 2, 3, 40]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gather_keeps_picked_rows() {
|
||||
// A 2x3 block of 1-byte elements; keep columns 0 and 2.
|
||||
let block = [1, 2, 3, 4, 5, 6];
|
||||
assert_eq!(
|
||||
gather_along(&block, &[2, 3], 1, &[0, 2], 1),
|
||||
vec![1, 3, 4, 6]
|
||||
);
|
||||
assert_eq!(consecutive_runs(&[]), vec![]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -374,9 +499,16 @@ mod tests {
|
||||
fields: vec![],
|
||||
scalar: false,
|
||||
};
|
||||
let (reads, list) = plan.reads(&[4, 3]);
|
||||
let (reads, list) = plan.reads(&[4, 3], None, 8);
|
||||
assert_eq!(list, None);
|
||||
assert_eq!(reads, vec![(Selection::All, vec![4, 3])]);
|
||||
assert_eq!(
|
||||
reads,
|
||||
vec![Read {
|
||||
sel: Selection::All,
|
||||
shape: vec![4, 3],
|
||||
pick: None
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -393,18 +525,19 @@ mod tests {
|
||||
fields: vec![],
|
||||
scalar: false,
|
||||
};
|
||||
let (reads, _) = plan.reads(&[4, 8]);
|
||||
let (reads, _) = plan.reads(&[4, 8], None, 8);
|
||||
assert_eq!(
|
||||
reads,
|
||||
vec![(
|
||||
Selection::Hyperslab {
|
||||
vec![Read {
|
||||
sel: Selection::Hyperslab {
|
||||
start: vec![2, 1],
|
||||
stride: vec![1, 3],
|
||||
count: vec![1, 2],
|
||||
block: vec![1, 1],
|
||||
},
|
||||
vec![1, 2]
|
||||
)]
|
||||
shape: vec![1, 2],
|
||||
pick: None
|
||||
}]
|
||||
);
|
||||
assert_eq!(plan.out_shape(), vec![2]);
|
||||
}
|
||||
|
||||
@@ -531,3 +531,35 @@ def test_compound_padding_bytes_match_h5py(pair):
|
||||
assert ours[name][key].tobytes() == theirs[name][key].tobytes(), f"{name}[{key!r}]"
|
||||
for key in [(slice(None), [0, 2]), ([0, 2, 3], slice(None)), ([1, 3], 1)]:
|
||||
assert ours["cmp/nested_2d"][key].tobytes() == theirs["cmp/nested_2d"][key].tobytes(), key
|
||||
|
||||
|
||||
def test_a_long_index_list_decodes_each_chunk_once(h5py, tmp_path):
|
||||
"""An index list is read one group of chunks at a time, not one
|
||||
hyperslab per run of indices: 5000 runs over 20 gzip chunks used to
|
||||
decode the chunks 5000 times (8 s, against h5py's 0.014 s)."""
|
||||
import time
|
||||
|
||||
path = str(tmp_path / "long_list.h5")
|
||||
data = np.arange(200000, dtype="<f8")
|
||||
grid = np.arange(400 * 3000, dtype="<i4").reshape(400, 3000)
|
||||
with h5py.File(path, "w") as f:
|
||||
f.create_dataset("d", data=data, chunks=(10000,), compression="gzip")
|
||||
f.create_dataset("grid", data=grid, chunks=(50, 100), compression="gzip")
|
||||
f.create_dataset("flat", data=data) # contiguous
|
||||
rng = np.random.default_rng(3)
|
||||
cases = [
|
||||
("d", list(range(0, 200000, 40))),
|
||||
("d", sorted(rng.choice(200000, 3000, replace=False).tolist())),
|
||||
("flat", list(range(0, 200000, 40))),
|
||||
("flat", [0, 7, 199999]),
|
||||
("grid", (slice(None), list(range(0, 3000, 3)))),
|
||||
("grid", (sorted(rng.choice(400, 150, replace=False).tolist()), slice(5, 2900, 7))),
|
||||
("grid", (7, [0, 1, 2, 2000, 2999])),
|
||||
]
|
||||
with h5py.File(path, "r") as theirs, clawhdf5.File(path, "r") as ours:
|
||||
for name, key in cases:
|
||||
t0 = time.perf_counter()
|
||||
got = ours[name][key]
|
||||
took = time.perf_counter() - t0
|
||||
assert_same(got, theirs[name][key], f"{name}[{len(key)}-key]")
|
||||
assert took < 2.0, f"{name}: {took:.2f} s"
|
||||
|
||||
Reference in New Issue
Block a user