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]>

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<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

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<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
 }

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<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]);
     }

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="<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"