format: read a contiguous selection's runs merged across small gaps

gather_storage merged only runs that touch, so a strided selection of a
contiguous dataset over a Storage became one range (and one owned Vec) per
element: a stride-2 read of 32M f32 through File::open_storage made
16,777,232 read_at calls, took 2.0 s and peaked at 2.09 GB.

The selection is now walked twice. The first walk checks the runs and plans
spans: runs in increasing order at most 4 KiB apart (GATHER_GAP_BYTES) are
read as one span up to 8 MiB (GATHER_SPAN_BYTES; a longer run is split), so
nothing is stored per run. The spans are fetched in RAW_BATCH_BYTES batches
while the second walk copies each run out of its span. Same checks and
errors as before.

The same read is now 32 reads and 0.31 s (File::open: 0.08 s).
contiguous_read_interop: every h5py-checked selection is also read through
File::open_storage and must give libhdf5's bytes; a new test bounds the
range reads of strided, blocked, column and point selections (stride 2: at
most 1 data read; 563,200 before).

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>

CHANGELOG.md

@@ -61,7 +61,10 @@
     `chunks x extent` bytes (`tests/raw_fetch_bounds.rs`). Chunks the
     file's chunk cache already holds are not fetched. A selection fetches
     only the chunks its bounding box overlaps; a contiguous selection only
-    its runs (adjacent ones merged). A global-heap collection is read once
+    its runs, merged into reads of up to 8 MiB across gaps of up to 4 KiB
+    (a stride-2 selection of 32M `f32` is 32 reads and 0.3 s over a
+    `CountingStorage`, where one read per element was 16.8M reads, 2.0 s
+    and 2.1 GB peak). A global-heap collection is read once
     per resolver and kept (within the resolver's 32 MiB budget).
   - Each extent's bounds error is the one the slice readers gave, reported
     when the read reaches that extent, so a damaged file fails with the

crates/clawhdf5-format/src/gather.rs

@@ -13,7 +13,7 @@ use alloc::{vec, vec::Vec};
 use crate::data_read::NativeElement;
 use crate::error::FormatError;
 use crate::selection::Selection;
-use crate::storage::Storage;
+use crate::storage::{ExtentBytes, ExtentReq, Storage, raw_batches};
 
 /// Row-major element strides of `dims` (the last dimension has stride 1).
 fn strides(dims: &[u64]) -> Vec<u64> {
@@ -262,12 +262,88 @@ pub(crate) fn gather<T: NativeElement>(
     Ok(out)
 }
 
+/// Largest gap between two of a selection's runs that [`gather_storage`]
+/// reads through rather than asking for the runs separately: skipping a
+/// few KiB costs a remote backend far less than another request (and a
+/// local one less than another call and allocation).
+pub(crate) const GATHER_GAP_BYTES: usize = 4 << 10;
+
+/// Largest single read [`gather_storage`] makes of a selection's runs: runs
+/// are merged into reads up to this size, and a longer run is split.
+pub(crate) const GATHER_SPAN_BYTES: usize = 8 << 20;
+
+/// Call `emit(first_element, element_count)` for each run of a validated
+/// hyperslab or point selection (in output order; see [`hyperslab_runs`]),
+/// or the error for a hyperslab of the wrong rank or a point outside `dims`
+/// (runs before that point have been emitted).
+fn selection_runs(
+    dims: &[u64],
+    selection: &Selection,
+    emit: &mut dyn FnMut(u64, u64),
+) -> Result<(), FormatError> {
+    match selection {
+        Selection::Hyperslab {
+            start,
+            stride,
+            count,
+            block,
+        } => {
+            let rank = dims.len();
+            if [start.len(), stride.len(), count.len(), block.len()] != [rank; 4] {
+                return Err(FormatError::SelectionOutOfBounds(
+                    "hyperslab rank does not match dataset rank".into(),
+                ));
+            }
+            hyperslab_runs(dims, start, stride, count, block, emit);
+        }
+        Selection::Points(points) => {
+            let strides = strides(dims);
+            let mut coalesce = Coalesce {
+                start: 0,
+                len: 0,
+                emit,
+            };
+            for p in points {
+                if p.len() != dims.len() || p.iter().zip(dims).any(|(c, n)| c >= n) {
+                    return Err(FormatError::SelectionOutOfBounds(
+                        "selection addresses elements outside the dataset".into(),
+                    ));
+                }
+                let at = p
+                    .iter()
+                    .zip(&strides)
+                    .fold(0u64, |acc, (c, s)| acc.wrapping_add(c.wrapping_mul(*s)));
+                coalesce.push(at, 1);
+            }
+            coalesce.flush();
+        }
+        Selection::None | Selection::All => {}
+    }
+    Ok(())
+}
+
+/// One read of [`gather_storage`]: bytes `[start, end)` of the dataset,
+/// which hold the output's bytes up to `out_end` (from where the previous
+/// span's end left off).
+#[derive(Clone, Copy)]
+struct Span {
+    start: usize,
+    end: usize,
+    out_end: usize,
+}
+
 /// [`gather`] of bytes (`T = u8`) from a dataset that is not in memory: the
 /// dataset's `src_len` bytes start at `base` in `file`, which must hold all
-/// of them (the caller checks). The selection's runs are collected first,
-/// adjacent ones merged, and fetched with one [`Storage::read_ranges`] call,
-/// so only the selected bytes are read. Same checks and errors as
-/// [`gather`].
+/// of them (the caller checks). Same checks and errors as [`gather`].
+///
+/// The selection's runs are walked twice. The first walk checks them and
+/// plans the reads: runs in increasing order with at most
+/// [`GATHER_GAP_BYTES`] between them are read as one span (the gap is read
+/// and dropped), up to [`GATHER_SPAN_BYTES`] per span. So a strided
+/// selection is a few large reads, not one per element, and nothing is
+/// allocated per run. The spans are fetched batch by batch (one
+/// [`Storage::read_ranges`] call per [`crate::storage::RAW_BATCH_BYTES`])
+/// while the second walk copies each run out of its span.
 pub(crate) fn gather_storage<S: Storage + ?Sized>(
     file: &S,
     base: u64,
@@ -297,11 +373,15 @@ pub(crate) fn gather_storage<S: Storage + ?Sized>(
         }
     };
     let out_bytes = crate::chunked_read::checked_byte_len(n_elements, elem_size)?;
-    // Runs as (byte offset in the dataset, byte length), in output order.
-    let mut runs: Vec<(usize, usize)> = Vec::new();
+    let outside = || {
+        FormatError::SelectionOutOfBounds("selection addresses elements outside the dataset".into())
+    };
+
+    // First walk: check every run and plan the spans.
+    let mut spans: Vec<Span> = Vec::new();
     let mut total = 0usize;
     let mut failed = false;
-    let mut collect = |first: u64, n: u64| {
+    selection_runs(dims, selection, &mut |first: u64, n: u64| {
         if failed {
             return;
         }
@@ -314,78 +394,123 @@ pub(crate) fn gather_storage<S: Storage + ?Sized>(
                     .and_then(|n| n.checked_mul(elem_size)),
             )
             .and_then(|(at, len)| Some((at, len, at.checked_add(len)?)));
-        match range {
-            Some((at, len, end)) if end <= src_len && total + len <= out_bytes => {
-                match runs.last_mut() {
-                    Some((a, l)) if *a + *l == at => *l += len,
-                    _ => runs.push((at, len)),
-                }
-                total += len;
-            }
-            _ => failed = true,
-        }
-    };
-    let mut bad_point = false;
-    match selection {
-        Selection::Hyperslab {
-            start,
-            stride,
-            count,
-            block,
-        } => {
-            let rank = dims.len();
-            if [start.len(), stride.len(), count.len(), block.len()] != [rank; 4] {
-                return Err(FormatError::SelectionOutOfBounds(
-                    "hyperslab rank does not match dataset rank".into(),
-                ));
-            }
-            hyperslab_runs(dims, start, stride, count, block, &mut collect);
-        }
-        Selection::Points(points) => {
-            let strides = strides(dims);
-            let mut coalesce = Coalesce {
-                start: 0,
-                len: 0,
-                emit: &mut collect,
-            };
-            for p in points {
-                if p.len() != dims.len() || p.iter().zip(dims).any(|(c, n)| c >= n) {
-                    bad_point = true;
-                    break;
-                }
-                let at = p
-                    .iter()
-                    .zip(&strides)
-                    .fold(0u64, |acc, (c, s)| acc.wrapping_add(c.wrapping_mul(*s)));
-                coalesce.push(at, 1);
-            }
-            coalesce.flush();
-        }
-        Selection::None | Selection::All => {}
-    }
-    if failed || bad_point || total != out_bytes {
-        return Err(FormatError::SelectionOutOfBounds(
-            "selection addresses elements outside the dataset".into(),
-        ));
-    }
-    let ranges: Vec<core::ops::Range<u64>> = runs
-        .iter()
-        .map(|&(at, len)| base + at as u64..base + (at + len) as u64)
-        .collect();
-    let fetched = file.read_ranges(&ranges)?;
-    if fetched.len() != ranges.len() {
-        return Err(FormatError::Storage(
-            "read_ranges returned the wrong number of ranges".into(),
-        ));
-    }
-    let mut out = crate::bulk_alloc::vec_for_bulk(out_bytes);
-    for (bytes, &(_, len)) in fetched.iter().zip(&runs) {
-        let Some(b) = bytes.get(..len) else {
-            return Err(FormatError::Storage(
-                "short read inside the file (the storage shrank or the backend failed)".into(),
-            ));
+        let Some((mut at, mut len)) = range
+            .filter(|&(_, len, end)| end <= src_len && len <= out_bytes - total)
+            .map(|(at, len, _)| (at, len))
+        else {
+            failed = true;
+            return;
         };
-        out.extend_from_slice(b);
+        while len > 0 {
+            let room = match spans.last_mut() {
+                Some(s)
+                    if at >= s.end
+                        && at - s.end <= GATHER_GAP_BYTES
+                        && at - s.start < GATHER_SPAN_BYTES =>
+                {
+                    let take = len.min(GATHER_SPAN_BYTES - (at - s.start));
+                    s.end = at + take;
+                    s.out_end += take;
+                    take
+                }
+                _ => {
+                    let take = len.min(GATHER_SPAN_BYTES);
+                    spans.push(Span {
+                        start: at,
+                        end: at + take,
+                        out_end: total + take,
+                    });
+                    take
+                }
+            };
+            total += room;
+            at += room;
+            len -= room;
+        }
+    })?;
+    if failed || total != out_bytes {
+        return Err(outside());
+    }
+
+    // The spans' reads, and the batches they are fetched in.
+    let reqs: Vec<ExtentReq> = spans
+        .iter()
+        .map(|s| ExtentReq {
+            addr: base + s.start as u64,
+            len: s.end - s.start,
+            fetch: Some(s.end - s.start),
+        })
+        .collect();
+    let batches = raw_batches(reqs.len(), false, |i| reqs[i].len);
+
+    // Second walk: copy each run out of its span, fetching each batch of
+    // spans when the walk reaches it (and dropping the previous one).
+    let mut out = crate::bulk_alloc::vec_for_bulk(out_bytes);
+    let mut span = 0usize;
+    let mut batch = 0usize;
+    let mut fetched: Option<ExtentBytes<'_>> = None;
+    let mut error: Option<FormatError> = None;
+    selection_runs(dims, selection, &mut |first: u64, n: u64| {
+        if error.is_some() {
+            return;
+        }
+        // Checked by the first walk.
+        let mut at = first as usize * elem_size;
+        let mut len = n as usize * elem_size;
+        while len > 0 {
+            while spans.get(span).is_some_and(|s| s.out_end <= out.len()) {
+                span += 1;
+            }
+            if fetched.is_none() || span >= batches[batch].end {
+                fetched = None;
+                while batches.get(batch).is_some_and(|b| span >= b.end) {
+                    batch += 1;
+                }
+                let (Some(b), Some(_)) = (batches.get(batch).cloned(), spans.get(span)) else {
+                    // The second walk emitted more than the first.
+                    error = Some(outside());
+                    return;
+                };
+                match ExtentBytes::fetch(file, &reqs[b.clone()], b.start) {
+                    Ok(f) => fetched = Some(f),
+                    Err(e) => {
+                        error = Some(e);
+                        return;
+                    }
+                }
+            }
+            let s = spans[span];
+            let take = len.min(s.out_end - out.len());
+            let bytes = match fetched
+                .as_ref()
+                .map(|f| f.get(span, &reqs[span]))
+                .unwrap_or_else(|| Err(outside()))
+            {
+                Ok(b) => b,
+                Err(e) => {
+                    error = Some(e);
+                    return;
+                }
+            };
+            match at
+                .checked_sub(s.start)
+                .and_then(|o| bytes.get(o..o.checked_add(take)?))
+            {
+                Some(b) => out.extend_from_slice(b),
+                None => {
+                    error = Some(outside());
+                    return;
+                }
+            }
+            at += take;
+            len -= take;
+        }
+    })?;
+    if let Some(e) = error {
+        return Err(e);
+    }
+    if out.len() != out_bytes {
+        return Err(outside());
     }
     Ok(out)
 }

crates/clawhdf5-format/src/storage.rs

@@ -420,7 +420,7 @@ pub(crate) enum Extent<'a> {
 impl<'a> ExtentBytes<'a> {
     /// Fetch `reqs`, extents `base..base + reqs.len()` of the read: the
     /// bytes of those wanted, and the bounds check of all of them.
-    fn fetch<S: Storage + ?Sized>(
+    pub(crate) fn fetch<S: Storage + ?Sized>(
         file: &'a S,
         reqs: &[ExtentReq],
         base: usize,

crates/clawhdf5/tests/contiguous_read_interop.rs

@@ -12,9 +12,11 @@
 
 use std::path::Path;
 use std::process::Command;
+use std::sync::Arc;
 
 use clawhdf5::File;
 use clawhdf5_format::selection::Selection;
+use clawhdf5_format::storage::CountingStorage;
 
 fn python() -> String {
     std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string())
@@ -364,9 +366,31 @@ with h5py.File("{path}", "r") as f:
     ));
 
     let file = File::open(&path).unwrap();
+    // The same file over a storage that serves only range reads: the
+    // selections read only their runs (see `strided_selections_over_storage_
+    // read_few_ranges`) and must give the same bytes.
+    let remote = File::open_storage(Arc::new(CountingStorage::new(
+        std::fs::read(&path).unwrap(),
+    )))
+    .unwrap();
     for (k, (name, code, sel)) in cases.iter().enumerate() {
         let ds = file.dataset(name).unwrap();
         let want_bytes = std::fs::read(dir.path().join(format!("sel_{k}.bin"))).unwrap();
+        let rds = remote.dataset(name).unwrap();
+        assert!(
+            rds.read_selection(sel).unwrap() == want_bytes,
+            "{name} {sel:?}: bytes over a range storage differ from libhdf5's"
+        );
+        assert_eq!(
+            rds.read_f64_selection(sel).unwrap(),
+            ds.read_f64_selection(sel).unwrap(),
+            "{name} {sel:?}: f64 over a range storage"
+        );
+        assert_eq!(
+            rds.read_i32_selection(sel).unwrap(),
+            ds.read_i32_selection(sel).unwrap(),
+            "{name} {sel:?}: i32 over a range storage"
+        );
         let got_bytes = ds.read_selection(sel).unwrap();
         assert!(
             got_bytes == want_bytes,
@@ -401,3 +425,64 @@ with h5py.File("{path}", "r") as f:
         );
     }
 }
+
+/// Over a storage without the file in memory, a selection of a contiguous
+/// dataset reads its runs merged across small gaps: a strided selection is
+/// a few large reads, not one per element (563 200 for the stride-2 case
+/// before), and the values are libhdf5's.
+#[test]
+fn strided_selections_over_storage_read_few_ranges() {
+    skip_if_no_python!();
+    let dir = tempfile::tempdir().unwrap();
+    let path = dir.path().join("contig.h5");
+    write_file(&path);
+    let storage = Arc::new(CountingStorage::new(std::fs::read(&path).unwrap()));
+    let remote = File::open_storage(storage.clone()).unwrap();
+    let local = File::open(&path).unwrap();
+    // f4le_big is 1100 x 1024 f32: 4 KiB rows, 4.4 MB in all.
+    let name = "f4le_big";
+    let every_100th: Vec<Vec<u64>> = (0..1100u64 * 1024)
+        .step_by(100)
+        .map(|i| vec![i / 1024, i % 1024])
+        .collect();
+    let backwards: Vec<Vec<u64>> = every_100th.iter().rev().take(50).cloned().collect();
+    // (selection, most range reads it may take)
+    let cases: Vec<(Selection, u64)> = vec![
+        // Every other element of every row: one read of the whole dataset.
+        (slab(&[(0, 1, 1100, 1), (0, 2, 512, 1)]), 1),
+        // Blocks of 3 every 7 on every other row: rows are 4 KiB apart, so
+        // one read per selected row at most.
+        (slab(&[(0, 2, 550, 1), (1, 7, 146, 3)]), 550),
+        // Every 100th element, in order: 400-byte gaps, one read.
+        (Selection::Points(every_100th), 1),
+        // Points going backwards are not merged.
+        (Selection::Points(backwards), 50),
+        // A column: 4 KiB apart, merged.
+        (slab(&[(0, 1, 1100, 1), (5, 1, 1, 1)]), 1),
+    ];
+    let ds = remote.dataset(name).unwrap();
+    for (sel, most) in cases {
+        storage.reset();
+        let got = ds.read_f32_selection(&sel).unwrap();
+        let reads = storage.reads();
+        assert_eq!(
+            got,
+            local
+                .dataset(name)
+                .unwrap()
+                .read_f32_selection(&sel)
+                .unwrap(),
+            "{sel:?}"
+        );
+        // A few reads of metadata besides the data.
+        assert!(reads <= most + 8, "{sel:?}: {reads} range reads");
+        storage.reset();
+        let bytes = ds.read_selection(&sel).unwrap();
+        assert!(
+            storage.reads() <= most + 8,
+            "{sel:?}: {} reads",
+            storage.reads()
+        );
+        assert_eq!(bytes.len(), got.len() * 4);
+    }
+}