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