8000 0 dim arrays: indexing (#1980) · AdamWill/zarr-python@65dc4cc · GitHub
[go: up one dir, main page]
Skip to content

Commit 65dc4cc

Browse files
d-v-bd-v-b
authored
0 dim arrays: indexing (zarr-developers#1980)
* fill in shapelike tests, handle negative values correctly, allow 0 * support indexing arrays with empty shape * fix behavior of AsyncArray._set_selection when converting a scalar input into an array. ensures that the dtype of the array is used. * remove array interface dict
1 parent 6fc05b7 commit 65dc4cc

File tree

2 files changed

+111
-88
lines changed

2 files changed

+111
-88
lines changed

‎src/zarr/array.py

Lines changed: 26 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -487,7 +487,7 @@ async def _set_selection(
487487

488488
# check value shape
489489
if np.isscalar(value):
490-
value = np.asanyarray(value)
490+
value = np.asanyarray(value, dtype=self.metadata.dtype)
491491
else:
492492
if not hasattr(value, "shape"):
493493
value = np.asarray(value, self.metadata.dtype)
@@ -700,6 +700,24 @@ def order(self) -> Literal["C", "F"]:
700700
def read_only(self) -> bool:
701701
return self._async_array.read_only
702702

703+
def __array__(
704+
self, dtype: npt.DTypeLike | None = None, copy: bool | None = None
705+
) -> NDArrayLike:
706+
"""
707+
This method is used by numpy when converting zarr.Array into a numpy array.
708+
For more information, see https://numpy.org/devdocs/user/basics.interoperability.html#the-array-method
709+
"""
710+
if copy is False:
711+
msg = "`copy=False` is not supported. This method always creates a copy."
712+
raise ValueError(msg)
713+
714+
arr_np = self[...]
715+
716+
if dtype is not None:
717+
arr_np = arr_np.astype(dtype)
718+
719+
return arr_np
720+
703721
def __getitem__(self, selection: Selection) -> NDArrayLike:
704722
"""Retrieve data for an item or region of the array.
705723
@@ -1062,17 +1080,14 @@ def get_basic_selection(
10621080
10631081
"""
10641082

1065-
if self.shape == ():
1066-
raise NotImplementedError
1067-
else:
1068-
return sync(
1069-
self._async_array._get_selection(
1070-
BasicIndexer(selection, self.shape, self.metadata.chunk_grid),
1071-
out=out,
1072-
fields=fields,
1073-
prototype=prototype,
1074-
)
1083+
return sync(
1084+
self._async_array._get_selection(
1085+
BasicIndexer(selection, self.shape, self.metadata.chunk_grid),
1086+
out=out,
1087+
fields=fields,
1088+
prototype=prototype,
10751089
)
1090+
)
10761091

10771092
def set_basic_selection(
10781093
self,

‎tests/v3/test_indexing.py

Lines changed: 85 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ def zarr_array_from_numpy_array(
4242
chunk_shape=chunk_shape or a.shape,
4343
chunk_key_encoding=("v2", "."),
4444
)
45-
z[:] = a
45+
z[()] = a
4646
return z
4747

4848

@@ -111,42 +111,55 @@ def test_replace_ellipsis():
111111
)
112112

113113

114-
@pytest.mark.xfail(reason="zero-dimension arrays are not supported in v3")
115-
def test_get_basic_selection_0d(store: StorePath):
114+
@pytest.mark.parametrize(
115+
"value, dtype",
116+
[
117+
(42, "uint8"),
118+
pytest.param(
119+
(b"aaa", 1, 4.2), [("foo", "S3"), ("bar", "i4"), ("baz", "f8")], marks=pytest.mark.xfail
120+
),
121+
],
122+
)
123+
@pytest.mark.parametrize("use_out", (True, False))
124+
def test_get_basic_selection_0d(store: StorePath, use_out: bool, value: Any, dtype: Any) -> None:
116125
# setup
117-
a = np.array(42)
118-
z = zarr_array_from_numpy_array(store, a)
126+
arr_np = np.array(value, dtype=dtype)
127+
arr_z = zarr_array_from_numpy_array(store, arr_np)
119128

120-
assert_array_equal(a, z.get_basic_selection(Ellipsis))
121-
assert_array_equal(a, z[...])
122-
assert 42 == z.get_basic_selection(())
123-
assert 42 == z[()]
129+
assert_array_equal(arr_np, arr_z.get_basic_selection(Ellipsis))
130+
assert_array_equal(arr_np, arr_z[...])
131+
assert value == arr_z.get_basic_selection(())
132+
assert value == arr_z[()]
124133

125-
# test out param
126-
b = NDBuffer.from_numpy_array(np.zeros_like(a))
127-
z.get_basic_selection(Ellipsis, out=b)
128-
assert_array_equal(a, b)
134+
if use_out:
135+
# test out param
136+
b = NDBuffer.from_numpy_array(np.zeros_like(arr_np))
137+
arr_z.get_basic_selection(Ellipsis, out=b)
138+
assert_array_equal(arr_np, b.as_ndarray_like())
139+
140+
# todo: uncomment the structured array tests when we can make them pass,
141+
# or delete them if we formally decide not to support structured dtypes.
129142

130143
# test structured array
131-
value = (b"aaa", 1, 4.2)
132-
a = np.array(value, dtype=[("foo", "S3"), ("bar", "i4"), ("baz", "f8")])
133-
z = zarr_array_from_numpy_array(store, a)
134-
z[()] = value
135-
assert_array_equal(a, z.get_basic_selection(Ellipsis))
136-
assert_array_equal(a, z[...])
137-
assert a[()] == z.get_basic_selection(())
138-
assert a[()] == z[()]
139-
assert b"aaa" == z.get_basic_selection((), fields="foo")
140-
assert b"aaa" == z["foo"]
141-
assert a[["foo", "bar"]] == z.get_basic_selection((), fields=["foo", "bar"])
142-
assert a[["foo", "bar"]] == z["foo", "bar"]
143-
# test out param
144-
b = NDBuffer.from_numpy_array(np.zeros_like(a))
145-
z.get_basic_selection(Ellipsis, out=b)
146-
assert_array_equal(a, b)
147-
c = NDBuffer.from_numpy_array(np.zeros_like(a[["foo", "bar"]]))
148-
z.get_basic_selection(Ellipsis, out=c, fields=["foo", "bar"])
149-
assert_array_equal(a[["foo", "bar"]], c)
144+
# value = (b"aaa", 1, 4.2)
145+
# a = np.array(value, dtype=[("foo", "S3"), ("bar", "i4"), ("baz", "f8")])
146+
# z = zarr_array_from_numpy_array(store, a)
147+
# z[()] = value
148+
# assert_array_equal(a, z.get_basic_selection(Ellipsis))
149+
# assert_array_equal(a, z[...])
150+
# assert a[()] == z.get_basic_selection(())
151+
# assert a[()] == z[()]
152+
# assert b"aaa" == z.get_basic_selection((), fields="foo")
153+
# assert b"aaa" == z["foo"]
154+
# assert a[["foo", "bar"]] == z.get_basic_selection((), fields=["foo", "bar"])
155+
# assert a[["foo", "bar"]] == z["foo", "bar"]
156+
# # test out param
157+
# b = NDBuffer.from_numpy_array(np.zeros_like(a))
158+
# z.get_basic_selection(Ellipsis, out=b)
159+
# assert_array_equal(a, b)
160+
# c = NDBuffer.from_numpy_array(np.zeros_like(a[["foo", "bar"]]))
161+
# z.get_basic_selection(Ellipsis, out=c, fields=["foo", "bar"])
162+
# assert_array_equal(a[["foo", "bar"]], c)
150163

151164

152165
basic_selections_1d = [
@@ -466,51 +479,46 @@ def test_fancy_indexing_doesnt_mix_with_implicit_slicing(store: StorePath):
466479
np.testing.assert_array_equal(z2[..., [1, 2, 3]], 0)
467480

468481

469-
@pytest.mark.xfail(reason="zero-dimension arrays are not supported in v3")
470-
def test_set_basic_selection_0d(store: StorePath):
471-
# setup
472-
v = np.array(42)
473-
a = np.zeros_like(v)
474-
z = zarr_array_from_numpy_array(store, v)
475-
assert_array_equal(a, z[:])
476-
477-
# tests
478-
z.set_basic_selection(Ellipsis, v)
479-
assert_array_equal(v, z[:])
480-
z[...] = 0
481-
assert_array_equal(a, z[:])
482-
z[...] = v
483-
assert_array_equal(v, z[:])
484-
485-
# test structured array
486-
value = (b"aaa", 1, 4.2)
487-
v = np.array(value, dtype=[("foo", "S3"), ("bar", "i4"), ("baz", "f8")])
488-
a = np.zeros_like(v)
489-
z = zarr_array_from_numpy_array(store, a)
490-
491-
# tests
492-
z.set_basic_selection(Ellipsis, v)
493-
assert_array_equal(v, z[:])
494-
z.set_basic_selection(Ellipsis, a)
495-
assert_array_equal(a, z[:])
496-
z[...] = v
497-
assert_array_equal(v, z[:])
498-
z[...] = a
499-
assert_array_equal(a, z[:])
500-
# with fields
501-
z.set_basic_selection(Ellipsis, v["foo"], fields="foo")
502-
assert v["foo"] == z["foo"]
503-
assert a["bar"] == z["bar"]
504-
assert a["baz"] == z["baz"]
505-
z["bar"] = v["bar"]
506-
assert v["foo"] == z["foo"]
507-
assert v["bar"] == z["bar"]
508-
assert a["baz"] == z["baz"]
509-
# multiple field assignment not supported
510-
with pytest.raises(IndexError):
511-
z.set_basic_selection(Ellipsis, v[["foo", "bar"]], fields=["foo", "bar"])
512-
with pytest.raises(IndexError):
513-
z[..., "foo", "bar"] = v[["foo", "bar"]]
482+
@pytest.mark.parametrize(
483+
"value, dtype",
484+
[
485+
(42, "uint8"),
486+
pytest.param(
487+
(b"aaa", 1, 4.2), [("foo", "S3"), ("bar", "i4"), ("baz", "f8")], marks=pytest.mark.xfail
488+
),
489+
],
490+
)
491+
def test_set_basic_selection_0d(
492+
store: StorePath, value: Any, dtype: str | list[tuple[str, str]]
493+
) -> None:
494+
arr_np = np.array(value, dtype=dtype)
495+
arr_np_zeros = np.zeros_like(arr_np, dtype=dtype)
496+
arr_z = zarr_array_from_numpy_array(store, arr_np_zeros)
497+
assert_array_equal(arr_np_zeros, arr_z)
498+
499+
arr_z.set_basic_selection(Ellipsis, value)
500+
assert_array_equal(value, arr_z)
501+
arr_z[...] = 0
502+
assert_array_equal(arr_np_zeros, arr_z)
503+
arr_z[...] = value
504+
assert_array_equal(value, arr_z)
505+
506+
# todo: uncomment the structured array tests when we can make them pass,
507+
# or delete them if we formally decide not to support structured dtypes.
508+
509+
# arr_z.set_basic_selection(Ellipsis, v["foo"], fields="foo")
510+
# assert v["foo"] == arr_z["foo"]
511+
# assert arr_np_zeros["bar"] == arr_z["bar"]
512+
# assert arr_np_zeros["baz"] == arr_z["baz"]
513+
# arr_z["bar"] = v["bar"]
514+
# assert v["foo"] == arr_z["foo"]
515+
# assert v["bar"] == arr_z["bar"]
516+
# assert arr_np_zeros["baz"] == arr_z["baz"]
517+
# # multiple field assignment not supported
518+
# with pytest.raises(IndexError):
519+
# arr_z.set_basic_selection(Ellipsis, v[["foo", "bar"]], fields=["foo", "bar"])
520+
# with pytest.raises(IndexError):
521+
# arr_z[..., "foo", "bar"] = v[["foo", "bar"]]
514522

515523

516524
def _test_get_orthogonal_selection(a, z, selection):

0 commit comments

Comments
 (0)
0