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:
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user