Chunked writing of h5py.Dataset and zarr.Array

[ivirshup]

• Closes Writing a h5py.Dataset loads the whole thing into memory #1623
• Tests added
  • This might already be good on tests, but I should check
• Release note added (or unnecessary)
• Add benchmarks

This PR fixes #1623 by writing backed dense arrays in chunks.

Very open to feedback on the logic of how chunking pattern of writes is selected. Maybe we should prioritize the chunking of the destination array over the chunking of the source array?

cc: @ebezzi

Some proof it works:

%load_ext memory_profiler

import h5py
from anndata.experimental import write_elem
import numpy as np

f = h5py.File("tmp.h5", "w")
X = np.ones((10_000, 10_000))

%memit write_elem(f, "X", X)
# peak memory: 945.67 MiB, increment: 0.56 MiB

%memit write_elem(f, "X2", f["X"])
peak memory: 1047.00 MiB, increment: 101.12 MiB

%memit write_elem(f, "X3", f["X"], dataset_kwargs={"compression":"gzip"})
peak memory: 1068.03 MiB, increment: 6.14 MiB

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 84.63%. Comparing base (df213f6) to head (c6afa80).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1624      +/-   ##
==========================================
- Coverage   87.06%   84.63%   -2.44%     
==========================================
  Files          40       40              
  Lines        6101     6116      +15     
==========================================
- Hits         5312     5176     -136     
- Misses        789      940     +151     

Files with missing lines	Coverage Δ
src/anndata/_io/specs/methods.py	88.54% <100.00%> (-0.06%)	⬇️

... and 7 files with indirect coverage changes

[flying-sheep]

should any of this be configurable?

[ilan-gold]

As things stand, the value is currently dependent on both the shape and an arbitrary cutoff in max....so Given the current implementation, we could make 2-3 things configurable which seems like overkill. Perhaps just n_rows should be a setting with 1000 as the default?

[flying-sheep]

We already have

anndata/src/anndata/_io/h5ad.py

Line 176 in df213f6

chunk_size: int = 6000, # TODO, probably make this 2d chunks

, documented as “Used only when loading sparse dataset that is stored as dense.”

Also

anndata/src/anndata/_io/zarr.py

Line 30 in df213f6

chunks: tuple[int, ...] | None = None,

Let’s not create multiple different/incompatible conventions / features under the same name.

[ilan-gold]

So this is an argument for not calling it chunk_size? I wasn't proposing literally calling it n_rows but just that variable being the settings as opposed to entry_chunk_size or the max value

[flying-sheep]

It’s an argument for keeping our terminology consistent when we get around to make this configurable. But we can also not do that for now.