Raise error on concatenation with partial boundary chunks (#594)

* Test concatenation with partial boundary chunk

* Add check on concetenation

* Add a note to custom reader design docs about boundary chunks

* Test roundtripping HDF5 with partial chunk

Co-authored-by: Tom Nicholas <[email protected]>

* Update error message

---------

Co-authored-by: Tom Nicholas <[email protected]>

Author: maxrjones

docs/custom_readers.md

@@ -90,6 +90,10 @@ Generally you want to follow steps like this:
 7. Group `ManifestArrays` up into one or more `ManifestGroup` objects. Ideally you would only have one group, but your format's data model may preclude that. If there is group-level metadata attach this to the `ManifestGroup` object as a `zarr.metadata.GroupMetadata` object. Remember that `ManifestGroups` can contain other groups as well as arrays.
 8. Instantiate the final `ManifestStore` using the top-most `ManifestGroup` and return it.
 
+```{note}
+The [regular chunk grid](https://github.com/zarr-developers/zarr-specs/blob/main/docs/v3/chunk-grids/regular-grid/index.rst) for Zarr V3 data expects that chunks at the border of an array always have the full chunk size, even when the array only covers parts of it. For example, having an array with ``"shape": [30, 30]`` and ``"chunk_shape": [16, 16]``, the chunk ``0,1`` would also contain unused values for the indices ``0-16, 30-31``. If the file format that you are virtualizing does not fill in partial chunks, it is recommended that you raise a `ValueError` until Zarr supports [variable chunk sizes](https://github.com/orgs/zarr-developers/discussions/52).
+```
+
 ### Parsing a pre-existing index file
 
 A custom reader can parse multiple files, perhaps by passing a glob string and looking for expected file naming conventions, or by passing additional reader-specific keyword arguments.

virtualizarr/manifests/array_api.py

@@ -7,6 +7,7 @@
 from .manifest import ChunkManifest
 from .utils import (
     check_combinable_zarr_arrays,
+    check_no_partial_chunks_on_concat_axis,
     check_same_ndims,
     check_same_shapes,
     check_same_shapes_except_on_concat_axis,
@@ -75,7 +76,9 @@ def concatenate(
         axis = axis % first_arr.ndim
 
     arr_shapes = [arr.shape for arr in arrays]
+    arr_chunks = [arr.chunks for arr in arrays]
     check_same_shapes_except_on_concat_axis(arr_shapes, axis)
+    check_no_partial_chunks_on_concat_axis(arr_shapes, arr_chunks, axis)
 
     # find what new array shape must be
     new_length_along_concat_axis = sum([shape[axis] for shape in arr_shapes])

virtualizarr/manifests/utils.py

@@ -157,6 +157,20 @@ def _remove_element_at_position(t: tuple[int, ...], pos: int) -> tuple[int, ...]
     return tuple(new_l)
 
 
+def check_no_partial_chunks_on_concat_axis(
+    shapes: list[tuple[int, ...]], chunks: list[tuple[int, ...]], axis: int
+):
+    """Check that there are no partial chunks along the concatenation axis"""
+    # loop over the arrays to be concatenated
+    for i, (shape, chunk_shape) in enumerate(zip(shapes, chunks)):
+        if shape[axis] % chunk_shape[axis] > 0:
+            raise ValueError(
+                "Cannot concatenate arrays with partial chunks because only regular chunk grids are currently supported. "
+                f"Concat input {i} has array length {shape[axis]} along the concatenation axis which is not "
+                f"evenly divisible by chunk length {chunk_shape[axis]}."
+            )
+
+
 def check_same_shapes_except_on_concat_axis(shapes: list[tuple[int, ...]], axis: int):
     """Check that shapes are compatible for concatenation"""
 

virtualizarr/tests/test_manifests/test_array.py

@@ -385,6 +385,17 @@ def test_refuse_combine(array_v3_metadata):
         with pytest.raises(ValueError, match="inconsistent dtypes"):
             func([marr1, marr2], axis=0)
 
+    metadata_variable_chunk = array_v3_metadata(
+        shape=shape,
+        chunks=(4, 1, 19),
+    )
+    marr = ManifestArray(metadata=metadata_variable_chunk, chunkmanifest=chunkmanifest2)
+    with pytest.raises(
+        ValueError,
+        match="Cannot concatenate arrays with partial chunks because only regular chunk grids are currently supported. Concat input 0 has array length 5 along the concatenation axis which is not evenly divisible by chunk length 4.",
+    ):
+        np.concatenate([marr, marr], axis=0)
+
 
 class TestIndexing:
     @pytest.mark.parametrize("dodgy_indexer", ["string", [], (5, "string")])

virtualizarr/tests/test_readers/test_hdf/test_hdf_integration.py

@@ -10,7 +10,9 @@
     requires_imagecodecs,
     requires_kerchunk,
 )
-from virtualizarr.tests.test_integration import roundtrip_as_in_memory_icechunk
+from virtualizarr.tests.test_integration import (
+    roundtrip_as_in_memory_icechunk,
+)
 
 
 @requires_kerchunk

virtualizarr/tests/test_readers/test_hdf/test_hdf_manifest_store.py

@@ -40,6 +40,22 @@ def test_rountrip_simple_virtualdataset(self, tmpdir, basic_ds):
         )
         xr.testing.assert_allclose(basic_ds, rountripped_ds)
 
+    def test_rountrip_partial_chunk_virtualdataset(self, tmpdir, basic_ds):
+        "Roundtrip a dataset to/from NetCDF with the HDF reader and ManifestStore with a single partial chunk"
+
+        filepath = f"{tmpdir}/basic_ds_roundtrip.nc"
+        encoding = {
+            "temperature": {"chunksizes": (90, 90), "original_shape": (100, 100)}
+        }
+        basic_ds.to_netcdf(filepath, engine="h5netcdf", encoding=encoding)
+        store = HDFVirtualBackend._create_manifest_store(
+            filepath=filepath,
+        )
+        rountripped_ds = xr.open_dataset(
+            store, engine="zarr", consolidated=False, zarr_format=3
+        )
+        xr.testing.assert_allclose(basic_ds, rountripped_ds)
+
     def test_rountrip_simple_virtualdataset_default_store(self, tmpdir, basic_ds):
         "Roundtrip a dataset to/from NetCDF with the HDF reader and ManifestStore"