chore: handle new zarr dtype API (#100)

* (chore): handle new zarr dtype API

* (fix): dep in pyproject toml for zarr from main

* (fix): use local store for v2

* fix: don't use `default_filters`/`default_compressor`

* fix: clean up/slim donw tests

* Update pyproject.toml

---------

Co-authored-by: Lachlan Deakin <[email protected]>

Author: ilan-gold

pyproject.toml

@@ -20,7 +20,7 @@ classifiers = [
 dynamic = ["version"]
 dependencies = [
     "numpy>=1.24",
-    "zarr>=3.0.3,<3.1",
+    "zarr>=3.1",
 ]
 
 [dependency-groups]

python/zarrs/pipeline.py

@@ -21,6 +21,7 @@
     from zarr.core.chunk_grids import ChunkGrid
     from zarr.core.common import ChunkCoords
     from zarr.core.indexing import SelectorTuple
+    from zarr.dtype import ZDType
 
 from ._internal import CodecPipelineImpl, codec_metadata_v2_to_v3
 from .utils import (
@@ -134,7 +135,7 @@ def __iter__(self) -> Iterator[Codec]:
         yield from self.codecs
 
     def validate(
-        self, *, shape: ChunkCoords, dtype: np.dtype[Any], chunk_grid: ChunkGrid
+        self, *, shape: ChunkCoords, dtype: ZDType, chunk_grid: ChunkGrid
     ) -> None:
         raise NotImplementedError("validate")
 
@@ -236,7 +237,7 @@ def _raise_error_on_unsupported_batch_dtype(
         # https://github.com/LDeakin/zarrs/blob/0532fe983b7b42b59dbf84e50a2fe5e6f7bad4ce/zarrs_metadata/src/v2_to_v3.rs#L289-L293 for VSUMm
         # Further, our pipeline does not support variable-length objects due to limitations on decode_into, so object/np.dtypes.StringDType is also out
         if any(
-            info.dtype.kind in {"V", "S", "U", "M", "m", "O", "T"}
+            info.dtype.to_native_dtype().kind in {"V", "S", "U", "M", "m", "O", "T"}
             for (_, info, _, _, _) in batch_info
         ):
             raise UnsupportedDataTypeError()

python/zarrs/utils.py

@@ -8,7 +8,6 @@
 import numpy as np
 from zarr.core.array_spec import ArraySpec
 from zarr.core.indexing import SelectorTuple, is_integer
-from zarr.core.metadata.v2 import _default_fill_value
 
 from zarrs._internal import Basic, WithSubset
 
@@ -17,6 +16,7 @@
     from types import EllipsisType
 
     from zarr.abc.store import ByteGetter, ByteSetter
+    from zarr.dtype import ZDType
 
 
 # adapted from https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ThreadPoolExecutor
@@ -139,9 +139,9 @@ def get_shape_for_selector(
     return resulting_shape_from_index(shape, selector_tuple, drop_axes, pad=pad)
 
 
-def get_implicit_fill_value(dtype: np.dtype, fill_value: Any) -> Any:
+def get_implicit_fill_value(dtype: ZDType, fill_value: Any) -> Any:
     if fill_value is None:
-        fill_value = _default_fill_value(dtype)
+        fill_value = dtype.default_scalar()
     return fill_value
 
 

src/chunk_item.rs

@@ -68,6 +68,7 @@ impl Basic {
         let chunk_shape = chunk_spec.getattr("shape")?.extract()?;
         let mut dtype: String = chunk_spec
             .getattr("dtype")?
+            .call_method0("to_native_dtype")?
             .call_method0("__str__")?
             .extract()?;
         if dtype == "object" {

tests/test_blosc.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import json
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
@@ -13,14 +12,14 @@
     TransposeCodec,
 )
 from zarr.core.buffer import default_buffer_prototype
-from zarr.core.indexing import Selection, morton_order_iter
 from zarr.storage import StorePath
 
 if TYPE_CHECKING:
     from zarr.abc.codec import Codec
     from zarr.abc.store import Store
     from zarr.core.buffer.core import NDArrayLike
     from zarr.core.common import MemoryOrder
+    from zarr.core.indexing import Selection
 
 
 @dataclass(frozen=True)
@@ -165,51 +164,6 @@ def test_order_implicit(
         assert read_data.flags["C_CONTIGUOUS"]
 
 
-def test_open(store: Store) -> None:
-    spath = StorePath(store)
-    a = Array.create(
-        spath,
-        shape=(16, 16),
-        chunk_shape=(16, 16),
-        dtype="int32",
-        fill_value=0,
-    )
-    b = Array.open(spath)
-    assert a.metadata == b.metadata
-
-
-def test_morton() -> None:
-    assert list(morton_order_iter((2, 2))) == [(0, 0), (1, 0), (0, 1), (1, 1)]
-    assert list(morton_order_iter((2, 2, 2))) == [
-        (0, 0, 0),
-        (1, 0, 0),
-        (0, 1, 0),
-        (1, 1, 0),
-        (0, 0, 1),
-        (1, 0, 1),
-        (0, 1, 1),
-        (1, 1, 1),
-    ]
-    assert list(morton_order_iter((2, 2, 2, 2))) == [
-        (0, 0, 0, 0),
-        (1, 0, 0, 0),
-        (0, 1, 0, 0),
-        (1, 1, 0, 0),
-        (0, 0, 1, 0),
-        (1, 0, 1, 0),
-        (0, 1, 1, 0),
-        (1, 1, 1, 0),
-        (0, 0, 0, 1),
-        (1, 0, 0, 1),
-        (0, 1, 0, 1),
-        (1, 1, 0, 1),
-        (0, 0, 1, 1),
-        (1, 0, 1, 1),
-        (0, 1, 1, 1),
-        (1, 1, 1, 1),
-    ]
-
-
 def test_write_partial_chunks(store: Store) -> None:
     data = np.arange(0, 256, dtype="uint16").reshape((16, 16))
     spath = StorePath(store)
@@ -241,41 +195,6 @@ async def test_delete_empty_chunks(store: Store) -> None:
     assert await store.get(f"{path}/c0/0", prototype=default_buffer_prototype()) is None
 
 
-async def test_dimension_names(store: Store) -> None:
-    data = np.arange(0, 256, dtype="uint16").reshape((16, 16))
-    path = "dimension_names"
-    spath = StorePath(store, path)
-    await AsyncArray.create(
-        spath,
-        shape=data.shape,
-        chunk_shape=(16, 16),
-        dtype=data.dtype,
-        fill_value=0,
-        dimension_names=("x", "y"),
-    )
-
-    assert (await AsyncArray.open(spath)).metadata.dimension_names == (
-        "x",
-        "y",
-    )
-    path2 = "dimension_names2"
-    spath2 = StorePath(store, path2)
-    await AsyncArray.create(
-        spath2,
-        shape=data.shape,
-        chunk_shape=(16, 16),
-        dtype=data.dtype,
-        fill_value=0,
-    )
-
-    assert (await AsyncArray.open(spath2)).metadata.dimension_names is None
-    zarr_json_buffer = await store.get(
-        f"{path2}/zarr.json", prototype=default_buffer_prototype()
-    )
-    assert zarr_json_buffer is not None
-    assert "dimension_names" not in json.loads(zarr_json_buffer.to_bytes())
-
-
 def test_invalid_metadata(store: Store) -> None:
     # LD: Disabled for `zarrs`. Including endianness for a single-byte data type is not invalid.
     # spath2 = StorePath(store, "invalid_endian")

tests/test_codecs.py

@@ -1,4 +1,3 @@
-import pickle
 from typing import Any
 
 import numpy as np
@@ -286,30 +285,6 @@ def test_nested_sharding(
     assert np.array_equal(data, read_data)
 
 
-def test_open_sharding(store: Store) -> None:
-    path = "open_sharding"
-    spath = StorePath(store, path)
-    a = Array.create(
-        spath,
-        shape=(16, 16),
-        chunk_shape=(16, 16),
-        dtype="int32",
-        fill_value=0,
-        codecs=[
-            ShardingCodec(
-                chunk_shape=(8, 8),
-                codecs=[
-                    TransposeCodec(order=order_from_dim("F", 2)),
-                    BytesCodec(),
-                    BloscCodec(),
-                ],
-            )
-        ],
-    )
-    b = Array.open(spath)
-    assert a.metadata == b.metadata
-
-
 def test_write_partial_sharded_chunks(store: Store) -> None:
     data = np.arange(0, 16 * 16, dtype="uint16").reshape((16, 16))
     spath = StorePath(store)
@@ -365,11 +340,6 @@ async def test_delete_empty_shards(store: Store) -> None:
     assert len(chunk_bytes) == 16 * 2 + 8 * 8 * 2 + 4
 
 
-def test_pickle() -> None:
-    codec = ShardingCodec(chunk_shape=(8, 8))
-    assert pickle.loads(pickle.dumps(codec)) == codec
-
-
 @pytest.mark.parametrize(
     "index_location", [ShardingCodecIndexLocation.start, ShardingCodecIndexLocation.end]
 )

tests/test_sharding.py

@@ -84,21 +84,3 @@ def test_transpose_non_self_inverse(store: Store, order: list[int]) -> None:
     a[:, :] = data
     read_data = a[:, :]
     assert np.array_equal(data, read_data)
-
-
-def test_transpose_invalid(
-    store: Store,
-) -> None:
-    data = np.arange(0, 256, dtype="uint16").reshape((1, 32, 8))
-    spath = StorePath(store, "transpose_invalid")
-    for order in [(1, 0), (3, 2, 1), (3, 3, 1)]:
-        with pytest.raises(ValueError, match=r".*order"):
-            Array.create(
-                spath,
-                shape=data.shape,
-                chunk_shape=(1, 32, 8),
-                dtype=data.dtype,
-                fill_value=0,
-                chunk_key_encoding=("v2", "."),
-                codecs=[TransposeCodec(order=order), BytesCodec()],
-            )