fix: dont override dataset kwargs in loop for sparse (#2445)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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

@@ -804,9 +804,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

@@ -948,6 +948,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(tmp_path: Path):
+    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
 def test_write_auto_sharded_does_not_override(tmp_path: Path):
     z = open_write_group(tmp_path / "arr.zarr", zarr_format=3)