From b086dc3c2bba49fbfa562c8f4ce0ded537a454a5 Mon Sep 17 00:00:00 2001 From: osobh Date: Sat, 26 Sep 2026 18:34:31 -0500 Subject: [PATCH] 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) --- CHANGELOG.md | 5 +- crates/clawhdf5-format/src/gather.rs | 283 +++++++++++++----- crates/clawhdf5-format/src/storage.rs | 2 +- .../clawhdf5/tests/contiguous_read_interop.rs | 85 ++++++ 4 files changed, 294 insertions(+), 81 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 21e857d..3c6a7b6 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/crates/clawhdf5-format/src/gather.rs b/crates/clawhdf5-format/src/gather.rs index 2b8a4ab..08480f1 100644 --- a/crates/clawhdf5-format/src/gather.rs +++ b/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 { @@ -262,12 +262,88 @@ pub(crate) fn gather( 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( file: &S, base: u64, @@ -297,11 +373,15 @@ pub(crate) fn gather_storage( } }; 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 = 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( .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> = 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 = 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> = None; + let mut error: Option = 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) } diff --git a/crates/clawhdf5-format/src/storage.rs b/crates/clawhdf5-format/src/storage.rs index 2b4b9f1..7425296 100644 --- a/crates/clawhdf5-format/src/storage.rs +++ b/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( + pub(crate) fn fetch( file: &'a S, reqs: &[ExtentReq], base: usize, diff --git a/crates/clawhdf5/tests/contiguous_read_interop.rs b/crates/clawhdf5/tests/contiguous_read_interop.rs index 188856e..ed4d4fa 100644 --- a/crates/clawhdf5/tests/contiguous_read_interop.rs +++ b/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> = (0..1100u64 * 1024) + .step_by(100) + .map(|i| vec![i / 1024, i % 1024]) + .collect(); + let backwards: Vec> = 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); + } +}