Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Return scalar when accessing zero dimensional array #2718

Open
wants to merge 69 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
69 commits
Select commit Hold shift + click to select a range
a9531ca
Return scalar when accessing zero dimensional array
brokkoli71 Jan 15, 2025
1a290c7
returning npt.ArrayLike instead of NDArrayLike because of scalar retu…
brokkoli71 Jan 15, 2025
3348439
returning npt.ArrayLike instead of NDArrayLike because of scalar retu…
brokkoli71 Jan 15, 2025
34bf260
fix mypy in tests
brokkoli71 Jan 16, 2025
75d6cdf
fix mypy in tests
brokkoli71 Jan 16, 2025
66c0798
fix mypy in tests
brokkoli71 Jan 16, 2025
a86d041
improve test_scalar_array
brokkoli71 Jan 23, 2025
f8393b0
Merge branch 'main' into return-scalar-for-zero-dim-indexing
brokkoli71 Jan 29, 2025
c2a0de0
fix typo
brokkoli71 Jan 29, 2025
2475960
add ScalarWrapper
brokkoli71 Jan 29, 2025
cced470
use ScalarWrapper as NDArrayLike
brokkoli71 Jan 29, 2025
d57e1f2
Revert "fix mypy in tests"
brokkoli71 Jan 29, 2025
842bf95
Revert "fix mypy in tests"
brokkoli71 Jan 29, 2025
ee6d62d
Revert "fix mypy in tests"
brokkoli71 Jan 29, 2025
e302ae6
format
brokkoli71 Jan 29, 2025
b55c8b3
Revert "returning npt.ArrayLike instead of NDArrayLike because of sca…
brokkoli71 Jan 29, 2025
f1cb6e1
Revert "returning npt.ArrayLike instead of NDArrayLike because of sca…
brokkoli71 Jan 29, 2025
c76be48
fix mypy for ScalarWrapper
brokkoli71 Jan 29, 2025
359eb66
add missing import NDArrayLike
brokkoli71 Jan 29, 2025
fc0937e
ignore unavoidable mypy error
brokkoli71 Jan 29, 2025
8454d7b
format
brokkoli71 Jan 29, 2025
15c5103
fix __array__
brokkoli71 Jan 29, 2025
a93ce00
extend tests
brokkoli71 Jan 29, 2025
805a8df
format
brokkoli71 Jan 29, 2025
f6b48ba
fix typing in test_scalar_array
brokkoli71 Jan 29, 2025
1b7966f
add dtype to ScalarWrapper
brokkoli71 Jan 29, 2025
81ab808
correct dtype type
brokkoli71 Jan 29, 2025
56c52ae
fix test_basic_indexing
brokkoli71 Jan 29, 2025
a3473fb
fix test_basic_indexing
brokkoli71 Jan 29, 2025
b98578b
Merge remote-tracking branch 'origin/return-scalar-for-zero-dim-index…
brokkoli71 Jan 29, 2025
c150462
fix test_basic_indexing for dtype=datetime64[Y]
brokkoli71 Jan 29, 2025
a7d4421
increase codecov
brokkoli71 Jan 29, 2025
632cad0
fix typing
brokkoli71 Jan 29, 2025
50fd5ff
document changes
brokkoli71 Jan 29, 2025
2aca6c2
move test_scalar_wrapper to test_buffer.py
brokkoli71 Jan 29, 2025
15ec284
Merge branch 'main' into return-scalar-for-zero-dim-indexing
brokkoli71 Jan 30, 2025
a842e88
Merge branch 'main' into return-scalar-for-zero-dim-indexing
brokkoli71 Feb 1, 2025
e0ef9b1
Merge branch 'main' into return-scalar-for-zero-dim-indexing
brokkoli71 Feb 5, 2025
24ed5c2
Merge branch 'main' into return-scalar-for-zero-dim-indexing
brokkoli71 Feb 9, 2025
ab7afbb
Merge branch 'main' into return-scalar-for-zero-dim-indexing
brokkoli71 Feb 11, 2025
98e13a8
remove ScalarWrapper usage
brokkoli71 Feb 12, 2025
862481a
Merge branch 'refs/heads/main' into return-scalar-for-zero-dim-indexing
brokkoli71 Feb 12, 2025
888304a
create NDArrayOrScalarLike
brokkoli71 Feb 12, 2025
422e0d8
fix NDArrayOrScalarLike
brokkoli71 Feb 12, 2025
a9c0eab
fix mypy
brokkoli71 Feb 12, 2025
b005620
fix mypy
brokkoli71 Feb 12, 2025
9284246
fix mypy
brokkoli71 Feb 12, 2025
73c835f
fix mypy in asynchronous.py
brokkoli71 Feb 12, 2025
ada708e
fix mypy in test_api.py
brokkoli71 Feb 12, 2025
0886e77
fix mypy in test_api.py and synchronous.py
brokkoli71 Feb 12, 2025
300d03a
fix mypy in test_api.py and test_array.py
brokkoli71 Feb 12, 2025
584762b
fix mypy in test_array.py
brokkoli71 Feb 12, 2025
3835768
fix mypy in test_array.py
brokkoli71 Feb 12, 2025
363f06b
fix mypy in test_array.py
brokkoli71 Feb 12, 2025
17d4a9b
fix mypy in test_array.py
brokkoli71 Feb 12, 2025
d4b8bfb
fix mypy in test_array.py, test_api.py, test_buffer.py, test_sharding.py
brokkoli71 Feb 12, 2025
2b18ab1
add bytes, str and datetime to ScalarType
brokkoli71 Feb 12, 2025
070c673
only support numpy datetime64 in ScalarType
brokkoli71 Feb 12, 2025
655d464
remove ScalarWrapper and update changes
brokkoli71 Feb 12, 2025
f03c026
undo wrong code changes
brokkoli71 Feb 12, 2025
a557456
rename ``NDArrayOrScalarLike`` to ``NDArrayLikeOrScalar``
brokkoli71 Feb 12, 2025
b23994c
rename ``NDArrayOrScalarLike`` to ``NDArrayLikeOrScalar``
brokkoli71 Feb 15, 2025
b4c53a9
Merge branch 'main' into return-scalar-for-zero-dim-indexing
brokkoli71 Feb 15, 2025
995a5d5
fix mypy in test_array.py
brokkoli71 Feb 15, 2025
8e5f47d
fix mypy in test_array.py
brokkoli71 Feb 15, 2025
b68c60e
Merge branch 'main' into return-scalar-for-zero-dim-indexing
brokkoli71 Feb 19, 2025
56d1dea
handle datetype scalars for different units
brokkoli71 Feb 19, 2025
276bdf3
fix mypy
brokkoli71 Feb 19, 2025
75a267b
fix mypy
brokkoli71 Feb 19, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions changes/2718.bugfix.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
0-dimensional arrays are now returning a scalar. Therefore, the return type of ``__getitem__`` changed
to NDArrayLikeOrScalar. This change is to make the behavior of 0-dimensional arrays consistent with
``numpy`` scalars.
3 changes: 2 additions & 1 deletion src/zarr/api/asynchronous.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
from collections.abc import Iterable

from zarr.abc.codec import Codec
from zarr.core.buffer import NDArrayLikeOrScalar
from zarr.core.chunk_key_encodings import ChunkKeyEncoding
from zarr.storage import StoreLike

Expand Down Expand Up @@ -232,7 +233,7 @@ async def load(
path: str | None = None,
zarr_format: ZarrFormat | None = None,
zarr_version: ZarrFormat | None = None,
) -> NDArrayLike | dict[str, NDArrayLike]:
) -> NDArrayLikeOrScalar | dict[str, NDArrayLikeOrScalar]:
"""Load data from an array or group into memory.

Parameters
Expand Down
4 changes: 2 additions & 2 deletions src/zarr/api/synchronous.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
ShardsLike,
)
from zarr.core.array_spec import ArrayConfigLike
from zarr.core.buffer import NDArrayLike
from zarr.core.buffer import NDArrayLike, NDArrayLikeOrScalar
from zarr.core.chunk_key_encodings import ChunkKeyEncoding, ChunkKeyEncodingLike
from zarr.core.common import (
JSON,
Expand Down Expand Up @@ -119,7 +119,7 @@ def load(
path: str | None = None,
zarr_format: ZarrFormat | None = None,
zarr_version: ZarrFormat | None = None,
) -> NDArrayLike | dict[str, NDArrayLike]:
) -> NDArrayLikeOrScalar | dict[str, NDArrayLikeOrScalar]:
"""Load data from an array or group into memory.

Parameters
Expand Down
48 changes: 26 additions & 22 deletions src/zarr/core/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
from zarr.core.buffer import (
BufferPrototype,
NDArrayLike,
NDArrayLikeOrScalar,
NDBuffer,
default_buffer_prototype,
)
Expand Down Expand Up @@ -1256,7 +1257,7 @@
prototype: BufferPrototype,
out: NDBuffer | None = None,
fields: Fields | None = None,
) -> NDArrayLike:
) -> NDArrayLikeOrScalar:
# check fields are sensible
out_dtype = check_fields(fields, self.dtype)

Expand Down Expand Up @@ -1298,14 +1299,16 @@
out_buffer,
drop_axes=indexer.drop_axes,
)
if isinstance(indexer, BasicIndexer) and indexer.shape == ():
return out_buffer.as_scalar()

Check warning on line 1303 in src/zarr/core/array.py

View check run for this annotation

Codecov / codecov/patch

src/zarr/core/array.py#L1303

Added line #L1303 was not covered by tests
return out_buffer.as_ndarray_like()

async def getitem(
self,
selection: BasicSelection,
*,
prototype: BufferPrototype | None = None,
) -> NDArrayLike:
) -> NDArrayLikeOrScalar:
"""
Asynchronous function that retrieves a subset of the array's data based on the provided selection.

Expand All @@ -1318,7 +1321,7 @@

Returns
-------
NDArrayLike
NDArrayLikeOrScalar
The retrieved subset of the array's data.

Examples
Expand Down Expand Up @@ -2270,14 +2273,15 @@
msg = "`copy=False` is not supported. This method always creates a copy."
raise ValueError(msg)

arr_np = self[...]
arr = self[...]
arr_np: NDArrayLike = np.array(arr, dtype=dtype)

Check warning on line 2277 in src/zarr/core/array.py

View check run for this annotation

Codecov / codecov/patch

src/zarr/core/array.py#L2277

Added line #L2277 was not covered by tests

if dtype is not None:
arr_np = arr_np.astype(dtype)

return arr_np

def __getitem__(self, selection: Selection) -> NDArrayLike:
def __getitem__(self, selection: Selection) -> NDArrayLikeOrScalar:
"""Retrieve data for an item or region of the array.

Parameters
Expand All @@ -2288,8 +2292,8 @@

Returns
-------
NDArrayLike
An array-like containing the data for the requested region.
NDArrayLikeOrScalar
An array-like or scalar containing the data for the requested region.

Examples
--------
Expand Down Expand Up @@ -2535,7 +2539,7 @@
out: NDBuffer | None = None,
prototype: BufferPrototype | None = None,
fields: Fields | None = None,
) -> NDArrayLike:
) -> NDArrayLikeOrScalar:
"""Retrieve data for an item or region of the array.

Parameters
Expand All @@ -2553,8 +2557,8 @@

Returns
-------
NDArrayLike
An array-like containing the data for the requested region.
NDArrayLikeOrScalar
An array-like or scalar containing the data for the requested region.

Examples
--------
Expand Down Expand Up @@ -2755,7 +2759,7 @@
out: NDBuffer | None = None,
fields: Fields | None = None,
prototype: BufferPrototype | None = None,
) -> NDArrayLike:
) -> NDArrayLikeOrScalar:
"""Retrieve data by making a selection for each dimension of the array. For
example, if an array has 2 dimensions, allows selecting specific rows and/or
columns. The selection for each dimension can be either an integer (indexing a
Expand All @@ -2777,8 +2781,8 @@

Returns
-------
NDArrayLike
An array-like containing the data for the requested selection.
NDArrayLikeOrScalar
An array-like or scalar containing the data for the requested selection.

Examples
--------
Expand Down Expand Up @@ -2991,7 +2995,7 @@
out: NDBuffer | None = None,
fields: Fields | None = None,
prototype: BufferPrototype | None = None,
) -> NDArrayLike:
) -> NDArrayLikeOrScalar:
"""Retrieve a selection of individual items, by providing a Boolean array of the
same shape as the array against which the selection is being made, where True
values indicate a selected item.
Expand All @@ -3011,8 +3015,8 @@

Returns
-------
NDArrayLike
An array-like containing the data for the requested selection.
NDArrayLikeOrScalar
An array-like or scalar containing the data for the requested selection.

Examples
--------
Expand Down Expand Up @@ -3153,7 +3157,7 @@
out: NDBuffer | None = None,
fields: Fields | None = None,
prototype: BufferPrototype | None = None,
) -> NDArrayLike:
) -> NDArrayLikeOrScalar:
"""Retrieve a selection of individual items, by providing the indices
(coordinates) for each selected item.

Expand All @@ -3171,8 +3175,8 @@

Returns
-------
NDArrayLike
An array-like containing the data for the requested coordinate selection.
NDArrayLikeOrScalar
An array-like or scalar containing the data for the requested coordinate selection.

Examples
--------
Expand Down Expand Up @@ -3341,7 +3345,7 @@
out: NDBuffer | None = None,
fields: Fields | None = None,
prototype: BufferPrototype | None = None,
) -> NDArrayLike:
) -> NDArrayLikeOrScalar:
"""Retrieve a selection of individual items, by providing the indices
(coordinates) for each selected item.

Expand All @@ -3359,8 +3363,8 @@

Returns
-------
NDArrayLike
An array-like containing the data for the requested block selection.
NDArrayLikeOrScalar
An array-like or scalar containing the data for the requested block selection.

Examples
--------
Expand Down
2 changes: 2 additions & 0 deletions src/zarr/core/buffer/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
Buffer,
BufferPrototype,
NDArrayLike,
NDArrayLikeOrScalar,
NDBuffer,
default_buffer_prototype,
)
Expand All @@ -13,6 +14,7 @@
"Buffer",
"BufferPrototype",
"NDArrayLike",
"NDArrayLikeOrScalar",
"NDBuffer",
"default_buffer_prototype",
"numpy_buffer_prototype",
Expand Down
19 changes: 19 additions & 0 deletions src/zarr/core/buffer/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,10 @@
"""


ScalarType = int | float | complex | bytes | str | bool | np.generic
NDArrayLikeOrScalar = ScalarType | NDArrayLike


def check_item_key_is_1d_contiguous(key: Any) -> None:
"""Raises error if `key` isn't a 1d contiguous slice"""
if not isinstance(key, slice):
Expand Down Expand Up @@ -419,6 +423,21 @@
"""
...

def as_scalar(self) -> ScalarType:
"""Returns the buffer as a scalar value"""
if self._data.size != 1:
raise ValueError("Buffer does not contain a single scalar value")
item = self.as_numpy_array().item()

Check warning on line 430 in src/zarr/core/buffer/core.py

View check run for this annotation

Codecov / codecov/patch

src/zarr/core/buffer/core.py#L428-L430

Added lines #L428 - L430 were not covered by tests
scalar: ScalarType

if np.issubdtype(self.dtype, np.datetime64):
unit: str = np.datetime_data(self.dtype)[0] # Extract the unit (e.g., 'Y', 'D', etc.)
scalar = np.datetime64(item, unit)

Check warning on line 435 in src/zarr/core/buffer/core.py

View check run for this annotation

Codecov / codecov/patch

src/zarr/core/buffer/core.py#L433-L435

Added lines #L433 - L435 were not covered by tests
else:
scalar = self.dtype.type(item) # Regular conversion for non-datetime types

Check warning on line 437 in src/zarr/core/buffer/core.py

View check run for this annotation

Codecov / codecov/patch

src/zarr/core/buffer/core.py#L437

Added line #L437 was not covered by tests

return scalar

Check warning on line 439 in src/zarr/core/buffer/core.py

View check run for this annotation

Codecov / codecov/patch

src/zarr/core/buffer/core.py#L439

Added line #L439 was not covered by tests

@property
def dtype(self) -> np.dtype[Any]:
return self._data.dtype
Expand Down
10 changes: 6 additions & 4 deletions src/zarr/core/indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@

if TYPE_CHECKING:
from zarr.core.array import Array
from zarr.core.buffer import NDArrayLike
from zarr.core.buffer import NDArrayLikeOrScalar
from zarr.core.chunk_grids import ChunkGrid
from zarr.core.common import ChunkCoords

Expand Down Expand Up @@ -937,7 +937,7 @@ class OIndex:
array: Array

# TODO: develop Array generic and move zarr.Array[np.intp] | zarr.Array[np.bool_] to ArrayOfIntOrBool
def __getitem__(self, selection: OrthogonalSelection | Array) -> NDArrayLike:
def __getitem__(self, selection: OrthogonalSelection | Array) -> NDArrayLikeOrScalar:
from zarr.core.array import Array

# if input is a Zarr array, we materialize it now.
Expand Down Expand Up @@ -1046,7 +1046,7 @@ def __iter__(self) -> Iterator[ChunkProjection]:
class BlockIndex:
array: Array

def __getitem__(self, selection: BasicSelection) -> NDArrayLike:
def __getitem__(self, selection: BasicSelection) -> NDArrayLikeOrScalar:
fields, new_selection = pop_fields(selection)
new_selection = ensure_tuple(new_selection)
new_selection = replace_lists(new_selection)
Expand Down Expand Up @@ -1236,7 +1236,9 @@ class VIndex:
array: Array

# TODO: develop Array generic and move zarr.Array[np.intp] | zarr.Array[np.bool_] to ArrayOfIntOrBool
def __getitem__(self, selection: CoordinateSelection | MaskSelection | Array) -> NDArrayLike:
def __getitem__(
self, selection: CoordinateSelection | MaskSelection | Array
) -> NDArrayLikeOrScalar:
from zarr.core.array import Array

# if input is a Zarr array, we materialize it now.
Expand Down
21 changes: 16 additions & 5 deletions tests/test_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
save_array,
save_group,
)
from zarr.core.buffer import NDArrayLike
from zarr.core.common import JSON, MemoryOrder, ZarrFormat
from zarr.errors import MetadataValidationError
from zarr.storage import MemoryStore
Expand Down Expand Up @@ -236,7 +237,9 @@ def test_open_with_mode_r(tmp_path: pathlib.Path) -> None:
z2 = zarr.open(store=tmp_path, mode="r")
assert isinstance(z2, Array)
assert z2.fill_value == 1
assert (z2[:] == 1).all()
result = z2[:]
assert isinstance(result, NDArrayLike)
assert (result == 1).all()
with pytest.raises(ValueError):
z2[:] = 3

Expand All @@ -248,7 +251,9 @@ def test_open_with_mode_r_plus(tmp_path: pathlib.Path) -> None:
zarr.ones(store=tmp_path, shape=(3, 3))
z2 = zarr.open(store=tmp_path, mode="r+")
assert isinstance(z2, Array)
assert (z2[:] == 1).all()
result = z2[:]
assert isinstance(result, NDArrayLike)
assert (result == 1).all()
z2[:] = 3


Expand All @@ -264,7 +269,9 @@ async def test_open_with_mode_a(tmp_path: pathlib.Path) -> None:
arr[...] = 1
z2 = zarr.open(store=tmp_path, mode="a")
assert isinstance(z2, Array)
assert (z2[:] == 1).all()
result = z2[:]
assert isinstance(result, NDArrayLike)
assert (result == 1).all()
z2[:] = 3


Expand All @@ -276,7 +283,9 @@ def test_open_with_mode_w(tmp_path: pathlib.Path) -> None:
arr[...] = 3
z2 = zarr.open(store=tmp_path, mode="w", shape=(3, 3))
assert isinstance(z2, Array)
assert not (z2[:] == 3).all()
result = z2[:]
assert isinstance(result, NDArrayLike)
assert not (result == 3).all()
z2[:] = 3


Expand Down Expand Up @@ -1120,7 +1129,9 @@ def test_open_array_with_mode_r_plus(store: Store) -> None:
zarr.ones(store=store, shape=(3, 3))
z2 = zarr.open_array(store=store, mode="r+")
assert isinstance(z2, Array)
assert (z2[:] == 1).all()
result = z2[:]
assert isinstance(result, NDArrayLike)
assert (result == 1).all()
z2[:] = 3


Expand Down
Loading