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

Attention: Patch coverage is 93.33333% with 1 line in your changes missing coverage. Please review.

Project coverage is 83.99%. Comparing base (8cc5a18) to head (32e008d).
Report is 41 commits behind head on main.

Files with missing lines	Patch %	Lines
src/anndata/_io/specs/methods.py	93.33%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1624      +/-   ##
==========================================
- Coverage   86.74%   83.99%   -2.75%     
==========================================
  Files          37       40       +3     
  Lines        5988     6031      +43     
==========================================
- Hits         5194     5066     -128     
- Misses        794      965     +171     

Flag	Coverage Δ
	83.99% <93.33%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
src/anndata/_io/specs/methods.py	88.77% <93.33%> (+0.37%)	⬆️

... and 23 files with indirect coverage changes

[ilan-gold]

Typing please!

[ilan-gold]

Suggested change

	elem,
	elem: GroupStorageType,

no?

[ivirshup]

ArrayStorageType I think

[ilan-gold]

Yes!

[ilan-gold]

Suggested change

	elem,
	elem: ArrayStorageType,