Add open support

Author: maxrjones

src/zarr/core/array.py

@@ -675,13 +675,30 @@ async def _create(
 
         if chunks is not None and chunk_shape is not None:
             raise ValueError("Only one of chunk_shape or chunks can be provided.")
-        item_size = 1
-        if isinstance(dtype_parsed, HasItemSize):
-            item_size = dtype_parsed.item_size
-        if chunks:
-            _chunks = normalize_chunks(chunks, shape, item_size)
+
+        # detect rectilinear (dask-style) chunks before normalize_chunks flattens them
+        from zarr.core.chunk_grids import _is_rectilinear_chunks
+
+        _raw_chunks = chunks if chunks is not None else chunk_shape
+        _rectilinear_chunk_grid: ChunkGridMetadata | None = None
+        _chunks: tuple[int, ...] | None = None
+        if _is_rectilinear_chunks(_raw_chunks):
+            from zarr.core.metadata.v3 import (
+                RectilinearChunkGrid as RectilinearChunkGridMeta,
+            )
+
+            _rectilinear_chunk_grid = RectilinearChunkGridMeta(
+                chunk_shapes=tuple(tuple(c) for c in _raw_chunks)
+            )
         else:
-            _chunks = normalize_chunks(chunk_shape, shape, item_size)
+            item_size = 1
+            if isinstance(dtype_parsed, HasItemSize):
+                item_size = dtype_parsed.item_size
+            if chunks:
+                _chunks = normalize_chunks(chunks, shape, item_size)
+            else:
+                _chunks = normalize_chunks(chunk_shape, shape, item_size)
+
         config_parsed = parse_array_config(config)
 
         result: AnyAsyncArray
@@ -714,6 +731,7 @@ async def _create(
                 attributes=attributes,
                 overwrite=overwrite,
                 config=config_parsed,
+                chunk_grid=_rectilinear_chunk_grid,
             )
         elif zarr_format == 2:
             if codecs is not None:
@@ -726,6 +744,9 @@ async def _create(
                 )
             if dimension_names is not None:
                 raise ValueError("dimension_names cannot be used for arrays with zarr_format 2.")
+            if _rectilinear_chunk_grid is not None:
+                raise ValueError("Zarr format 2 does not support rectilinear chunk grids.")
+            assert _chunks is not None
 
             if order is None:
                 order_parsed = config_parsed.order
@@ -760,7 +781,7 @@ async def _create(
     def _create_metadata_v3(
         shape: ShapeLike,
         dtype: ZDType[TBaseDType, TBaseScalar],
-        chunk_shape: tuple[int, ...],
+        chunk_shape: tuple[int, ...] | None = None,
         fill_value: Any | None = DEFAULT_FILL_VALUE,
         chunk_key_encoding: ChunkKeyEncodingLike | None = None,
         codecs: Iterable[Codec | dict[str, JSON]] | None = None,
@@ -771,8 +792,12 @@ def _create_metadata_v3(
         """
         Create an instance of ArrayV3Metadata.
 
-        If `chunk_grid` is provided, it takes precedence over `chunk_shape`.
+        Exactly one of ``chunk_shape`` or ``chunk_grid`` must be provided.
         """
+        if chunk_shape is not None and chunk_grid is not None:
+            raise ValueError("Only one of chunk_shape or chunk_grid can be provided.")
+        if chunk_shape is None and chunk_grid is None:
+            raise ValueError("One of chunk_shape or chunk_grid must be provided.")
         filters: tuple[ArrayArrayCodec, ...]
         compressors: tuple[BytesBytesCodec, ...]
 
@@ -799,11 +824,12 @@ def _create_metadata_v3(
         else:
             fill_value_parsed = fill_value
 
-        chunk_grid_meta = (
-            chunk_grid
-            if chunk_grid is not None
-            else RegularChunkGrid(chunk_shape=parse_shapelike(chunk_shape))
-        )
+        chunk_grid_meta: ChunkGridMetadata
+        if chunk_grid is not None:
+            chunk_grid_meta = chunk_grid
+        else:
+            assert chunk_shape is not None  # validated above
+            chunk_grid_meta = RegularChunkGrid(chunk_shape=parse_shapelike(chunk_shape))
         return ArrayV3Metadata(
             shape=shape,
             data_type=dtype,
@@ -822,7 +848,7 @@ async def _create_v3(
         *,
         shape: ShapeLike,
         dtype: ZDType[TBaseDType, TBaseScalar],
-        chunk_shape: tuple[int, ...],
+        chunk_shape: tuple[int, ...] | None = None,
         config: ArrayConfig,
         fill_value: Any | None = DEFAULT_FILL_VALUE,
         chunk_key_encoding: (
@@ -835,7 +861,12 @@ async def _create_v3(
         dimension_names: DimensionNamesLike = None,
         attributes: dict[str, JSON] | None = None,
         overwrite: bool = False,
+        chunk_grid: ChunkGridMetadata | None = None,
     ) -> AsyncArrayV3:
+        if chunk_shape is not None and chunk_grid is not None:
+            raise ValueError("Only one of chunk_shape or chunk_grid can be provided.")
+        if chunk_shape is None and chunk_grid is None:
+            raise ValueError("One of chunk_shape or chunk_grid must be provided.")
         if overwrite:
             if store_path.store.supports_deletes:
                 await store_path.delete_dir()
@@ -860,6 +891,7 @@ async def _create_v3(
             codecs=codecs,
             dimension_names=dimension_names,
             attributes=attributes,
+            chunk_grid=chunk_grid,
         )
 
         array = cls(metadata=metadata, store_path=store_path, config=config)
@@ -4902,7 +4934,7 @@ async def init_array(
             shape=shape_parsed,
             dtype=zdtype,
             fill_value=fill_value,
-            chunk_shape=chunks_out,
+            chunk_shape=chunks_out if rectilinear_meta is None else None,
             chunk_key_encoding=chunk_key_encoding_parsed,
             codecs=codecs_out,
             dimension_names=dimension_names,

tests/test_api.py

@@ -280,6 +280,19 @@ async def test_open_array(memory_store: MemoryStore, zarr_format: ZarrFormat) ->
         zarr.api.synchronous.open(store="doesnotexist", mode="r", zarr_format=zarr_format)
 
 
+def test_open_array_rectilinear_chunks(tmp_path: Path) -> None:
+    """zarr.open with rectilinear (dask-style) chunks preserves the chunk grid."""
+    from zarr.core.metadata.v3 import RectilinearChunkGrid
+
+    chunks = ((3, 3, 4), (5, 5))
+    with zarr.config.set({"array.rectilinear_chunks": True}):
+        z = zarr.open(store=tmp_path, shape=(10, 10), dtype="float64", chunks=chunks, mode="w")
+    assert isinstance(z, Array)
+    assert z.shape == (10, 10)
+    assert isinstance(z.metadata.chunk_grid, RectilinearChunkGrid)
+    assert z.read_chunk_sizes == ((3, 3, 4), (5, 5))
+
+
 @pytest.mark.asyncio
 async def test_async_array_open_array_not_found() -> None:
     """Test that AsyncArray.open raises ArrayNotFoundError when array doesn't exist"""