concurrent streaming for equal chunk sizes

Author: brokkoli71

src/zarr/api/asynchronous.py

@@ -27,7 +27,6 @@
 from zarr.core.group import AsyncGroup, ConsolidatedMetadata, GroupMetadata
 from zarr.core.metadata import ArrayMetadataDict, ArrayV2Metadata, ArrayV3Metadata
 from zarr.core.metadata.v2 import _default_filters_and_compressor
-from zarr.core.sync import sync
 from zarr.errors import NodeTypeValidationError
 from zarr.storage import (
     StoreLike,
@@ -562,16 +561,19 @@ async def array(
 
         new_array = await create(data.shape, **kwargs)
 
-        async def _copy_chunk(chunk_coords: ChunkCoords) -> None:
+        async def _copy_chunk(chunk_coords: ChunkCoords|slice) -> None:
             arr = await data._async_array.getitem(chunk_coords)
             await new_array.setitem(chunk_coords, arr)
 
-        # Stream data from the source array to the new array
-        await concurrent_map(
-            [(region,) for region in data._iter_chunk_regions()],
-            _copy_chunk,
-            config.get("async.concurrency"),
-        )
+        if new_array.chunks == data.chunks:
+            # Stream data from the source array to the new array
+            await concurrent_map(
+                [(region,) for region in data._iter_chunk_regions()],
+                _copy_chunk,
+                config.get("async.concurrency"),
+            )
+        else:
+            await _copy_chunk(slice(None))
         return new_array
 
     # ensure data is array-like

tests/conftest.py

@@ -76,6 +76,12 @@ async def store(request: pytest.FixtureRequest, tmpdir: LEGACY_PATH) -> Store:
     return await parse_store(param, str(tmpdir))
 
 
+@pytest.fixture
+async def store2(request: pytest.FixtureRequest, tmpdir: LEGACY_PATH) -> Store:
+    param = request.param
+    return await parse_store(param, str(tmpdir / "store2"))
+
+
 @pytest.fixture(params=["local", "memory", "zip"])
 def sync_store(request: pytest.FixtureRequest, tmp_path: LEGACY_PATH) -> Store:
     result = sync(parse_store(request.param, str(tmp_path)))

tests/test_array.py

@@ -2,15 +2,13 @@
 import json
 import math
 import pickle
-import time
 from itertools import accumulate
 from typing import Any, Literal
 
 import numcodecs
 import numpy as np
 import pytest
 from numcodecs import Zstd
-from numpy import dtype
 from numpy.ma.testutils import assert_array_equal
 
 import zarr.api.asynchronous
@@ -893,31 +891,40 @@ async def test_scalar_array() -> None:
     assert arr.shape == ()
 
 
-async def test_creation_from_other_zarr(tmpdir):
+@pytest.mark.parametrize("store", ["memory", "local"], indirect=True)
+@pytest.mark.parametrize("store2", ["memory", "local"], indirect=True)
+@pytest.mark.parametrize("src_chunks", [(1, 2), (5, 5), (5, 10)])
+@pytest.mark.parametrize("new_chunks", [(1, 2), (5, 5), (5, 10)])
+async def test_creation_from_other_zarr(store, store2, src_chunks, new_chunks):
     src_fill_value = 2
     src_dtype = np.dtype("uint8")
     src_attributes = {}
-    src_chunks = (2, 2)
 
-    src = zarr.create((10, 10), chunks=src_chunks, dtype=src_dtype, store=LocalStore(str(tmpdir)), fill_value = src_fill_value, attributes=src_attributes)
-    src[:] = np.arange(100).reshape((10,10))
-
-    result = zarr.array(src, store=MemoryStore())
-    assert_array_equal(result[:], src[:])
-    assert result.fill_value == src_fill_value
-    assert result.dtype==src_dtype
-    assert result.attrs.asdict() == src_attributes
-    assert result.chunks == src_chunks
+    src = zarr.create(
+        (10, 10),
+        chunks=src_chunks,
+        dtype=src_dtype,
+        store=store,
+        fill_value=src_fill_value,
+        attributes=src_attributes,
+    )
+    src[:] = np.arange(100).reshape((10, 10))
 
     new_fill_value = 3
     new_dtype = np.dtype("uint16")
-    new_attributes = {"foo":"bar"}
-    new_chunks = (5, 10)
-
-    result2 = zarr.array(src, store=MemoryStore(), chunks=new_chunks, dtype=new_dtype, fill_value = new_fill_value, attributes=new_attributes)
+    new_attributes = {"foo": "bar"}
+
+    result2 = zarr.array(
+        src,
+        store=store2,
+        chunks=new_chunks,
+        dtype=new_dtype,
+        fill_value=new_fill_value,
+        attributes=new_attributes,
+    )
 
     assert_array_equal(result2[:], src[:])
     assert result2.fill_value == new_fill_value
     assert result2.dtype == new_dtype
     assert result2.attrs == new_attributes
-    assert result2.chunks == new_chunks
+    assert result2.chunks == new_chunks

tests/test_indexing.py

@@ -1872,7 +1872,7 @@ def test_iter_grid(
     """
     Test that iter_grid works as expected for 1, 2, and 3 dimensions.
     """
-    grid_shape = (10,2,7)[:ndim]
+    grid_shape = (10, 2, 7)[:ndim]
 
     if origin_0d is not None:
         origin_kwarg = origin_0d * ndim
@@ -1954,3 +1954,13 @@ def test_vectorized_indexing_incompatible_shape(store) -> None:
     )
     with pytest.raises(ValueError, match="Attempting to set"):
         arr[np.array([1, 2]), np.array([1, 2])] = np.array([[-1, -2], [-3, -4]])
+
+
+def test_iter_chunk_regions():
+    chunks = (2, 3)
+    a = zarr.create((10, 10), chunks=chunks)
+    a[:] = 1
+    for region in a._iter_chunk_regions():
+        assert_array_equal(a[region], np.ones_like(a[region]))
+        a[region] = 0
+        assert_array_equal(a[region], np.zeros_like(a[region]))