fix: avoid copying 0-d chunk data when indexing CPU NDBuffers (#4466)

Co-authored-by: Davis Bennett <davis.v.bennett@gmail.com>

Author: janickm

changes/4464.bugfix.md

@@ -0,0 +1 @@
+Fixed slow reads of 0-d arrays with large fixed-length string dtypes (e.g. `|S1000000`) caused by copying the chunk data multiple times.

src/zarr/core/buffer/cpu.py

@@ -220,6 +220,10 @@ def as_numpy_array(self) -> npt.NDArray[Any]:
         return np.asanyarray(self._data)
 
     def __getitem__(self, key: Any) -> Self:
+        # a trailing Ellipsis makes numpy return a 0-d view instead of a scalar, which
+        # np.asanyarray would copy (slow for large fixed-length string dtypes)
+        if isinstance(key, tuple) and not any(k is Ellipsis for k in key):
+            key = (*key, Ellipsis)
         return self.__class__(np.asanyarray(self._data.__getitem__(key)))
 
     def __setitem__(self, key: Any, value: Any) -> None:

tests/test_buffer.py

@@ -198,6 +198,15 @@ def test_gpu_buffer_prototype() -> None:
         ndbuffer.as_scalar()
 
 
+@pytest.mark.parametrize(("shape", "key"), [((), ()), ((3, 4), (1, 2))])
+def test_cpu_ndbuffer_getitem_returns_view(shape: tuple[int, ...], key: tuple[int, ...]) -> None:
+    # https://github.com/zarr-developers/zarr-python/issues/4464
+    data = np.full(shape, b"x" * 10)
+    item = cpu.NDBuffer.from_numpy_array(data)[key].as_numpy_array()
+    assert item.shape == ()
+    assert np.shares_memory(item, data)
+
+
 # TODO: the same test for other buffer classes
 def test_cpu_buffer_as_scalar() -> None:
     buf = cpu.buffer_prototype.nd_buffer.create(shape=(), dtype="int64")