diff --git a/doc/internals/chunked-arrays.rst b/doc/internals/chunked-arrays.rst index ba7ce72c834..eb0a8b197e7 100644 --- a/doc/internals/chunked-arrays.rst +++ b/doc/internals/chunked-arrays.rst @@ -91,7 +91,8 @@ Once the chunkmanager subclass has been registered, xarray objects wrapping the The latter two methods ultimately call the chunkmanager's implementation of ``.from_array``, to which they pass the ``from_array_kwargs`` dict. The ``chunked_array_type`` kwarg selects which registered chunkmanager subclass to dispatch to. It defaults to ``'dask'`` if Dask is installed, otherwise it defaults to whichever chunkmanager is registered if only one is registered. -If multiple chunkmanagers are registered it will raise an error by default. +If multiple chunkmanagers are registered, the ``chunk_manager`` configuration option (which can be set using :py:func:`set_options`) +will be used to determine which chunkmanager to use, defaulting to ``'dask'``. Parallel processing without chunks ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 67e58e52619..c8f3a40e87f 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -22,7 +22,8 @@ v2024.07.1 (unreleased) New Features ~~~~~~~~~~~~ - +- Make chunk manager an option in ``set_options`` (:pull:`9362`). + By `Tom White `_. Breaking changes ~~~~~~~~~~~~~~~~ diff --git a/xarray/core/options.py b/xarray/core/options.py index f5614104357..f31413a2a1a 100644 --- a/xarray/core/options.py +++ b/xarray/core/options.py @@ -10,6 +10,7 @@ Options = Literal[ "arithmetic_join", + "chunk_manager", "cmap_divergent", "cmap_sequential", "display_max_rows", @@ -36,6 +37,7 @@ class T_Options(TypedDict): arithmetic_broadcast: bool arithmetic_join: Literal["inner", "outer", "left", "right", "exact"] + chunk_manager: str cmap_divergent: str | Colormap cmap_sequential: str | Colormap display_max_rows: int @@ -62,6 +64,7 @@ class T_Options(TypedDict): OPTIONS: T_Options = { "arithmetic_broadcast": True, "arithmetic_join": "inner", + "chunk_manager": "dask", "cmap_divergent": "RdBu_r", "cmap_sequential": "viridis", "display_max_rows": 12, @@ -172,7 +175,9 @@ class set_options: - "override": if indexes are of same size, rewrite indexes to be those of the first object with that dimension. Indexes for the same dimension must have the same size in all objects. - + chunk_manager : str, default: "dask" + Chunk manager to use for chunked array computations when multiple + options are installed. cmap_divergent : str or matplotlib.colors.Colormap, default: "RdBu_r" Colormap to use for divergent data plots. If string, must be matplotlib built-in colormap. Can also be a Colormap object diff --git a/xarray/namedarray/parallelcompat.py b/xarray/namedarray/parallelcompat.py index 3394aca9a5e..b90e0f99782 100644 --- a/xarray/namedarray/parallelcompat.py +++ b/xarray/namedarray/parallelcompat.py @@ -14,6 +14,7 @@ import numpy as np +from xarray.core.options import OPTIONS from xarray.core.utils import emit_user_level_warning from xarray.namedarray.pycompat import is_chunked_array @@ -101,8 +102,8 @@ def guess_chunkmanager( # use the only option available manager = next(iter(chunkmanagers.keys())) else: - # default to trying to use dask - manager = "dask" + # use the one in options (default dask) + manager = OPTIONS["chunk_manager"] if isinstance(manager, str): if manager not in chunkmanagers: diff --git a/xarray/tests/test_parallelcompat.py b/xarray/tests/test_parallelcompat.py index dbe40be710c..67c68aac534 100644 --- a/xarray/tests/test_parallelcompat.py +++ b/xarray/tests/test_parallelcompat.py @@ -6,6 +6,7 @@ import numpy as np import pytest +from xarray import set_options from xarray.core.types import T_Chunks, T_DuckArray, T_NormalizedChunks from xarray.namedarray._typing import _Chunks from xarray.namedarray.daskmanager import DaskManager @@ -152,6 +153,11 @@ def test_get_chunkmanger(self, register_dummy_chunkmanager) -> None: chunkmanager = guess_chunkmanager("dummy") assert isinstance(chunkmanager, DummyChunkManager) + def test_get_chunkmanger_via_set_options(self, register_dummy_chunkmanager) -> None: + with set_options(chunk_manager="dummy"): + 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"): guess_chunkmanager("foo")