Merge branch 'main' into fix/cache-store-byte-range

Author: d-v-b

changes/3657.bugfix.md

@@ -0,0 +1 @@
+Fix obstore _transform_list_dir implementation to correctly relativize paths (removing lstrip usage).

changes/3704.misc.md

@@ -0,0 +1 @@
+Remove an expensive `isinstance` check from the bytes codec decoding routine.

changes/3705.bugfix.md

@@ -0,0 +1 @@
+Fix a performance bug in morton curve generation.

changes/3706.misc.md

@@ -0,0 +1 @@
+Allow NumPy ints as input when declaring a shape.

src/zarr/codecs/bytes.py

@@ -5,10 +5,8 @@
 from enum import Enum
 from typing import TYPE_CHECKING
 
-import numpy as np
-
 from zarr.abc.codec import ArrayBytesCodec
-from zarr.core.buffer import Buffer, NDArrayLike, NDBuffer
+from zarr.core.buffer import Buffer, NDBuffer
 from zarr.core.common import JSON, parse_enum, parse_named_configuration
 from zarr.core.dtype.common import HasEndianness
 
@@ -72,20 +70,15 @@ async def _decode_single(
         chunk_bytes: Buffer,
         chunk_spec: ArraySpec,
     ) -> NDBuffer:
-        assert isinstance(chunk_bytes, Buffer)
         # TODO: remove endianness enum in favor of literal union
         endian_str = self.endian.value if self.endian is not None else None
         if isinstance(chunk_spec.dtype, HasEndianness):
             dtype = replace(chunk_spec.dtype, endianness=endian_str).to_native_dtype()  # type: ignore[call-arg]
         else:
             dtype = chunk_spec.dtype.to_native_dtype()
         as_array_like = chunk_bytes.as_array_like()
-        if isinstance(as_array_like, NDArrayLike):
-            as_nd_array_like = as_array_like
-        else:
-            as_nd_array_like = np.asanyarray(as_array_like)
         chunk_array = chunk_spec.prototype.nd_buffer.from_ndarray_like(
-            as_nd_array_like.view(dtype=dtype)
+            as_array_like.view(dtype=dtype)  # type: ignore[attr-defined]
         )
 
         # ensure correct chunk shape

src/zarr/core/common.py

@@ -21,6 +21,7 @@
     overload,
 )
 
+import numpy as np
 from typing_extensions import ReadOnly
 
 from zarr.core.config import config as zarr_config
@@ -37,7 +38,7 @@
 ZMETADATA_V2_JSON = ".zmetadata"
 
 BytesLike = bytes | bytearray | memoryview
-ShapeLike = Iterable[int] | int
+ShapeLike = Iterable[int | np.integer[Any]] | int | np.integer[Any]
 # For backwards compatibility
 ChunkCoords = tuple[int, ...]
 ZarrFormat = Literal[2, 3]
@@ -185,23 +186,28 @@ def parse_named_configuration(
 
 
 def parse_shapelike(data: ShapeLike) -> tuple[int, ...]:
-    if isinstance(data, int):
+    """
+    Parse a shape-like input into an explicit shape.
+    """
+    if isinstance(data, int | np.integer):
         if data < 0:
             raise ValueError(f"Expected a non-negative integer. Got {data} instead")
-        return (data,)
+        return (int(data),)
     try:
         data_tuple = tuple(data)
     except TypeError as e:
         msg = f"Expected an integer or an iterable of integers. Got {data} instead."
         raise TypeError(msg) from e
 
-    if not all(isinstance(v, int) for v in data_tuple):
+    if not all(isinstance(v, int | np.integer) for v in data_tuple):
         msg = f"Expected an iterable of integers. Got {data} instead."
         raise TypeError(msg)
     if not all(v > -1 for v in data_tuple):
         msg = f"Expected all values to be non-negative. Got {data} instead."
         raise ValueError(msg)
-    return data_tuple
+
+    # cast NumPy scalars to plain python ints
+    return tuple(int(x) for x in data_tuple)
 
 
 def parse_fill_value(data: Any) -> Any:

src/zarr/core/indexing.py

@@ -7,7 +7,7 @@
 from collections.abc import Iterator, Sequence
 from dataclasses import dataclass
 from enum import Enum
-from functools import reduce
+from functools import lru_cache, reduce
 from types import EllipsisType
 from typing import (
     TYPE_CHECKING,
@@ -1467,16 +1467,21 @@ def decode_morton(z: int, chunk_shape: tuple[int, ...]) -> tuple[int, ...]:
     return tuple(out)
 
 
-def morton_order_iter(chunk_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]:
-    i = 0
+@lru_cache
+def _morton_order(chunk_shape: tuple[int, ...]) -> tuple[tuple[int, ...], ...]:
+    n_total = product(chunk_shape)
     order: list[tuple[int, ...]] = []
-    while len(order) < product(chunk_shape):
+    i = 0
+    while len(order) < n_total:
         m = decode_morton(i, chunk_shape)
-        if m not in order and all(x < y for x, y in zip(m, chunk_shape, strict=False)):
+        if all(x < y for x, y in zip(m, chunk_shape, strict=False)):
             order.append(m)
         i += 1
-    for j in range(product(chunk_shape)):
-        yield order[j]
+    return tuple(order)
+
+
+def morton_order_iter(chunk_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]:
+    return iter(_morton_order(tuple(chunk_shape)))
 
 
 def c_order_iter(chunks_per_shard: tuple[int, ...]) -> Iterator[tuple[int, ...]]:

src/zarr/storage/_obstore.py

@@ -4,6 +4,8 @@
 import contextlib
 import pickle
 from collections import defaultdict
+from itertools import chain
+from operator import itemgetter
 from typing import TYPE_CHECKING, Generic, Self, TypedDict, TypeVar
 
 from zarr.abc.store import (
@@ -15,6 +17,7 @@
 )
 from zarr.core.common import concurrent_map
 from zarr.core.config import config
+from zarr.storage._utils import _relativize_path
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Coroutine, Iterable, Sequence
@@ -263,10 +266,11 @@ async def _transform_list_dir(
     # We assume that the underlying object-store implementation correctly handles the
     # prefix, so we don't double-check that the returned results actually start with the
     # given prefix.
-    prefixes = [obj.lstrip(prefix).lstrip("/") for obj in list_result["common_prefixes"]]
-    objects = [obj["path"].removeprefix(prefix).lstrip("/") for obj in list_result["objects"]]
-    for item in prefixes + objects:
-        yield item
+    prefix = prefix.rstrip("/")
+    for path in chain(
+        list_result["common_prefixes"], map(itemgetter("path"), list_result["objects"])
+    ):
+        yield _relativize_path(path=path, prefix=prefix)
 
 
 class _BoundedRequest(TypedDict):

src/zarr/testing/store.py

@@ -492,24 +492,36 @@ async def test_list_empty_path(self, store: S) -> None:
         assert observed_prefix_sorted == expected_prefix_sorted
 
     async def test_list_dir(self, store: S) -> None:
-        root = "foo"
-        store_dict = {
-            root + "/zarr.json": self.buffer_cls.from_bytes(b"bar"),
-            root + "/c/1": self.buffer_cls.from_bytes(b"\x01"),
-        }
+        roots_and_keys: list[tuple[str, dict[str, Buffer]]] = [
+            (
+                "foo",
+                {
+                    "foo/zarr.json": self.buffer_cls.from_bytes(b"bar"),
+                    "foo/c/1": self.buffer_cls.from_bytes(b"\x01"),
+                },
+            ),
+            (
+                "foo/bar",
+                {
+                    "foo/bar/foobar_first_child": self.buffer_cls.from_bytes(b"1"),
+                    "foo/bar/foobar_second_child/zarr.json": self.buffer_cls.from_bytes(b"2"),
+                },
+            ),
+        ]
 
         assert await _collect_aiterator(store.list_dir("")) == ()
-        assert await _collect_aiterator(store.list_dir(root)) == ()
 
-        await store._set_many(store_dict.items())
+        for root, store_dict in roots_and_keys:
+            assert await _collect_aiterator(store.list_dir(root)) == ()
 
-        keys_observed = await _collect_aiterator(store.list_dir(root))
-        keys_expected = {k.removeprefix(root + "/").split("/")[0] for k in store_dict}
+            await store._set_many(store_dict.items())
 
-        assert sorted(keys_observed) == sorted(keys_expected)
+            keys_observed = await _collect_aiterator(store.list_dir(root))
+            keys_expected = {k.removeprefix(root + "/").split("/")[0] for k in store_dict}
+            assert sorted(keys_observed) == sorted(keys_expected)
 
-        keys_observed = await _collect_aiterator(store.list_dir(root + "/"))
-        assert sorted(keys_expected) == sorted(keys_observed)
+            keys_observed = await _collect_aiterator(store.list_dir(root + "/"))
+            assert sorted(keys_expected) == sorted(keys_observed)
 
     async def test_set_if_not_exists(self, store: S) -> None:
         key = "k"

tests/test_codecs/test_codecs.py

@@ -18,7 +18,7 @@
     TransposeCodec,
 )
 from zarr.core.buffer import default_buffer_prototype
-from zarr.core.indexing import BasicSelection, morton_order_iter
+from zarr.core.indexing import BasicSelection, decode_morton, morton_order_iter
 from zarr.core.metadata.v3 import ArrayV3Metadata
 from zarr.dtype import UInt8
 from zarr.errors import ZarrUserWarning
@@ -171,7 +171,8 @@ def test_open(store: Store) -> None:
     assert a.metadata == b.metadata
 
 
-def test_morton() -> None:
+def test_morton_exact_order() -> None:
+    """Test exact morton ordering for power-of-2 shapes."""
     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),
@@ -206,21 +207,58 @@ def test_morton() -> None:
 @pytest.mark.parametrize(
     "shape",
     [
-        [2, 2, 2],
-        [5, 2],
-        [2, 5],
-        [2, 9, 2],
-        [3, 2, 12],
-        [2, 5, 1],
-        [4, 3, 6, 2, 7],
-        [3, 2, 1, 6, 4, 5, 2],
+        (2, 2, 2),
+        (5, 2),
+        (2, 5),
+        (2, 9, 2),
+        (3, 2, 12),
+        (2, 5, 1),
+        (4, 3, 6, 2, 7),
+        (3, 2, 1, 6, 4, 5, 2),
+        (1,),
+        (1, 1),
+        (5, 1, 3),
+        (1, 4, 1, 2),
     ],
 )
-def test_morton2(shape: tuple[int, ...]) -> None:
+def test_morton_is_permutation(shape: tuple[int, ...]) -> None:
+    """Test that morton_order_iter produces every valid coordinate exactly once."""
+    import itertools
+
+    from zarr.core.common import product
+
+    order = list(morton_order_iter(shape))
+    expected_len = product(shape)
+    # completeness: every valid coordinate is present
+    assert len(order) == expected_len
+    # no duplicates
+    assert len(set(order)) == expected_len
+    # all coordinates are within bounds
+    assert all(all(c < s for c, s in zip(coord, shape, strict=True)) for coord in order)
+    # the set of coordinates equals the full cartesian product
+    assert set(order) == set(itertools.product(*(range(s) for s in shape)))
+
+
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (2, 2),
+        (4, 4),
+        (2, 2, 2),
+        (4, 4, 4),
+        (2, 2, 2, 2),
+    ],
+)
+def test_morton_ordering(shape: tuple[int, ...]) -> None:
+    """Test that the iteration order matches consecutive decode_morton outputs.
+
+    For power-of-2 shapes, every decode_morton output is in-bounds,
+    so the ordering should be exactly decode_morton(0), decode_morton(1), ...
+    """
+
     order = list(morton_order_iter(shape))
-    for i, x in enumerate(order):
-        assert x not in order[:i]  # no duplicates
-        assert all(x[j] < shape[j] for j in range(len(shape)))  # all indices are within bounds
+    for i, coord in enumerate(order):
+        assert coord == decode_morton(i, shape)
 
 
 @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])