Backport PR #2445: fix: dont override dataset kwargs in loop for sparse

Author: ilan-gold

docs/release-notes/2445.fix.md

@@ -0,0 +1 @@
+Make sure `indices`, `data`, and `indptr` have `zarr.config.set({'array.target_shard_size_bytes'})` applied instead of having one override the other's setting when writing sparse matrices. {user}`ilan-gold`

src/anndata/_io/specs/methods.py

@@ -838,9 +838,9 @@ def write_sparse_compressed(
         else:
             with zarr_v3_sharding(
                 dataset_kwargs, format=f.metadata.zarr_format
-            ) as dataset_kwargs:
+            ) as dataset_kwargs_local:
                 arr = g.create_array(
-                    attr_name, shape=attr.shape, dtype=dtype, **dataset_kwargs
+                    attr_name, shape=attr.shape, dtype=dtype, **dataset_kwargs_local
                 )
             # see https://github.com/zarr-developers/zarr-python/discussions/2712
             arr[...] = attr[...]

tests/test_io_elementwise.py

@@ -969,6 +969,27 @@ def test_write_auto_sharded_default_warns(tmp_path: Path):
         adata.write_zarr(path)
 
 
+@pytest.mark.zarr_io
+@pytest.mark.skipif(
+    Version(version("zarr")) < Version("3.1.4"),
+    reason="autosharding with chosen size was not available",
+)
+def test_write_auto_sharded_size_sparse():
+    path = "memory://check_shards.zarr"
+    z = zarr.open(path)
+    mat = sparse.random(
+        1000, 1000, density=0.5, format="csr", random_state=np.random.default_rng(42)
+    )
+    ad.io.write_elem(z, "two_shards_per_sub_element", mat)
+    # i.e., there are at most two shards since one shard will contain two chunks,
+    # and the other the last elements, since the target size is 1GB uncompressed.
+    for sub_element in ["indices", "data", "indptr"]:
+        assert (
+            z["two_shards_per_sub_element"][sub_element].shape[0]
+            / z["two_shards_per_sub_element"][sub_element].shards[0]
+        ) < 2, sub_element
+
+
 @pytest.mark.zarr_io
 @pytest.mark.skipif(is_zarr_v2(), reason="auto sharding is allowed only for zarr v3.")
 def test_write_auto_sharded_does_not_override(tmp_path: Path):