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]>
This commit is contained in:
osobh
2026-09-26 08:57:16 -05:00
co-authored by Claude Opus 5.5
parent 8c51b05b9c
commit b43bd2e67f
5 changed files with 259 additions and 36 deletions
@@ -531,3 +531,35 @@ def test_compound_padding_bytes_match_h5py(pair):
assert ours[name][key].tobytes() == theirs[name][key].tobytes(), f"{name}[{key!r}]"
for key in [(slice(None), [0, 2]), ([0, 2, 3], slice(None)), ([1, 3], 1)]:
assert ours["cmp/nested_2d"][key].tobytes() == theirs["cmp/nested_2d"][key].tobytes(), key
def test_a_long_index_list_decodes_each_chunk_once(h5py, tmp_path):
"""An index list is read one group of chunks at a time, not one
hyperslab per run of indices: 5000 runs over 20 gzip chunks used to
decode the chunks 5000 times (8 s, against h5py's 0.014 s)."""
import time
path = str(tmp_path / "long_list.h5")
data = np.arange(200000, dtype="<f8")
grid = np.arange(400 * 3000, dtype="<i4").reshape(400, 3000)
with h5py.File(path, "w") as f:
f.create_dataset("d", data=data, chunks=(10000,), compression="gzip")
f.create_dataset("grid", data=grid, chunks=(50, 100), compression="gzip")
f.create_dataset("flat", data=data) # contiguous
rng = np.random.default_rng(3)
cases = [
("d", list(range(0, 200000, 40))),
("d", sorted(rng.choice(200000, 3000, replace=False).tolist())),
("flat", list(range(0, 200000, 40))),
("flat", [0, 7, 199999]),
("grid", (slice(None), list(range(0, 3000, 3)))),
("grid", (sorted(rng.choice(400, 150, replace=False).tolist()), slice(5, 2900, 7))),
("grid", (7, [0, 1, 2, 2000, 2999])),
]
with h5py.File(path, "r") as theirs, clawhdf5.File(path, "r") as ours:
for name, key in cases:
t0 = time.perf_counter()
got = ours[name][key]
took = time.perf_counter() - t0
assert_same(got, theirs[name][key], f"{name}[{len(key)}-key]")
assert took < 2.0, f"{name}: {took:.2f} s"