edit: shrink by visiting the chunks that exist

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]>
2026-09-26 18:50:23 -05:00
parent 159e588550
commit 930921e8cb
3 changed files with 245 additions and 68 deletions
@@ -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}");
+    }
+}