Repository: cubed-dev/cubed
Pull request #669 (2025-01-15 11:14:48)

Incorrect dtypes in `map_selection`


Patch:
diff --git a/cubed/core/ops.py b/cubed/core/ops.py
index 37202559..d9154183 100644
--- a/cubed/core/ops.py
+++ b/cubed/core/ops.py
@@ -816,9 +816,9 @@ def map_selection(
         key_function,
         x,
         shapes=[shape],
-        dtypes=[x.dtype],
+        dtypes=[dtype],
         chunkss=[chunks],
-        extra_func_kwargs=dict(func=func, dtype=dtype),
+        extra_func_kwargs=dict(func=func, dtype=x.dtype),
         num_input_blocks=num_input_blocks,
         iterable_input_blocks=iterable_input_blocks,
         selection_function=selection_function,
diff --git a/cubed/primitive/blockwise.py b/cubed/primitive/blockwise.py
index bae129a1..31de3af0 100644
--- a/cubed/primitive/blockwise.py
+++ b/cubed/primitive/blockwise.py
@@ -16,7 +16,7 @@ from cubed.backend_array_api import (
     numpy_array_to_backend_array,
 )
 from cubed.runtime.types import CubedPipeline
-from cubed.storage.zarr import T_ZarrArray, lazy_zarr_array
+from cubed.storage.zarr import LazyZarrArray, T_ZarrArray, lazy_zarr_array
 from cubed.types import T_Chunks, T_DType, T_Shape, T_Store
 from cubed.utils import array_memory, chunk_memory, get_item, map_nested
 from cubed.utils import numblocks as compute_numblocks
@@ -321,7 +321,7 @@ def general_blockwise(
 
     write_proxies = []
     output_chunk_memory = 0
-    target_array = []
+    target_arrays = []
 
     numblocks0 = None
     for i, target_store in enumerate(target_stores):
@@ -335,6 +335,7 @@ def general_blockwise(
                 raise ValueError(
                     f"All outputs must have matching number of blocks in each dimension. Chunks specified: {chunkss}"
                 )
+        ta: Union[zarr.Array, LazyZarrArray]
         if isinstance(target_store, zarr.Array):
             ta = target_store
         else:
@@ -347,7 +348,7 @@ def general_blockwise(
                 storage_options=storage_options,
                 compressor=compressor,
             )
-        target_array.append(ta)
+        target_arrays.append(ta)
 
         write_proxies.append(CubedArrayProxy(ta, chunksize))
 
@@ -356,9 +357,6 @@ def general_blockwise(
             output_chunk_memory, array_memory(dtypes[i], chunksize) * 2
         )
 
-    if len(target_array) == 1:
-        target_array = target_array[0]
-
     spec = BlockwiseSpec(
         key_function,
         func_with_kwargs,
@@ -404,7 +402,7 @@ def general_blockwise(
     return PrimitiveOperation(
         pipeline=pipeline,
         source_array_names=array_names,
-        target_array=target_array,
+        target_array=target_arrays[0] if len(target_arrays) == 1 else target_arrays,
         projected_mem=projected_mem,
         allowed_mem=allowed_mem,
         reserved_mem=reserved_mem,
diff --git a/cubed/primitive/rechunk.py b/cubed/primitive/rechunk.py
index cbc6741c..06245898 100644
--- a/cubed/primitive/rechunk.py
+++ b/cubed/primitive/rechunk.py
@@ -1,7 +1,7 @@
 import itertools
 import math
 from math import ceil, prod
-from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, Tuple
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple
 
 import numpy as np
 
@@ -68,10 +68,9 @@ def rechunk(
         storage_options=storage_options,
     )
 
-    intermediate = int_proxy.array
     target = write_proxy.array
 
-    if intermediate is None:
+    if int_proxy is None:
         copy_spec = CubedCopySpec(read_proxy, write_proxy)
         num_tasks = total_chunks(write_proxy.array.shape, write_proxy.chunks)
         return [
@@ -88,6 +87,7 @@ def rechunk(
 
     else:
         # break spec into two if there's an intermediate
+        intermediate = int_proxy.array
         copy_spec1 = CubedCopySpec(read_proxy, int_proxy)
         num_tasks = total_chunks(copy_spec1.write.array.shape, copy_spec1.write.chunks)
         op1 = spec_to_primitive_op(
@@ -123,7 +123,7 @@ def _setup_array_rechunk(
     target_store: T_Store,
     temp_store: Optional[T_Store] = None,
     storage_options: Optional[Dict[str, Any]] = None,
-) -> Tuple[CubedArrayProxy, CubedArrayProxy, CubedArrayProxy]:
+) -> Tuple[CubedArrayProxy, Optional[CubedArrayProxy], CubedArrayProxy]:
     shape = source_array.shape
     source_chunks = source_array.chunks
     dtype = source_array.dtype
@@ -162,7 +162,9 @@ def _setup_array_rechunk(
         )
 
     read_proxy = CubedArrayProxy(source_array, read_chunks)
-    int_proxy = CubedArrayProxy(int_array, int_chunks)
+    int_proxy = (
+        CubedArrayProxy(int_array, int_chunks) if int_array is not None else None
+    )
     write_proxy = CubedArrayProxy(target_array, write_chunks)
     return read_proxy, int_proxy, write_proxy
 
@@ -200,7 +202,7 @@ class ChunkKeys(Iterable[Tuple[slice, ...]]):
         return chunk_keys(self.shape, self.chunks)
 
 
-def copy_read_to_write(chunk_key: Sequence[slice], *, config: CubedCopySpec) -> None:
+def copy_read_to_write(chunk_key: Tuple[slice], *, config: CubedCopySpec) -> None:
     # workaround limitation of lithops.utils.verify_args
     if isinstance(chunk_key, list):
         chunk_key = tuple(chunk_key)
diff --git a/cubed/storage/backend.py b/cubed/storage/backend.py
index 3844a079..07ae0092 100644
--- a/cubed/storage/backend.py
+++ b/cubed/storage/backend.py
@@ -33,7 +33,7 @@ def open_backend_array(
     storage_name = backend_storage_name()
 
     if storage_name == "zarr-python":
-        from cubed.storage.backends.zarr_python import open_zarr_array
+        from cubed.storage.backends.zarr_python import open_zarr_array  # type: ignore
 
         open_func = open_zarr_array
 
diff --git a/cubed/storage/backends/zarr_python.py b/cubed/storage/backends/zarr_python.py
index 0c42dc7a..c811ad8f 100644
--- a/cubed/storage/backends/zarr_python.py
+++ b/cubed/storage/backends/zarr_python.py
@@ -1,3 +1,4 @@
+# type: ignore  # Zarr Python v2 types are different to v3
 from typing import Optional, Union
 
 import zarr
diff --git a/pyproject.toml b/pyproject.toml
index 64620738..41e66842 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -36,7 +36,7 @@ dependencies = [
     "psutil",
     "tenacity",
     "toolz",
-    "zarr < 3",
+    "zarr",
 ]
 
 [project.optional-dependencies]


Test patch:
diff --git a/.github/workflows/zarr-v3-tests.yml b/.github/workflows/zarr-v2-tests.yml
similarity index 91%
rename from .github/workflows/zarr-v3-tests.yml
rename to .github/workflows/zarr-v2-tests.yml
index f1649061..18b36be4 100644
--- a/.github/workflows/zarr-v3-tests.yml
+++ b/.github/workflows/zarr-v2-tests.yml
@@ -1,4 +1,4 @@
-name: Zarr v3 Tests
+name: Zarr v2 Tests
 
 on:
   push:
@@ -44,7 +44,7 @@ jobs:
       - name: Install
         run: |
           python -m pip install -e .[test,lithops]
-          python -m pip install -U git+https://github.com/zarr-developers/zarr-python.git
+          python -m pip install -U 'zarr<3'
 
       - name: Run tests
         env:
diff --git a/cubed/tests/test_overlap.py b/cubed/tests/test_overlap.py
index 5c29e3a7..a388a3ca 100644
--- a/cubed/tests/test_overlap.py
+++ b/cubed/tests/test_overlap.py
@@ -39,6 +39,24 @@ def test_map_overlap_1d_single_chunk():
     assert_array_equal(b.compute(), np.array([0, 0, 1, 2, 3, 4, 5, 0]))
 
 
+def test_map_overlap_1d_change_dtype():
+    x = np.arange(6)
+    a = xp.asarray(x, chunks=(3,))
+
+    b = cubed.map_overlap(
+        lambda x: x.astype(np.float64),
+        a,
+        dtype=np.float64,
+        chunks=((5, 5),),
+        depth=1,
+        boundary=0,
+        trim=False,
+    )
+
+    assert b.dtype == np.float64
+    assert_array_equal(b.compute(), np.array([0, 0, 1, 2, 3, 2, 3, 4, 5, 0]))
+
+
 def test_map_overlap_2d():
     x = np.arange(36).reshape((6, 6))
     a = xp.asarray(x, chunks=(3, 3))