Feat: readd manually changes on align_chunks from pydata#10516

Author: lbesnard

xarray/backends/chunks.py

@@ -4,20 +4,18 @@
 
 
 def align_nd_chunks(
-    nd_var_chunks: tuple[tuple[int, ...], ...],
+    nd_v_chunks: tuple[tuple[int, ...], ...],
     nd_backend_chunks: tuple[tuple[int, ...], ...],
 ) -> tuple[tuple[int, ...], ...]:
-    if len(nd_backend_chunks) != len(nd_var_chunks):
+    if len(nd_backend_chunks) != len(nd_v_chunks):
         raise ValueError(
             "The number of dimensions on the backend and the variable must be the same."
         )
 
     nd_aligned_chunks: list[tuple[int, ...]] = []
-    for backend_chunks, var_chunks in zip(
-        nd_backend_chunks, nd_var_chunks, strict=True
-    ):
+    for backend_chunks, v_chunks in zip(nd_backend_chunks, nd_v_chunks, strict=True):
         # Validate that they have the same number of elements
-        if sum(backend_chunks) != sum(var_chunks):
+        if sum(backend_chunks) != sum(v_chunks):
             raise ValueError(
                 "The number of elements in the backend does not "
                 "match the number of elements in the variable. "
@@ -42,39 +40,39 @@ def align_nd_chunks(
             nd_aligned_chunks.append(backend_chunks)
             continue
 
-        if len(var_chunks) == 1:
-            nd_aligned_chunks.append(var_chunks)
+        if len(v_chunks) == 1:
+            nd_aligned_chunks.append(v_chunks)
             continue
 
         # Size of the chunk on the backend
         fixed_chunk = max(backend_chunks)
 
         # The ideal size of the chunks is the maximum of the two; this would avoid
         # that we use more memory than expected
-        max_chunk = max(fixed_chunk, max(var_chunks))
+        max_chunk = max(fixed_chunk, *v_chunks)
 
         # The algorithm assumes that the chunks on this array are aligned except the last one
         # because it can be considered a partial one
         aligned_chunks: list[int] = []
 
         # For simplicity of the algorithm, let's transform the Array chunks in such a way that
         # we remove the partial chunks. To achieve this, we add artificial data to the borders
-        t_var_chunks = list(var_chunks)
-        t_var_chunks[0] += fixed_chunk - backend_chunks[0]
-        t_var_chunks[-1] += fixed_chunk - backend_chunks[-1]
+        t_v_chunks = list(v_chunks)
+        t_v_chunks[0] += fixed_chunk - backend_chunks[0]
+        t_v_chunks[-1] += fixed_chunk - backend_chunks[-1]
 
         # The unfilled_size is the amount of space that has not been filled on the last
         # processed chunk; this is equivalent to the amount of data that would need to be
         # added to a partial Zarr chunk to fill it up to the fixed_chunk size
         unfilled_size = 0
 
-        for var_chunk in t_var_chunks:
+        for v_chunk in t_v_chunks:
             # Ideally, we should try to preserve the original Dask chunks, but this is only
             # possible if the last processed chunk was aligned (unfilled_size == 0)
-            ideal_chunk = var_chunk
+            ideal_chunk = v_chunk
             if unfilled_size:
                 # If that scenario is not possible, the best option is to merge the chunks
-                ideal_chunk = var_chunk + aligned_chunks[-1]
+                ideal_chunk = v_chunk + aligned_chunks[-1]
 
             while ideal_chunk:
                 if not unfilled_size:
@@ -105,27 +103,27 @@ def align_nd_chunks(
             border_size = fixed_chunk - backend_chunks[::order][0]
             aligned_chunks = aligned_chunks[::order]
             aligned_chunks[0] -= border_size
-            t_var_chunks = t_var_chunks[::order]
-            t_var_chunks[0] -= border_size
+            t_v_chunks = t_v_chunks[::order]
+            t_v_chunks[0] -= border_size
             if (
                 len(aligned_chunks) >= 2
                 and aligned_chunks[0] + aligned_chunks[1] <= max_chunk
-                and aligned_chunks[0] != t_var_chunks[0]
+                and aligned_chunks[0] != t_v_chunks[0]
             ):
                 # The artificial data added to the border can introduce inefficient chunks
                 # on the borders, for that reason, we will check if we can merge them or not
                 # Example:
                 # backend_chunks = [6, 6, 1]
-                # var_chunks = [6, 7]
-                # t_var_chunks = [6, 12]
-                # The ideal output should preserve the same var_chunks, but the previous loop
+                # v_chunks = [6, 7]
+                # t_v_chunks = [6, 12]
+                # The ideal output should preserve the same v_chunks, but the previous loop
                 # is going to produce aligned_chunks = [6, 6, 6]
                 # And after removing the artificial data, we will end up with aligned_chunks = [6, 6, 1]
                 # which is not ideal and can be merged into a single chunk
                 aligned_chunks[1] += aligned_chunks[0]
                 aligned_chunks = aligned_chunks[1:]
 
-            t_var_chunks = t_var_chunks[::order]
+            t_v_chunks = t_v_chunks[::order]
             aligned_chunks = aligned_chunks[::order]
 
         nd_aligned_chunks.append(tuple(aligned_chunks))
@@ -141,9 +139,14 @@ def build_grid_chunks(
     if region is None:
         region = slice(0, size)
 
-    region_start = region.start if region.start else 0
+    region_start = region.start or 0
     # Generate the zarr chunks inside the region of this dim
     chunks_on_region = [chunk_size - (region_start % chunk_size)]
+    if chunks_on_region[0] >= size:
+        # This is useful for the scenarios where the chunk_size are bigger
+        # than the variable chunks, which can happens when the user specifies
+        # the enc_chunks manually.
+        return (size,)
     chunks_on_region.extend([chunk_size] * ((size - chunks_on_region[0]) // chunk_size))
     if (size - chunks_on_region[0]) % chunk_size != 0:
         chunks_on_region.append((size - chunks_on_region[0]) % chunk_size)
@@ -155,45 +158,45 @@ def grid_rechunk(
     enc_chunks: tuple[int, ...],
     region: tuple[slice, ...],
 ) -> Variable:
-    nd_var_chunks = v.chunks
-    if not nd_var_chunks:
+    nd_v_chunks = v.chunks
+    if not nd_v_chunks:
         return v
 
     nd_grid_chunks = tuple(
         build_grid_chunks(
-            sum(var_chunks),
+            v_size,
             region=interval,
             chunk_size=chunk_size,
         )
-        for var_chunks, chunk_size, interval in zip(
-            nd_var_chunks, enc_chunks, region, strict=True
+        for v_size, chunk_size, interval in zip(
+            v.shape, enc_chunks, region, strict=True
         )
     )
 
     nd_aligned_chunks = align_nd_chunks(
-        nd_var_chunks=nd_var_chunks,
+        nd_v_chunks=nd_v_chunks,
         nd_backend_chunks=nd_grid_chunks,
     )
     v = v.chunk(dict(zip(v.dims, nd_aligned_chunks, strict=True)))
     return v
 
 
 def validate_grid_chunks_alignment(
-    nd_var_chunks: tuple[tuple[int, ...], ...] | None,
+    nd_v_chunks: tuple[tuple[int, ...], ...] | None,
     enc_chunks: tuple[int, ...],
     backend_shape: tuple[int, ...],
     region: tuple[slice, ...],
     allow_partial_chunks: bool,
     name: str,
 ):
-    if nd_var_chunks is None:
+    if nd_v_chunks is None:
         return
     base_error = (
         "Specified Zarr chunks encoding['chunks']={enc_chunks!r} for "
         "variable named {name!r} would overlap multiple Dask chunks. "
-        "Check the chunk at position {var_chunk_pos}, which has a size of "
-        "{var_chunk_size} on dimension {dim_i}. It is unaligned with "
-        "backend chunks of size {chunk_size} in region {region}. "
+        "Please check the Dask chunks at position {v_chunk_pos} and "
+        "{v_chunk_pos_next}, on axis {axis}, they are overlapped "
+        "on the same Zarr chunk in the region {region}. "
         "Writing this array in parallel with Dask could lead to corrupted data. "
         "To resolve this issue, consider one of the following options: "
         "- Rechunk the array using `chunk()`. "
@@ -202,58 +205,61 @@ def validate_grid_chunks_alignment(
         "- Enable automatic chunks alignment with `align_chunks=True`."
     )
 
-    for dim_i, chunk_size, var_chunks, interval, size in zip(
+    for axis, chunk_size, v_chunks, interval, size in zip(
         range(len(enc_chunks)),
         enc_chunks,
-        nd_var_chunks,
+        nd_v_chunks,
         region,
         backend_shape,
         strict=True,
     ):
-        for i, chunk in enumerate(var_chunks[1:-1]):
+        for i, chunk in enumerate(v_chunks[1:-1]):
             if chunk % chunk_size:
                 raise ValueError(
                     base_error.format(
-                        var_chunk_pos=i + 1,
-                        var_chunk_size=chunk,
+                        v_chunk_pos=i + 1,
+                        v_chunk_pos_next=i + 2,
+                        v_chunk_size=chunk,
+                        axis=axis,
                         name=name,
-                        dim_i=dim_i,
                         chunk_size=chunk_size,
                         region=interval,
                         enc_chunks=enc_chunks,
                     )
                 )
 
-        interval_start = interval.start if interval.start else 0
+        interval_start = interval.start or 0
 
-        if len(var_chunks) > 1:
+        if len(v_chunks) > 1:
             # The first border size is the amount of data that needs to be updated on the
             # first chunk taking into account the region slice.
             first_border_size = chunk_size
             if allow_partial_chunks:
                 first_border_size = chunk_size - interval_start % chunk_size
 
-            if (var_chunks[0] - first_border_size) % chunk_size:
+            if (v_chunks[0] - first_border_size) % chunk_size:
                 raise ValueError(
                     base_error.format(
-                        var_chunk_pos=0,
-                        var_chunk_size=var_chunks[0],
+                        v_chunk_pos=0,
+                        v_chunk_pos_next=0,
+                        v_chunk_size=v_chunks[0],
+                        axis=axis,
                         name=name,
-                        dim_i=dim_i,
                         chunk_size=chunk_size,
                         region=interval,
                         enc_chunks=enc_chunks,
                     )
                 )
 
         if not allow_partial_chunks:
-            region_stop = interval.stop if interval.stop else size
+            region_stop = interval.stop or size
 
             error_on_last_chunk = base_error.format(
-                var_chunk_pos=len(var_chunks) - 1,
-                var_chunk_size=var_chunks[-1],
+                v_chunk_pos=len(v_chunks) - 1,
+                v_chunk_pos_next=len(v_chunks) - 1,
+                v_chunk_size=v_chunks[-1],
+                axis=axis,
                 name=name,
-                dim_i=dim_i,
                 chunk_size=chunk_size,
                 region=interval,
                 enc_chunks=enc_chunks,
@@ -267,7 +273,7 @@ def validate_grid_chunks_alignment(
                 # If the region is covering the last chunk then check
                 # if the reminder with the default chunk size
                 # is equal to the size of the last chunk
-                if var_chunks[-1] % chunk_size != size % chunk_size:
+                if v_chunks[-1] % chunk_size != size % chunk_size:
                     raise ValueError(error_on_last_chunk)
-            elif var_chunks[-1] % chunk_size:
+            elif v_chunks[-1] % chunk_size:
                 raise ValueError(error_on_last_chunk)

xarray/backends/zarr.py

@@ -1186,7 +1186,7 @@ def set_variables(
                 # threads
                 shape = zarr_shape if zarr_shape else v.shape
                 validate_grid_chunks_alignment(
-                    nd_var_chunks=v.chunks,
+                    nd_v_chunks=v.chunks,
                     enc_chunks=encoding["chunks"],
                     region=region,
                     allow_partial_chunks=self._mode != "r+",

xarray/core/dataset.py

@@ -2303,6 +2303,7 @@ def to_zarr(
             append_dim=append_dim,
             region=region,
             safe_chunks=safe_chunks,
+            align_chunks=align_chunks,
             zarr_version=zarr_version,
             zarr_format=zarr_format,
             write_empty_chunks=write_empty_chunks,