From d9fb6b9725a86e02d47945e189765cbd4e6a1f11 Mon Sep 17 00:00:00 2001 From: Mark Harfouche Date: Mon, 12 Feb 2024 21:08:17 -0500 Subject: [PATCH 1/6] Add test --- xarray/tests/test_variable.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index 6575f2e1740..59bfc8e990b 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -2722,6 +2722,19 @@ def __init__(self, array): orig = Variable(dims=(), data=array) assert isinstance(orig._data.item(), CustomWithValuesAttr) + def test_duck_array_with_chunks(self): + # Non indexable type + class CustomArray(NDArrayMixin, indexing.ExplicitlyIndexed): + def __init__(self, array): + self.array = array + + @property + def chunks(self): + return self.shape + + array = CustomArray(np.arange(3)) + orig = Variable(dims=("x"), data=array, attrs={"foo": "bar"}) + assert isinstance(orig._data, np.ndarray) # should not be CustomArray def test_raise_no_warning_for_nan_in_binary_ops(): with assert_no_warnings(): From 791e1245e3e2958e24eb101ba7d5578ae501f1ca Mon Sep 17 00:00:00 2001 From: Mark Harfouche Date: Mon, 12 Feb 2024 21:44:45 -0500 Subject: [PATCH 2/6] Add a simple test for showing missing ChunkedArray Manager --- xarray/tests/test_variable.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index 59bfc8e990b..c1835b31416 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -29,7 +29,7 @@ from xarray.core.types import T_DuckArray from xarray.core.utils import NDArrayMixin from xarray.core.variable import as_compatible_data, as_variable -from xarray.namedarray.pycompat import array_type +from xarray.namedarray.pycompat import array_type, is_chunked_array from xarray.tests import ( assert_allclose, assert_array_equal, @@ -2724,7 +2724,7 @@ def __init__(self, array): def test_duck_array_with_chunks(self): # Non indexable type - class CustomArray(NDArrayMixin, indexing.ExplicitlyIndexed): + class CustomArray(NDArrayMixin): def __init__(self, array): self.array = array @@ -2732,9 +2732,17 @@ def __init__(self, array): def chunks(self): return self.shape + def __array_function__(self, *args, **kwargs): + return NotImplemented + + def __array_ufunc__(self, *args, **kwargs): + return NotImplemented + + array = CustomArray(np.arange(3)) - orig = Variable(dims=("x"), data=array, attrs={"foo": "bar"}) - assert isinstance(orig._data, np.ndarray) # should not be CustomArray + assert is_chunked_array(array) + var = Variable(dims=("x"), data=array) + var.load() def test_raise_no_warning_for_nan_in_binary_ops(): with assert_no_warnings(): From c402f1224c13625465554818f8f8cfa275492c81 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 13 Feb 2024 02:47:19 +0000 Subject: [PATCH 3/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/tests/test_variable.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index c1835b31416..12a170ad75e 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -2738,12 +2738,12 @@ def __array_function__(self, *args, **kwargs): def __array_ufunc__(self, *args, **kwargs): return NotImplemented - array = CustomArray(np.arange(3)) assert is_chunked_array(array) var = Variable(dims=("x"), data=array) var.load() + def test_raise_no_warning_for_nan_in_binary_ops(): with assert_no_warnings(): Variable("x", [1, 2, np.nan]) > 0 From afbd9b68185428f5374d5196ddf7a7f3878364b0 Mon Sep 17 00:00:00 2001 From: Mark Harfouche Date: Mon, 12 Feb 2024 22:16:08 -0500 Subject: [PATCH 4/6] Really simple --- xarray/namedarray/parallelcompat.py | 4 +--- xarray/namedarray/pycompat.py | 8 +++++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/xarray/namedarray/parallelcompat.py b/xarray/namedarray/parallelcompat.py index c6263bff4ff..161c9d7b7ac 100644 --- a/xarray/namedarray/parallelcompat.py +++ b/xarray/namedarray/parallelcompat.py @@ -155,9 +155,7 @@ def get_chunked_array_type(*args: Any) -> ChunkManagerEntrypoint[Any]: if chunkmanager.is_chunked_array(chunked_arr) ] if not selected: - raise TypeError( - f"Could not find a Chunk Manager which recognises type {type(chunked_arr)}" - ) + return None elif len(selected) >= 2: raise TypeError(f"Multiple ChunkManagers recognise type {type(chunked_arr)}") else: diff --git a/xarray/namedarray/pycompat.py b/xarray/namedarray/pycompat.py index 3ce33d4d8ea..84854ffb33a 100644 --- a/xarray/namedarray/pycompat.py +++ b/xarray/namedarray/pycompat.py @@ -108,7 +108,8 @@ def to_numpy( # TODO first attempt to call .to_numpy() once some libraries implement it if is_chunked_array(data): chunkmanager = get_chunked_array_type(data) - data, *_ = chunkmanager.compute(data, **kwargs) + if chunkmanager is not None: + data, *_ = chunkmanager.compute(data, **kwargs) if isinstance(data, array_type("cupy")): data = data.get() # pint has to be imported dynamically as pint imports xarray @@ -127,8 +128,9 @@ def to_duck_array(data: Any, **kwargs: dict[str, Any]) -> duckarray[_ShapeType, if is_chunked_array(data): chunkmanager = get_chunked_array_type(data) - loaded_data, *_ = chunkmanager.compute(data, **kwargs) # type: ignore[var-annotated] - return loaded_data + if chunkmanager is not None: + loaded_data, *_ = chunkmanager.compute(data, **kwargs) # type: ignore[var-annotated] + return loaded_data if isinstance(data, ExplicitlyIndexed): return data.get_duck_array() # type: ignore[no-untyped-call, no-any-return] From ad1072280f53689ac2b5d18a98e2e961fba0568b Mon Sep 17 00:00:00 2001 From: Mark Harfouche Date: Mon, 12 Feb 2024 22:19:48 -0500 Subject: [PATCH 5/6] mypy --- xarray/namedarray/parallelcompat.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/xarray/namedarray/parallelcompat.py b/xarray/namedarray/parallelcompat.py index 161c9d7b7ac..379cef20c5c 100644 --- a/xarray/namedarray/parallelcompat.py +++ b/xarray/namedarray/parallelcompat.py @@ -11,7 +11,7 @@ from abc import ABC, abstractmethod from collections.abc import Iterable, Sequence from importlib.metadata import EntryPoint, entry_points -from typing import TYPE_CHECKING, Any, Callable, Generic, Protocol, TypeVar +from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, Protocol, TypeVar import numpy as np @@ -121,11 +121,13 @@ def guess_chunkmanager( ) -def get_chunked_array_type(*args: Any) -> ChunkManagerEntrypoint[Any]: +def get_chunked_array_type(*args: Any) -> Optional[ChunkManagerEntrypoint[Any]]: """ Detects which parallel backend should be used for given set of arrays. Also checks that all arrays are of same chunking type (i.e. not a mix of cubed and dask). + + Returns None if no matching ChunkManager is found. """ # TODO this list is probably redundant with something inside xarray.apply_ufunc From cc505c77930130bd527d330f43fe21bf9cd6c182 Mon Sep 17 00:00:00 2001 From: Mark Harfouche Date: Mon, 12 Feb 2024 22:30:03 -0500 Subject: [PATCH 6/6] Just assert --- xarray/namedarray/parallelcompat.py | 4 ++-- xarray/tests/test_parallelcompat.py | 5 +---- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/xarray/namedarray/parallelcompat.py b/xarray/namedarray/parallelcompat.py index 379cef20c5c..c923e55db8b 100644 --- a/xarray/namedarray/parallelcompat.py +++ b/xarray/namedarray/parallelcompat.py @@ -11,7 +11,7 @@ from abc import ABC, abstractmethod from collections.abc import Iterable, Sequence from importlib.metadata import EntryPoint, entry_points -from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, Protocol, TypeVar +from typing import TYPE_CHECKING, Any, Callable, Generic, Protocol, TypeVar import numpy as np @@ -121,7 +121,7 @@ def guess_chunkmanager( ) -def get_chunked_array_type(*args: Any) -> Optional[ChunkManagerEntrypoint[Any]]: +def get_chunked_array_type(*args: Any) -> ChunkManagerEntrypoint[Any] | None: """ Detects which parallel backend should be used for given set of arrays. diff --git a/xarray/tests/test_parallelcompat.py b/xarray/tests/test_parallelcompat.py index d4a4e273bc0..a96ed31c04b 100644 --- a/xarray/tests/test_parallelcompat.py +++ b/xarray/tests/test_parallelcompat.py @@ -195,10 +195,7 @@ def test_raise_if_no_arrays_chunked(self, register_dummy_chunkmanager) -> None: def test_raise_if_no_matching_chunkmanagers(self) -> None: dummy_arr = DummyChunkedArray([1, 2, 3]) - with pytest.raises( - TypeError, match="Could not find a Chunk Manager which recognises" - ): - get_chunked_array_type(dummy_arr) + assert get_chunked_array_type(dummy_arr) is None @requires_dask def test_detect_dask_if_installed(self) -> None: