diff --git a/CHANGELOG.md b/CHANGELOG.md index f19acc4..a29e25a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/crates/clawhdf5-py/src/dataset.rs b/crates/clawhdf5-py/src/dataset.rs index 36467d8..aa75e7d 100644 --- a/crates/clawhdf5-py/src/dataset.rs +++ b/crates/clawhdf5-py/src/dataset.rs @@ -31,6 +31,8 @@ pub struct PyDataset { path: String, /// `None` for a dataset with a null dataspace (h5py's `Empty`). shape: Option>, + /// The chunk shape, for a chunked dataset. + chunks: Option>, datatype: Datatype, /// Why the datatype cannot be read into numpy, if it cannot. conv: Result, @@ -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 { 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::() * 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 diff --git a/crates/clawhdf5-py/src/node.rs b/crates/clawhdf5-py/src/node.rs index 0482603..bb6c2d2 100644 --- a/crates/clawhdf5-py/src/node.rs +++ b/crates/clawhdf5-py/src/node.rs @@ -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> { + 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 } diff --git a/crates/clawhdf5-py/src/select.rs b/crates/clawhdf5-py/src/select.rs index dac62ab..6b358b9 100644 --- a/crates/clawhdf5-py/src/select.rs +++ b/crates/clawhdf5-py/src/select.rs @@ -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 { + 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)>, Option) { - 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, + elem_size: usize, + ) -> (Vec, Option) { + 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 = count.iter().map(|&c| c as usize).collect(); + let shape: Vec = 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, + /// For a list: the positions along the list axis, within the block, to + /// keep (`None`: all of them). + pub pick: Option>, +} + +/// 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 { + let outer: usize = shape[..axis].iter().product(); + let inner: usize = shape[axis + 1..].iter().product::() * 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]); } diff --git a/crates/clawhdf5-py/tests/test_read_vs_h5py.py b/crates/clawhdf5-py/tests/test_read_vs_h5py.py index c5876d2..f51c534 100644 --- a/crates/clawhdf5-py/tests/test_read_vs_h5py.py +++ b/crates/clawhdf5-py/tests/test_read_vs_h5py.py @@ -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="