rename (Async)Array.create to _create

Author: normanrz

src/zarr/api/asynchronous.py

@@ -425,7 +425,7 @@ async def save_array(
     shape = arr.shape
     chunks = getattr(arr, "chunks", None)  # for array-likes with chunks attribute
     overwrite = kwargs.pop("overwrite", None) or _infer_overwrite(mode)
-    new = await AsyncArray.create(
+    new = await AsyncArray._create(
         store_path,
         zarr_format=zarr_format,
         shape=shape,
@@ -1041,7 +1041,7 @@ async def create(
 
     config_parsed = ArrayConfig.from_dict(config_dict)
 
-    return await AsyncArray.create(
+    return await AsyncArray._create(
         store_path,
         shape=shape,
         chunks=chunks,

src/zarr/core/array.py

@@ -7,7 +7,7 @@
 from dataclasses import dataclass, field
 from itertools import starmap
 from logging import getLogger
-from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, cast, overload
+from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, TypedDict, cast, overload
 from warnings import warn
 
 import numcodecs
@@ -109,8 +109,10 @@
     from typing import Self
 
     from zarr.abc.codec import CodecPipeline
+    from zarr.codecs.sharding import ShardingCodecIndexLocation
     from zarr.core.group import AsyncGroup
 
+
 # Array and AsyncArray are defined in the base ``zarr`` namespace
 __all__ = ["create_codec_pipeline", "parse_array_metadata"]
 
@@ -271,7 +273,7 @@ def __init__(
     # this overload defines the function signature when zarr_format is 2
     @overload
     @classmethod
-    async def create(
+    async def _create(
         cls,
         store: StoreLike,
         *,
@@ -295,7 +297,7 @@ async def create(
     # this overload defines the function signature when zarr_format is 3
     @overload
     @classmethod
-    async def create(
+    async def _create(
         cls,
         store: StoreLike,
         *,
@@ -323,7 +325,7 @@ async def create(
 
     @overload
     @classmethod
-    async def create(
+    async def _create(
         cls,
         store: StoreLike,
         *,
@@ -350,7 +352,7 @@ async def create(
     ) -> AsyncArray[ArrayV3Metadata]: ...
     @overload
     @classmethod
-    async def create(
+    async def _create(
         cls,
         store: StoreLike,
         *,
@@ -383,7 +385,7 @@ async def create(
     ) -> AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]: ...
 
     @classmethod
-    async def create(
+    async def _create(
         cls,
         store: StoreLike,
         *,
@@ -494,19 +496,6 @@ async def create(
         -------
         AsyncArray
             The created asynchronous array instance.
-
-        Examples
-        --------
-        >>> import zarr
-        >>> store = zarr.storage.MemoryStore(mode='w')
-        >>> async_arr = await zarr.core.array.AsyncArray.create(
-        >>>     store=store,
-        >>>     shape=(100,100),
-        >>>     chunks=(10,10),
-        >>>     dtype='i4',
-        >>>     fill_value=0)
-        <AsyncArray memory://140349042942400 shape=(100, 100) dtype=int32>
-
         """
         store_path = await make_store_path(store)
 
@@ -1148,7 +1137,7 @@ async def getitem(
         --------
         >>> import zarr
         >>>  store = zarr.storage.MemoryStore(mode='w')
-        >>>  async_arr = await zarr.core.array.AsyncArray.create(
+        >>>  async_arr = await zarr.api.asynchronous.create_array(
         ...      store=store,
         ...      shape=(100,100),
         ...      chunks=(10,10),
@@ -1542,7 +1531,7 @@ class Array:
 
     @classmethod
     @_deprecate_positional_args
-    def create(
+    def _create(
         cls,
         store: StoreLike,
         *,
@@ -1643,7 +1632,7 @@ def create(
             Array created from the store.
         """
         async_array = sync(
-            AsyncArray.create(
+            AsyncArray._create(
                 store=store,
                 shape=shape,
                 dtype=dtype,
@@ -2025,10 +2014,10 @@ def __getitem__(self, selection: Selection) -> NDArrayLike:
             >>> import zarr
             >>> import numpy as np
             >>> data = np.arange(100, dtype="uint16")
-            >>> z = Array.create(
+            >>> z = zarr.create_array(
             >>>        StorePath(MemoryStore(mode="w")),
             >>>        shape=data.shape,
-            >>>        chunk_shape=(10,),
+            >>>        chunks=(10,),
             >>>        dtype=data.dtype,
             >>>        )
             >>> z[:] = data
@@ -2059,10 +2048,10 @@ def __getitem__(self, selection: Selection) -> NDArrayLike:
         Setup a 2-dimensional array::
 
             >>> data = np.arange(100, dtype="uint16").reshape(10, 10)
-            >>> z = Array.create(
+            >>> z = zarr.create_array(
             >>>        StorePath(MemoryStore(mode="w")),
             >>>        shape=data.shape,
-            >>>        chunk_shape=(10, 10),
+            >>>        chunks=(10, 10),
             >>>        dtype=data.dtype,
             >>>        )
             >>> z[:] = data
@@ -2290,10 +2279,10 @@ def get_basic_selection(
             >>> import zarr
             >>> import numpy as np
             >>> data = np.arange(100, dtype="uint16")
-            >>> z = Array.create(
+            >>> z = zarr.create_array(
             >>>        StorePath(MemoryStore(mode="w")),
             >>>        shape=data.shape,
-            >>>        chunk_shape=(3,),
+            >>>        chunks=(3,),
             >>>        dtype=data.dtype,
             >>>        )
             >>> z[:] = data
@@ -2319,10 +2308,10 @@ def get_basic_selection(
         Setup a 3-dimensional array::
 
             >>> data = np.arange(1000).reshape(10, 10, 10)
-            >>> z = Array.create(
+            >>> z = zarr.create_array(
             >>>        StorePath(MemoryStore(mode="w")),
             >>>        shape=data.shape,
-            >>>        chunk_shape=(5, 5, 5),
+            >>>        chunks=(5, 5, 5),
             >>>        dtype=data.dtype,
             >>>        )
             >>> z[:] = data
@@ -2514,10 +2503,10 @@ def get_orthogonal_selection(
             >>> import zarr
             >>> import numpy as np
             >>> data = np.arange(100).reshape(10, 10)
-            >>> z = Array.create(
+            >>> z = zarr.create_array(
             >>>        StorePath(MemoryStore(mode="w")),
             >>>        shape=data.shape,
-            >>>        chunk_shape=data.shape,
+            >>>        chunks=data.shape,
             >>>        dtype=data.dtype,
             >>>        )
             >>> z[:] = data
@@ -2748,10 +2737,10 @@ def get_mask_selection(
             >>> import zarr
             >>> import numpy as np
             >>> data = np.arange(100).reshape(10, 10)
-            >>> z = Array.create(
+            >>> z = zarr.create_array(
             >>>        StorePath(MemoryStore(mode="w")),
             >>>        shape=data.shape,
-            >>>        chunk_shape=data.shape,
+            >>>        chunks=data.shape,
             >>>        dtype=data.dtype,
             >>>        )
             >>> z[:] = data
@@ -2908,10 +2897,10 @@ def get_coordinate_selection(
             >>> import zarr
             >>> import numpy as np
             >>> data = np.arange(0, 100, dtype="uint16").reshape((10, 10))
-            >>> z = Array.create(
+            >>> z = zarr.create_array(
             >>>        StorePath(MemoryStore(mode="w")),
             >>>        shape=data.shape,
-            >>>        chunk_shape=(3, 3),
+            >>>        chunks=(3, 3),
             >>>        dtype=data.dtype,
             >>>        )
             >>> z[:] = data
@@ -3096,10 +3085,10 @@ def get_block_selection(
             >>> import zarr
             >>> import numpy as np
             >>> data = np.arange(0, 100, dtype="uint16").reshape((10, 10))
-            >>> z = Array.create(
+            >>> z = zarr.create_array(
             >>>        StorePath(MemoryStore(mode="w")),
             >>>        shape=data.shape,
-            >>>        chunk_shape=(3, 3),
+            >>>        chunks=(3, 3),
             >>>        dtype=data.dtype,
             >>>        )
             >>> z[:] = data
@@ -3517,14 +3506,22 @@ def _get_default_codecs(
 )
 
 
+class ShardsConfigParam(TypedDict):
+    shape: ChunkCoords
+    index_location: ShardingCodecIndexLocation | None
+
+
+ShardsParam: TypeAlias = ChunkCoords | ShardsConfigParam | Literal["auto"]
+
+
 async def create_array(
     store: str | StoreLike,
     *,
     name: str | None = None,
     shape: ShapeLike,
     dtype: npt.DTypeLike,
     chunks: ChunkCoords | Literal["auto"] = "auto",
-    shards: ChunkCoords | Literal["auto"] | None = None,
+    shards: ShardsParam | None = None,
     filters: FiltersParam = "auto",
     compressors: CompressorsParam = "auto",
     fill_value: Any | None = 0,
@@ -3616,13 +3613,24 @@ async def create_array(
     -------
     z : array
         The array.
+
+    Examples
+    --------
+    >>> import zarr
+    >>> store = zarr.storage.MemoryStore(mode='w')
+    >>> async_arr = await zarr.api.asynchronous.create_array(
+    >>>     store=store,
+    >>>     shape=(100,100),
+    >>>     chunks=(10,10),
+    >>>     dtype='i4',
+    >>>     fill_value=0)
+    <AsyncArray memory://140349042942400 shape=(100, 100) dtype=int32>
     """
 
     if zarr_format is None:
         zarr_format = _default_zarr_version()
 
-    # TODO: figure out why putting these imports at top-level causes circular imports
-    from zarr.codecs.sharding import ShardingCodec
+    from zarr.codecs.sharding import ShardingCodec, ShardingCodecIndexLocation
 
     mode: Literal["a"] = "a"
     dtype_parsed = parse_dtype(dtype, zarr_format=zarr_format)
@@ -3677,7 +3685,14 @@ async def create_array(
         sub_codecs = cast(tuple[Codec, ...], (*array_array, array_bytes, *bytes_bytes))
         codecs_out: tuple[Codec, ...]
         if shard_shape_parsed is not None:
-            sharding_codec = ShardingCodec(chunk_shape=chunk_shape_parsed, codecs=sub_codecs)
+            index_location = None
+            if isinstance(shards, dict):
+                index_location = ShardingCodecIndexLocation(shards.get("index_location", None))
+            if index_location is None:
+                index_location = ShardingCodecIndexLocation.end
+            sharding_codec = ShardingCodec(
+                chunk_shape=chunk_shape_parsed, codecs=sub_codecs, index_location=index_location
+            )
             sharding_codec.validate(
                 shape=chunk_shape_parsed,
                 dtype=dtype_parsed,

src/zarr/core/chunk_grids.py

@@ -27,6 +27,8 @@
     from collections.abc import Iterator
     from typing import Self
 
+    from zarr.core.array import ShardsParam
+
 
 def _guess_chunks(
     shape: ShapeLike,
@@ -201,7 +203,7 @@ def _auto_partition(
     *,
     array_shape: tuple[int, ...],
     chunk_shape: tuple[int, ...] | Literal["auto"],
-    shard_shape: tuple[int, ...] | Literal["auto"] | None,
+    shard_shape: ShardsParam | None,
     dtype: np.dtype[Any],
 ) -> tuple[tuple[int, ...] | None, tuple[int, ...]]:
     """
@@ -241,6 +243,8 @@ def _auto_partition(
                     _shards_out += (c_shape * 2,)
                 else:
                     _shards_out += (c_shape,)
+        elif isinstance(shard_shape, dict):
+            _shards_out = tuple(shard_shape["shape"])
         else:
             _shards_out = shard_shape