Skip to content

Commit

Permalink
friendlier error messages for missing chunk managers (#9676)
Browse files Browse the repository at this point in the history
* raise an error message while guessing if there's no chunkmanager available

* don't skip the no chunkmanager test if dask is not installed

* whats-new

* ensure at least one chunk manager is available

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove additional blank line from a bad merge

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* improve the wording

Co-authored-by: Tom Nicholas <[email protected]>

* switch to ImportError

* raise a helpful `ImportError` for known chunk managers

* make sure the new `ImportError` is actually raised

* check that the more specific error message is preferred

* prefer the more specific error

* also use `ImportError` as indicator for `chunks=None`

* move and improve the whats-new entry

* captialize global variable KNOWN_CHUNKMANAGERS

* chunkmanagers -> available_chunkmanagers

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* also use the string repr for printing `manager`

* reword

* more repr

* reflow

* adapt the test to the new error message

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Tom Nicholas <[email protected]>
Co-authored-by: Deepak Cherian <[email protected]>
  • Loading branch information
4 people authored Dec 29, 2024
1 parent 6875c64 commit dead70f
Show file tree
Hide file tree
Showing 4 changed files with 54 additions and 14 deletions.
2 changes: 2 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@ v.2024.12.0 (unreleased)

New Features
~~~~~~~~~~~~
- Improve the error message raised when using chunked-array methods if no chunk manager is available or if the requested chunk manager is missing (:pull:`9676`)
By `Justus Magin <https://github.com/keewis>`_. (:pull:`9676`)
- Better support wrapping additional array types (e.g. ``cupy`` or ``jax``) by calling generalized
duck array operations throughout more xarray methods. (:issue:`7848`, :pull:`9798`).
By `Sam Levang <https://github.com/slevang>`_.
Expand Down
2 changes: 1 addition & 1 deletion xarray/backends/zarr.py
Original file line number Diff line number Diff line change
Expand Up @@ -1471,7 +1471,7 @@ def open_zarr(
) # attempt to import that parallel backend

chunks = {}
except ValueError:
except (ValueError, ImportError):
chunks = None

if kwargs:
Expand Down
33 changes: 26 additions & 7 deletions xarray/namedarray/parallelcompat.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,12 @@ def compute(

T_ChunkedArray = TypeVar("T_ChunkedArray", bound=ChunkedArrayMixinProtocol)

KNOWN_CHUNKMANAGERS = {
"dask": "dask",
"cubed": "cubed-xarray",
"arkouda": "arkouda-xarray",
}


@functools.lru_cache(maxsize=1)
def list_chunkmanagers() -> dict[str, ChunkManagerEntrypoint[Any]]:
Expand Down Expand Up @@ -95,29 +101,42 @@ def guess_chunkmanager(
Else use whatever is installed, defaulting to dask if there are multiple options.
"""

chunkmanagers = list_chunkmanagers()
available_chunkmanagers = list_chunkmanagers()

if manager is None:
if len(chunkmanagers) == 1:
if len(available_chunkmanagers) == 1:
# use the only option available
manager = next(iter(chunkmanagers.keys()))
manager = next(iter(available_chunkmanagers.keys()))
else:
# use the one in options (default dask)
manager = OPTIONS["chunk_manager"]

if isinstance(manager, str):
if manager not in chunkmanagers:
if manager not in available_chunkmanagers and manager in KNOWN_CHUNKMANAGERS:
raise ImportError(
f"chunk manager {manager!r} is not available."
f" Please make sure {KNOWN_CHUNKMANAGERS[manager]!r} is installed"
" and importable."
)
elif len(available_chunkmanagers) == 0:
raise ImportError(
"no chunk managers available. Try installing `dask` or another package"
" that provides a chunk manager."
)
elif manager not in available_chunkmanagers:
raise ValueError(
f"unrecognized chunk manager {manager} - must be one of: {list(chunkmanagers)}"
f"unrecognized chunk manager {manager!r} - must be one of the installed"
f" chunk managers: {list(available_chunkmanagers)}"
)

return chunkmanagers[manager]
return available_chunkmanagers[manager]
elif isinstance(manager, ChunkManagerEntrypoint):
# already a valid ChunkManager so just pass through
return manager
else:
raise TypeError(
f"manager must be a string or instance of ChunkManagerEntrypoint, but received type {type(manager)}"
"manager must be a string or instance of ChunkManagerEntrypoint,"
f" but received type {type(manager)}"
)


Expand Down
31 changes: 25 additions & 6 deletions xarray/tests/test_parallelcompat.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,14 @@
from xarray.namedarray._typing import _Chunks
from xarray.namedarray.daskmanager import DaskManager
from xarray.namedarray.parallelcompat import (
KNOWN_CHUNKMANAGERS,
ChunkManagerEntrypoint,
get_chunked_array_type,
guess_chunkmanager,
list_chunkmanagers,
load_chunkmanagers,
)
from xarray.tests import has_dask, requires_dask
from xarray.tests import requires_dask


class DummyChunkedArray(np.ndarray):
Expand Down Expand Up @@ -158,18 +159,36 @@ def test_get_chunkmanger_via_set_options(self, register_dummy_chunkmanager) -> N
chunkmanager = guess_chunkmanager(None)
assert isinstance(chunkmanager, DummyChunkManager)

def test_fail_on_nonexistent_chunkmanager(self) -> None:
with pytest.raises(ValueError, match="unrecognized chunk manager foo"):
def test_fail_on_known_but_missing_chunkmanager(
self, register_dummy_chunkmanager, monkeypatch
) -> None:
monkeypatch.setitem(KNOWN_CHUNKMANAGERS, "test", "test-package")
with pytest.raises(
ImportError, match="chunk manager 'test' is not available.+test-package"
):
guess_chunkmanager("test")

def test_fail_on_nonexistent_chunkmanager(
self, register_dummy_chunkmanager
) -> None:
with pytest.raises(ValueError, match="unrecognized chunk manager 'foo'"):
guess_chunkmanager("foo")

@requires_dask
def test_get_dask_if_installed(self) -> None:
chunkmanager = guess_chunkmanager(None)
assert isinstance(chunkmanager, DaskManager)

@pytest.mark.skipif(has_dask, reason="requires dask not to be installed")
def test_dont_get_dask_if_not_installed(self) -> None:
with pytest.raises(ValueError, match="unrecognized chunk manager dask"):
def test_no_chunk_manager_available(self, monkeypatch) -> None:
monkeypatch.setattr("xarray.namedarray.parallelcompat.list_chunkmanagers", dict)
with pytest.raises(ImportError, match="no chunk managers available"):
guess_chunkmanager("foo")

def test_no_chunk_manager_available_but_known_manager_requested(
self, monkeypatch
) -> None:
monkeypatch.setattr("xarray.namedarray.parallelcompat.list_chunkmanagers", dict)
with pytest.raises(ImportError, match="chunk manager 'dask' is not available"):
guess_chunkmanager("dask")

@requires_dask
Expand Down

0 comments on commit dead70f

Please sign in to comment.