quantumclaw/clawhdf5
3
1
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases 1 Wiki Activity

edit: shrink by visiting the chunks that exist

Browse Source 
prune_plan stored one Vec<u64> for every chunk coordinate of the region a
shrink cuts off, existing or not, so a sparse dataset exhausted memory
(about 62 bytes per coordinate; (4, 2e7) with chunks (1, 1) took 2.5 GB,
larger extents never finished). It now places each existing chunk in
H5D__chunk_prune_by_extent's walk (its pass, then its coordinates) and
sorts, which gives the same chunks, order and actions in memory and time
proportional to the chunks that exist.

A unit test checks the plan against the full walk (kept as the test's
reference) for 3000 random extents and chunk subsets. The interop test
shrinks a (4, 10^12) dataset with chunks (1, 1) and 9 chunks (v1 and v2
B-tree): 0.56 s and 43 MB peak; the old code aborted on allocation under an
8 GB limit.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
osobh
2026-09-26 18:50:23 -05:00
co-authored by Claude Opus 5.5
parent 159e588550
commit 930921e8cb
3 changed files with 245 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/clawhdf5-tools/tests/edit_coverage_interop.rs
+62
View File
@@ -1566,3 +1566,65 @@ fn resize_clawhdf5_written_datasets() {
}
check_tools(&path, true);
}
/// Shrinking a huge, sparse dataset costs time and memory in the chunks
/// that exist, not in the chunk coordinates cut off. The editor once stored
/// every coordinate of the cut-off region (about 62 bytes each), so this
/// 2 x 10^12-coordinate shrink ran out of memory.
#[test]
fn shrinking_a_huge_sparse_dataset_is_bounded() {
if !tools_ok() {
return;
}
let dir = tmpdir();
const N: u64 = 1_000_000_000_000;
for libver in ["earliest", "v110"] {
let path = dir.path().join(format!("sparse_{libver}.h5"));
py(&format!(
"import h5py\n\
with h5py.File({p:?}, 'w', libver=({libver:?}, 'latest')) as f:\n\
\x20 d = f.create_dataset('b', shape=(4, {N}), maxshape=(None, None), chunks=(1, 1), dtype='<i4')\n\
\x20 d[0, 0:5] = [1, 2, 3, 4, 5]\n\
\x20 d[0, 900000000000] = 6\n\
\x20 d[1, {N} - 1] = 7\n\
\x20 d[2, 7] = 8\n\
\x20 d[3, 500000000000] = 9\n\
\x20 assert d.id.get_num_chunks() == 9\n",
p = path.to_str().unwrap(),
));
let start = std::time::Instant::now();
let mut ed = FileEditor::open(&path).unwrap();
ed.resize("b", &[2, 600_000_000_000]).unwrap();
drop(ed);
let took = start.elapsed();
assert!(
took < std::time::Duration::from_secs(30),
"{libver}: shrink took {took:?}"
);
py(&format!(
"import h5py\n\
with h5py.File({p:?}, 'r') as f:\n\
\x20 d = f['b']\n\
\x20 assert d.shape == (2, 600000000000), d.shape\n\
\x20 assert d.id.get_num_chunks() == 5, d.id.get_num_chunks()\n\
\x20 assert list(d[0, 0:6]) == [1, 2, 3, 4, 5, 0], d[0, 0:6]\n\
\x20 assert d[1, 599999999999] == 0\n\
with h5py.File({p:?}, 'r+') as f:\n\
\x20 d = f['b']\n\
\x20 d.resize((4, {N}))\n\
\x20 assert d[0, 900000000000] == 0 and d[1, {N} - 1] == 0\n\
\x20 assert d[2, 7] == 0 and d[3, 500000000000] == 0\n",
p = path.to_str().unwrap(),
));
let f = File::open(&path).unwrap();
let ds = f.dataset("b").unwrap();
let got = ds.read_selection(&block(&[0, 0], &[2, 6])).unwrap();
let got: Vec<i32> = got
.as_chunks::<4>()
.0
.iter()
.map(|c| i32::from_le_bytes(*c))
.collect();
assert_eq!(got, [1, 2, 3, 4, 5, 0, 0, 0, 0, 0, 0, 0], "{libver}");
}
}
crates/clawhdf5/src/edit/mod.rs
+179 -67
View File
@@ -1330,16 +1330,31 @@ enum Prune {
Remove,
}
/// The chunks `H5D__chunk_prune_by_extent` visits when the extent shrinks
/// from `old` to `new`, in its order.
fn prune_plan(old: &[u64], new: &[u64], cd: &[u64]) -> Vec<(Vec<u64>, Prune)> {
/// The chunks of `existing` (keyed by scaled coordinates) that
/// `H5D__chunk_prune_by_extent` visits when the extent shrinks from `old`
/// to `new`, in its order, with what it does to each.
///
/// libhdf5 walks every chunk coordinate of the region cut off, one pass per
/// shrunk dimension `op` in order: the coordinates with `scaled[op]` at or
/// past the new extent's chunk and every earlier shrunk dimension below its
/// own, row-major. It looks each one up in the index and stores nothing,
/// so a sparse dataset costs it time, not memory. Only chunks that exist
/// can be acted on, so this places each existing chunk in that walk (its
/// pass, then its coordinates) and sorts: the same chunks in the same
/// order, in memory and time proportional to the chunks that exist.
fn prune_plan<'c>(
old: &[u64],
new: &[u64],
cd: &[u64],
existing: &'c HashMap<Vec<u64>, ChunkInfo>,
) -> Vec<(&'c [u64], &'c ChunkInfo, Prune)> {
let rank = new.len();
let mut out = Vec::new();
if old.contains(&0) {
return out;
if old.contains(&0) || cd.contains(&0) {
return Vec::new();
}
let shrunk: Vec<bool> = (0..rank).map(|d| new[d] < old[d]).collect();
let mut max_mod: Vec<u64> = (0..rank).map(|d| (old[d] - 1) / cd[d]).collect();
let max_mod: Vec<u64> = (0..rank).map(|d| (old[d] - 1) / cd[d]).collect();
// The last chunk index still (partly) inside the new extent; -1 when
// the dimension shrank to nothing.
let max_fill: Vec<i64> = (0..rank)
.map(|d| {
if new[d] == 0 {
@@ -1350,60 +1365,25 @@ fn prune_plan(old: &[u64], new: &[u64], cd: &[u64]) -> Vec<(Vec<u64>, Prune)> {
})
.collect();
let min_mod: Vec<u64> = (0..rank).map(|d| new[d] / cd[d]).collect();
let fill_dim: Vec<bool> = (0..rank)
.map(|d| shrunk[d] && min_mod[d] as i64 == max_fill[d])
.collect();
for op in 0..rank {
if !shrunk[op] {
continue;
}
let mut scaled = vec![0u64; rank];
scaled[op] = min_mod[op];
let mut outside: Vec<bool> = (0..rank).map(|u| scaled[u] as i64 > max_fill[u]).collect();
let mut n_out = outside.iter().filter(|&&o| o).count();
loop {
if n_out == 0 {
out.push((scaled.clone(), Prune::Fill));
let mut out: Vec<(usize, &[u64], &ChunkInfo, Prune)> = existing
.iter()
.filter(|(s, _)| s.len() == rank && s.iter().zip(&max_mod).all(|(c, m)| c <= m))
.filter_map(|(s, info)| {
// The first shrunk dimension whose pass reaches this chunk; the
// passes before it covered only coordinates below their minimum.
let op = (0..rank).find(|&d| new[d] < old[d] && s[d] >= min_mod[d])?;
// Chunks with any coordinate past the last partly kept chunk go;
// the others lose the part outside the new extent.
let what = if (0..rank).all(|u| s[u] as i64 <= max_fill[u]) {
Prune::Fill
} else {
out.push((scaled.clone(), Prune::Remove));
}
let mut carry = true;
for i in (0..rank).rev() {
scaled[i] += 1;
if scaled[i] > max_mod[i] {
if i == op {
scaled[i] = min_mod[i];
if outside[i] && fill_dim[i] {
outside[i] = false;
n_out -= 1;
}
} else {
scaled[i] = 0;
if outside[i] && max_fill[i] >= 0 {
outside[i] = false;
n_out -= 1;
}
}
} else {
if !outside[i] && scaled[i] as i64 > max_fill[i] {
outside[i] = true;
n_out += 1;
}
carry = false;
break;
}
}
if carry {
break;
}
}
if min_mod[op] == 0 {
// Every chunk was visited (the dimension shrank to nothing).
break;
}
max_mod[op] = min_mod[op] - 1;
}
out
Prune::Remove
};
Some((op, s.as_slice(), info, what))
})
.collect();
out.sort_unstable_by(|a, b| (a.0, a.1).cmp(&(b.0, b.1)));
out.into_iter().map(|(_, s, i, w)| (s, i, w)).collect()
}
/// The chunk work of a resize: under early allocation, allocate and fill
@@ -1468,12 +1448,9 @@ fn resize_chunks(
})?;
}
if shrink && !existing.is_empty() {
for (scaled, what) in prune_plan(old, new, &cd) {
let Some(info) = existing.get(&scaled) else {
continue;
};
for (scaled, info, what) in prune_plan(old, new, &cd, &existing) {
match what {
Prune::Remove => ce.remove(img, &scaled, info)?,
Prune::Remove => ce.remove(img, scaled, info)?,
Prune::Fill => {
let mut buf = decode_chunk(img_read(f, info)?, t, info, chunk_bytes)?;
// Keep [0, count) in each dimension; fill the rest.
@@ -1495,7 +1472,7 @@ fn resize_chunks(
buf[at..at + es].copy_from_slice(&t.fill);
}
}
store_chunk(img, &mut ce, Some(info), &scaled, buf)?;
store_chunk(img, &mut ce, Some(info), scaled, buf)?;
}
}
}
@@ -1610,6 +1587,141 @@ fn decode_chunk(
#[cfg(test)]
mod tests {
use super::{Prune, prune_plan};
use clawhdf5_format::chunked_read::ChunkInfo;
use std::collections::HashMap;
/// `H5D__chunk_prune_by_extent`'s walk over every chunk coordinate of the
/// region cut off, stored (the implementation before `prune_plan`
/// visited only existing chunks).
fn prune_walk(old: &[u64], new: &[u64], cd: &[u64]) -> Vec<(Vec<u64>, Prune)> {
let rank = new.len();
let mut out = Vec::new();
if old.contains(&0) {
return out;
}
let shrunk: Vec<bool> = (0..rank).map(|d| new[d] < old[d]).collect();
let mut max_mod: Vec<u64> = (0..rank).map(|d| (old[d] - 1) / cd[d]).collect();
let max_fill: Vec<i64> = (0..rank)
.map(|d| {
if new[d] == 0 {
-1
} else {
((new[d].min(old[d]) - 1) / cd[d]) as i64
}
})
.collect();
let min_mod: Vec<u64> = (0..rank).map(|d| new[d] / cd[d]).collect();
let fill_dim: Vec<bool> = (0..rank)
.map(|d| shrunk[d] && min_mod[d] as i64 == max_fill[d])
.collect();
for op in 0..rank {
if !shrunk[op] {
continue;
}
let mut scaled = vec![0u64; rank];
scaled[op] = min_mod[op];
let mut outside: Vec<bool> =
(0..rank).map(|u| scaled[u] as i64 > max_fill[u]).collect();
let mut n_out = outside.iter().filter(|&&o| o).count();
loop {
if n_out == 0 {
out.push((scaled.clone(), Prune::Fill));
} else {
out.push((scaled.clone(), Prune::Remove));
}
let mut carry = true;
for i in (0..rank).rev() {
scaled[i] += 1;
if scaled[i] > max_mod[i] {
if i == op {
scaled[i] = min_mod[i];
if outside[i] && fill_dim[i] {
outside[i] = false;
n_out -= 1;
}
} else {
scaled[i] = 0;
if outside[i] && max_fill[i] >= 0 {
outside[i] = false;
n_out -= 1;
}
}
} else {
if !outside[i] && scaled[i] as i64 > max_fill[i] {
outside[i] = true;
n_out += 1;
}
carry = false;
break;
}
}
if carry {
break;
}
}
if min_mod[op] == 0 {
// Every chunk was visited (the dimension shrank to nothing).
break;
}
max_mod[op] = min_mod[op] - 1;
}
out
}
/// `prune_plan` over a random subset of chunks gives the chunks of the
/// full walk that exist, in its order and with its actions.
#[test]
fn prune_plan_follows_the_full_walk() {
let mut x: u64 = 0x1234_5678;
let mut rnd = |n: u64| {
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
x % n
};
for _ in 0..3000 {
let rank = 1 + rnd(3) as usize;
let cd: Vec<u64> = (0..rank).map(|_| 1 + rnd(4)).collect();
let old: Vec<u64> = (0..rank).map(|_| 1 + rnd(14)).collect();
let new: Vec<u64> = old
.iter()
.map(|&o| if rnd(2) == 0 { o } else { rnd(o + 1) })
.collect();
let grid: Vec<u64> = (0..rank).map(|d| old[d].div_ceil(cd[d])).collect();
let mut existing = HashMap::new();
let total: u64 = grid.iter().product();
for flat in 0..total {
if rnd(3) == 0 {
continue;
}
let mut r = flat;
let mut s = vec![0; rank];
for d in (0..rank).rev() {
s[d] = r % grid[d];
r /= grid[d];
}
let info = ChunkInfo {
chunk_size: 1,
filter_mask: 0,
offsets: s.iter().zip(&cd).map(|(a, b)| a * b).collect(),
address: flat,
};
existing.insert(s, info);
}
let want: Vec<(Vec<u64>, bool)> = prune_walk(&old, &new, &cd)
.into_iter()
.filter(|(s, _)| existing.contains_key(s))
.map(|(s, w)| (s, matches!(w, Prune::Fill)))
.collect();
let got: Vec<(Vec<u64>, bool)> = prune_plan(&old, &new, &cd, &existing)
.into_iter()
.map(|(s, _, w)| (s.to_vec(), matches!(w, Prune::Fill)))
.collect();
assert_eq!(got, want, "old {old:?} new {new:?} chunks {cd:?}");
}
}
use std::cell::Cell;
use std::path::Path;