zarr.array from from an existing zarr.Array (#2622)

* add creation from other zarr

* remove duplicated tests

* improve test

* test_iter_grid for non-squares

* concurrent streaming for equal chunk sizes

* fix merge

* fix mypy

* fix mypy

* fix test_iter_grid

* extract to zarr.from_array

* fix mypy

* fix mypy

* format

* fix test_creation_from_other_zarr_format

* distinguish between keep and auto for from_array arguments

* partition concurrency along new_array chunks

* fix mypy

* improve test_creation_from_other_zarr_format

* add typing in test

* Update src/zarr/core/array.py

Co-authored-by: Norman Rzepka <[email protected]>

* add from_array with npt.ArrayLike

* add write_data argument

* improve tests

* improve docstrings and add examples

* fix mypy and readthedocs

* fix mypy and readthedocs

* fix mypy and readthedocs

* fix mypy and readthedocs

* fix readthedocs ERROR: Unexpected indentation

* add release notes

* format docstring examples

* add write_data attr to synchronous.create_array

* `create_array` calls `from_array` calls `init_array`

* document changes

* fix serializer from_array v2 to v3

* fix mypy

* improve codecov

* fix mypy

* from_array: copy zarr format on default

* in ``from_array`` make all arguments except ``store`` keyword-only, to match ``create_array``

* in ``from_array`` default shards="keep"

* redundant ``ChunkKeyEncoding | ChunkKeyEncodingLike``

* fix argument order in calls of `from_array`

* fix numpydoc-validation

* add docstring to store2 pytest fixture

* extract `_parse_keep_array_attr` from `from_array`

* extract `_parse_keep_array_attr` from `from_array`

* correct _parse_keep_array_attr

* fix merge

* fix merge

---------

Co-authored-by: Norman Rzepka <[email protected]>

Author: brokkoli71

changes/2622.feature.rst

@@ -0,0 +1 @@
+Add ``zarr.from_array`` using concurrent streaming of source data

docs/release-notes.rst

@@ -145,6 +145,8 @@ Other
 3.0.1 (Jan. 17, 2025)
 ---------------------
 
+* Implement ``zarr.from_array`` using concurrent streaming (:issue:`2622`).
+
 Bug fixes
 ~~~~~~~~~
 * Fixes ``order`` argument for Zarr format 2 arrays (:issue:`2679`).

src/zarr/__init__.py

@@ -11,6 +11,7 @@
     create_hierarchy,
     empty,
     empty_like,
+    from_array,
     full,
     full_like,
     group,
@@ -54,6 +55,7 @@
     "create_hierarchy",
     "empty",
     "empty_like",
+    "from_array",
     "full",
     "full_like",
     "group",

src/zarr/api/asynchronous.py

@@ -9,7 +9,7 @@
 import numpy.typing as npt
 from typing_extensions import deprecated
 
-from zarr.core.array import Array, AsyncArray, create_array, get_array_metadata
+from zarr.core.array import Array, AsyncArray, create_array, from_array, get_array_metadata
 from zarr.core.array_spec import ArrayConfig, ArrayConfigLike, ArrayConfigParams
 from zarr.core.buffer import NDArrayLike
 from zarr.core.common import (
@@ -57,6 +57,7 @@
     "create_hierarchy",
     "empty",
     "empty_like",
+    "from_array",
     "full",
     "full_like",
     "group",
@@ -534,7 +535,7 @@ async def tree(grp: AsyncGroup, expand: bool | None = None, level: int | None =
 
 
 async def array(
-    data: npt.ArrayLike, **kwargs: Any
+    data: npt.ArrayLike | Array, **kwargs: Any
 ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array filled with `data`.
 
@@ -551,13 +552,16 @@ async def array(
         The new array.
     """
 
+    if isinstance(data, Array):
+        return await from_array(data=data, **kwargs)
+
     # ensure data is array-like
     if not hasattr(data, "shape") or not hasattr(data, "dtype"):
         data = np.asanyarray(data)
 
     # setup dtype
     kw_dtype = kwargs.get("dtype")
-    if kw_dtype is None:
+    if kw_dtype is None and hasattr(data, "dtype"):
         kwargs["dtype"] = data.dtype
     else:
         kwargs["dtype"] = kw_dtype

src/zarr/api/synchronous.py

@@ -50,6 +50,7 @@
     "create_hierarchy",
     "empty",
     "empty_like",
+    "from_array",
     "full",
     "full_like",
     "group",
@@ -359,7 +360,7 @@ def tree(grp: Group, expand: bool | None = None, level: int | None = None) -> An
 
 
 # TODO: add type annotations for kwargs
-def array(data: npt.ArrayLike, **kwargs: Any) -> Array:
+def array(data: npt.ArrayLike | Array, **kwargs: Any) -> Array:
     """Create an array filled with `data`.
 
     Parameters
@@ -759,11 +760,12 @@ def create_array(
     order: MemoryOrder | None = None,
     zarr_format: ZarrFormat | None = 3,
     attributes: dict[str, JSON] | None = None,
-    chunk_key_encoding: ChunkKeyEncoding | ChunkKeyEncodingLike | None = None,
+    chunk_key_encoding: ChunkKeyEncodingLike | None = None,
     dimension_names: Iterable[str] | None = None,
     storage_options: dict[str, Any] | None = None,
     overwrite: bool = False,
     config: ArrayConfigLike | None = None,
+    write_data: bool = True,
 ) -> Array:
     """Create an array.
 
@@ -857,6 +859,11 @@ def create_array(
         Whether to overwrite an array with the same name in the store, if one exists.
     config : ArrayConfigLike, optional
         Runtime configuration for the array.
+    write_data : bool
+        If a pre-existing array-like object was provided to this function via the ``data`` parameter
+        then ``write_data`` determines whether the values in that array-like object should be
+        written to the Zarr array created by this function. If ``write_data`` is ``False``, then the
+        array will be left empty.
 
     Returns
     -------
@@ -866,7 +873,7 @@ def create_array(
     Examples
     --------
     >>> import zarr
-    >>> store = zarr.storage.MemoryStore(mode='w')
+    >>> store = zarr.storage.MemoryStore()
     >>> arr = await zarr.create_array(
     >>>     store=store,
     >>>     shape=(100,100),
@@ -897,6 +904,220 @@ def create_array(
                 storage_options=storage_options,
                 overwrite=overwrite,
                 config=config,
+                write_data=write_data,
+            )
+        )
+    )
+
+
+def from_array(
+    store: str | StoreLike,
+    *,
+    data: Array | npt.ArrayLike,
+    write_data: bool = True,
+    name: str | None = None,
+    chunks: Literal["auto", "keep"] | ChunkCoords = "keep",
+    shards: ShardsLike | None | Literal["keep"] = "keep",
+    filters: FiltersLike | Literal["keep"] = "keep",
+    compressors: CompressorsLike | Literal["keep"] = "keep",
+    serializer: SerializerLike | Literal["keep"] = "keep",
+    fill_value: Any | None = None,
+    order: MemoryOrder | None = None,
+    zarr_format: ZarrFormat | None = None,
+    attributes: dict[str, JSON] | None = None,
+    chunk_key_encoding: ChunkKeyEncodingLike | None = None,
+    dimension_names: Iterable[str] | None = None,
+    storage_options: dict[str, Any] | None = None,
+    overwrite: bool = False,
+    config: ArrayConfigLike | None = None,
+) -> Array:
+    """Create an array from an existing array or array-like.
+
+    Parameters
+    ----------
+    store : str or Store
+        Store or path to directory in file system or name of zip file for the new array.
+    data : Array | array-like
+        The array to copy.
+    write_data : bool, default True
+        Whether to copy the data from the input array to the new array.
+        If ``write_data`` is ``False``, the new array will be created with the same metadata as the
+        input array, but without any data.
+    name : str or None, optional
+        The name of the array within the store. If ``name`` is ``None``, the array will be located
+        at the root of the store.
+    chunks : ChunkCoords or "auto" or "keep", optional
+        Chunk shape of the array.
+        Following values are supported:
+
+        - "auto": Automatically determine the chunk shape based on the array's shape and dtype.
+        - "keep": Retain the chunk shape of the data array if it is a zarr Array.
+        - ChunkCoords: A tuple of integers representing the chunk shape.
+
+        If not specified, defaults to "keep" if data is a zarr Array, otherwise "auto".
+    shards : ChunkCoords, optional
+        Shard shape of the array.
+        Following values are supported:
+
+        - "auto": Automatically determine the shard shape based on the array's shape and chunk shape.
+        - "keep": Retain the shard shape of the data array if it is a zarr Array.
+        - ChunkCoords: A tuple of integers representing the shard shape.
+        - None: No sharding.
+
+        If not specified, defaults to "keep" if data is a zarr Array, otherwise None.
+    filters : Iterable[Codec] or "auto" or "keep", optional
+        Iterable of filters to apply to each chunk of the array, in order, before serializing that
+        chunk to bytes.
+
+        For Zarr format 3, a "filter" is a codec that takes an array and returns an array,
+        and these values must be instances of ``ArrayArrayCodec``, or dict representations
+        of ``ArrayArrayCodec``.
+
+        For Zarr format 2, a "filter" can be any numcodecs codec; you should ensure that the
+        the order if your filters is consistent with the behavior of each filter.
+
+        Following values are supported:
+
+        - Iterable[Codec]: List of filters to apply to the array.
+        - "auto": Automatically determine the filters based on the array's dtype.
+        - "keep": Retain the filters of the data array if it is a zarr Array.
+
+        If no ``filters`` are provided, defaults to "keep" if data is a zarr Array, otherwise "auto".
+    compressors : Iterable[Codec] or "auto" or "keep", optional
+        List of compressors to apply to the array. Compressors are applied in order, and after any
+        filters are applied (if any are specified) and the data is serialized into bytes.
+
+        For Zarr format 3, a "compressor" is a codec that takes a bytestream, and
+        returns another bytestream. Multiple compressors my be provided for Zarr format 3.
+
+        For Zarr format 2, a "compressor" can be any numcodecs codec. Only a single compressor may
+        be provided for Zarr format 2.
+
+        Following values are supported:
+
+        - Iterable[Codec]: List of compressors to apply to the array.
+        - "auto": Automatically determine the compressors based on the array's dtype.
+        - "keep": Retain the compressors of the input array if it is a zarr Array.
+
+        If no ``compressors`` are provided, defaults to "keep" if data is a zarr Array, otherwise "auto".
+    serializer : dict[str, JSON] | ArrayBytesCodec or "auto" or "keep", optional
+        Array-to-bytes codec to use for encoding the array data.
+        Zarr format 3 only. Zarr format 2 arrays use implicit array-to-bytes conversion.
+
+        Following values are supported:
+
+        - dict[str, JSON]: A dict representation of an ``ArrayBytesCodec``.
+        - ArrayBytesCodec: An instance of ``ArrayBytesCodec``.
+        - "auto": a default serializer will be used. These defaults can be changed by modifying the value of
+          ``array.v3_default_serializer`` in :mod:`zarr.core.config`.
+        - "keep": Retain the serializer of the input array if it is a zarr Array.
+
+    fill_value : Any, optional
+        Fill value for the array.
+        If not specified, defaults to the fill value of the data array.
+    order : {"C", "F"}, optional
+        The memory of the array (default is "C").
+        For Zarr format 2, this parameter sets the memory order of the array.
+        For Zarr format 3, this parameter is deprecated, because memory order
+        is a runtime parameter for Zarr format 3 arrays. The recommended way to specify the memory
+        order for Zarr format 3 arrays is via the ``config`` parameter, e.g. ``{'config': 'C'}``.
+        If not specified, defaults to the memory order of the data array.
+    zarr_format : {2, 3}, optional
+        The zarr format to use when saving.
+        If not specified, defaults to the zarr format of the data array.
+    attributes : dict, optional
+        Attributes for the array.
+        If not specified, defaults to the attributes of the data array.
+    chunk_key_encoding : ChunkKeyEncoding, optional
+        A specification of how the chunk keys are represented in storage.
+        For Zarr format 3, the default is ``{"name": "default", "separator": "/"}}``.
+        For Zarr format 2, the default is ``{"name": "v2", "separator": "."}}``.
+        If not specified and the data array has the same zarr format as the target array,
+        the chunk key encoding of the data array is used.
+    dimension_names : Iterable[str], optional
+        The names of the dimensions (default is None).
+        Zarr format 3 only. Zarr format 2 arrays should not use this parameter.
+        If not specified, defaults to the dimension names of the data array.
+    storage_options : dict, optional
+        If using an fsspec URL to create the store, these will be passed to the backend implementation.
+        Ignored otherwise.
+    overwrite : bool, default False
+        Whether to overwrite an array with the same name in the store, if one exists.
+    config : ArrayConfig or ArrayConfigLike, optional
+        Runtime configuration for the array.
+
+    Returns
+    -------
+    Array
+        The array.
+
+    Examples
+    --------
+    Create an array from an existing Array::
+
+        >>> import zarr
+        >>> store = zarr.storage.MemoryStore()
+        >>> store2 = zarr.storage.LocalStore('example.zarr')
+        >>> arr = zarr.create_array(
+        >>>     store=store,
+        >>>     shape=(100,100),
+        >>>     chunks=(10,10),
+        >>>     dtype='int32',
+        >>>     fill_value=0)
+        >>> arr2 = zarr.from_array(store2, data=arr)
+        <Array file://example.zarr shape=(100, 100) dtype=int32>
+
+    Create an array from an existing NumPy array::
+
+        >>> import numpy as np
+        >>> arr3 = zarr.from_array(
+                zarr.storage.MemoryStore(),
+        >>>     data=np.arange(10000, dtype='i4').reshape(100, 100),
+        >>> )
+        <Array memory://125477403529984 shape=(100, 100) dtype=int32>
+
+    Create an array from any array-like object::
+
+        >>> arr4 = zarr.from_array(
+        >>>     zarr.storage.MemoryStore(),
+        >>>     data=[[1, 2], [3, 4]],
+        >>> )
+        <Array memory://125477392154368 shape=(2, 2) dtype=int64>
+        >>> arr4[...]
+        array([[1, 2],[3, 4]])
+
+    Create an array from an existing Array without copying the data::
+
+        >>> arr5 = zarr.from_array(
+        >>>     zarr.storage.MemoryStore(),
+        >>>     data=arr4,
+        >>>     write_data=False,
+        >>> )
+        <Array memory://140678602965568 shape=(2, 2) dtype=int64>
+        >>> arr5[...]
+        array([[0, 0],[0, 0]])
+    """
+    return Array(
+        sync(
+            zarr.core.array.from_array(
+                store,
+                data=data,
+                write_data=write_data,
+                name=name,
+                chunks=chunks,
+                shards=shards,
+                filters=filters,
+                compressors=compressors,
+                serializer=serializer,
+                fill_value=fill_value,
+                order=order,
+                zarr_format=zarr_format,
+                attributes=attributes,
+                chunk_key_encoding=chunk_key_encoding,
+                dimension_names=dimension_names,
+                storage_options=storage_options,
+                overwrite=overwrite,
+                config=config,
             )
         )
     )