use integer multiple of shard shape when rechunking in to_zarr (#12106)

d-v-b · dcherian · web-flow · commit 9915a79b0fc0 · 2025-10-29T11:51:50.000-06:00
Co-authored-by: Deepak Cherian &lt;dcherian@users.noreply.github.com&gt;
diff --git a/dask/array/core.py b/dask/array/core.py
@@ -3889,22 +3889,47 @@ def to_zarr(
                     "Cannot store into in memory Zarr Array using "
                     "the distributed scheduler."
                 )
+        zarr_write_chunks = _get_zarr_write_chunks(z)
+        dask_write_chunks = normalize_chunks(
+            chunks="auto",
+            shape=z.shape,
+            dtype=z.dtype,
+            previous_chunks=zarr_write_chunks,
+        )
 
+        for ax, (dw, zw) in enumerate(
+            zip(dask_write_chunks, zarr_write_chunks, strict=True)
+        ):
+            if len(dw) >= 1:
+                nominal_dask_chunk_size = dw[0]
+                if not nominal_dask_chunk_size % zw == 0:
+                    safe_chunk_size = np.prod(zarr_write_chunks) * max(
+                        1, z.dtype.itemsize
+                    )
+                    msg = (
+                        f"The input Dask array will be rechunked along axis {ax} with chunk size "
+                        f"{nominal_dask_chunk_size}, but a chunk size divisible by {zw} is "
+                        f"required for Dask to write safely to the Zarr array {z}. "
+                        "To avoid risk of data loss when writing to this Zarr array, set the "
+                        '"array.chunk-size" configuration parameter to at least the size in'
+                        " bytes of a single on-disk "
+                        f"chunk (or shard) of the Zarr array, which in this case is "
+                        f"{safe_chunk_size} bytes. "
+                        f'E.g., dask.config.set({{"array.chunk-size": {safe_chunk_size}}})'
+                    )
+                    raise PerformanceWarning(msg)
+                    break
         if region is None:
             # Get the appropriate write granularity (shard shape if sharding, else chunk shape)
-            write_chunks = _get_zarr_write_chunks(z)
-            arr = arr.rechunk(write_chunks)
+            arr = arr.rechunk(dask_write_chunks)
             regions = None
         else:
             from dask.array.slicing import new_blockdim, normalize_index
 
-            # For regions, use the appropriate write granularity
-            write_chunks = _get_zarr_write_chunks(z)
-            old_chunks = normalize_chunks(write_chunks, z.shape)
             index = normalize_index(region, z.shape)
             chunks = tuple(
                 tuple(new_blockdim(s, c, r))
-                for s, c, r in zip(z.shape, old_chunks, index)
+                for s, c, r in zip(z.shape, dask_write_chunks, index)
             )
             arr = arr.rechunk(chunks)
             regions = [region]
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
@@ -5,6 +5,7 @@
 import pathlib
 import re
 import xml.etree.ElementTree
+from typing import Literal
 
 import pytest
 
@@ -31,6 +32,7 @@
 from dask.array.core import (
     Array,
     BlockView,
+    PerformanceWarning,
     blockdims_from_blockshape,
     broadcast_chunks,
     broadcast_shapes,
@@ -5122,44 +5124,93 @@ def test_from_array_respects_zarr_shards():
     assert all(c % s == 0 for c, s in zip(dz.chunksize, z.shards))
 
 
-def test_zarr_chunk_shards_mismatch_warns():
+@pytest.mark.parametrize("region_spec", [None, "all", "half"])
+def test_zarr_to_zarr_shards(region_spec: None | Literal["all", "half"]):
     """
     Test that calling to_zarr with a dask array with chunks that do not match the
-    shard shape of the zarr array automatically rechunks to the shard shape to ensure
-    safe writes.
+    shard shape of the zarr array automatically rechunks to a multiple of the
+    shard shape to ensure safe writes.
+
+    This test is parametrized over different regions, because the rechunking logic in
+    to_zarr contains an branch depending on whether a region parameter was specified.
     """
     zarr = pytest.importorskip("zarr", minversion="3.0.0")
-    import numpy as np
 
-    shape = (24,)
-    dask_chunks = (10,)  # Not aligned with shard boundaries
-    zarr_chunk_shape = (4,)  # Inner chunk shape
-    zarr_shard_shape = (12,)  # Shard contains 3 chunks of size 4
+    shape = (100,)
+    dask_chunks = (10,)
+    zarr_chunk_shape = (1,)
+    zarr_shard_shape = (2,)
 
     # Create a dask array with chunks that don't align with shards
     arr = da.arange(shape[0], chunks=dask_chunks)
 
+    # the region parameter we will pass into to_zarr
+    region: tuple[slice, ...] | None
+
+    # The region of the zarr array we will write into
+    sel: tuple[slice, ...]
+
+    if region_spec is None:
+        sel = (slice(None),)
+        region = None
+    elif region_spec == "all":
+        sel = (slice(None),)
+        region = sel
+    else:
+        sel = (slice(shape[0] // 2),)
+        region = sel
+        # crop the source data
+        arr = arr[sel]
+
     # Create a sharded zarr array
     # In Zarr v3: chunks = inner chunk shape, shards = shard shape
     z = zarr.create_array(
-        store={},  # Use in-memory store
+        store={},
         shape=shape,
         chunks=zarr_chunk_shape,
         shards=zarr_shard_shape,
         dtype=arr.dtype,
     )
 
-    # to_zarr should automatically rechunk to shard boundaries
-    result = arr.to_zarr(z, compute=False)
+    # to_zarr should automatically rechunk to a multiple of the shard shape
+    result = arr.to_zarr(z, region=region, compute=False)
 
     # Verify the array was rechunked to the shard shape
-    assert result.chunks == (
-        (zarr_shard_shape[0], zarr_shard_shape[0]),
-    ), f"Expected chunks {((zarr_shard_shape[0], zarr_shard_shape[0]),)}, got {result.chunks}"
+    assert all(c % s == 0 for c, s in zip(result.chunksize, zarr_shard_shape))
 
     # Verify data correctness
     result.compute()
-    assert_eq(z[:], np.arange(shape[0]))
+    assert_eq(z[sel], arr.compute())
+
+
+def test_zarr_risky_shards_warns():
+    """
+    Test that we see a performance warning when dask chooses a chunk size that will cause data loss
+    for zarr arrays.
+    """
+    zarr = pytest.importorskip("zarr", minversion="3.0.0")
+
+    shape = (100,)
+    dask_chunks = (10,)
+    zarr_chunk_shape = (3,)
+    zarr_shard_shape = (6,)
+
+    arr = da.arange(shape[0], chunks=dask_chunks)
+
+    z = zarr.create_array(
+        store={},
+        shape=shape,
+        chunks=zarr_chunk_shape,
+        shards=zarr_shard_shape,
+        dtype=arr.dtype,
+    )
+
+    with dask.config.set({"array.chunk-size": 1}):
+        with pytest.raises(
+            PerformanceWarning,
+            match="The input Dask array will be rechunked along axis",
+        ):
+            arr.to_zarr(z)
 
 
 def test_zarr_nocompute():
