Skip to content

Commit

Permalink
Make chunk manager an option in set_options (#9362)
Browse files Browse the repository at this point in the history
* Make chunk manager an option in `set_options`

* Add docs

* Add unit test

* Add what's new entry
  • Loading branch information
tomwhite authored Aug 16, 2024
1 parent 6b6a3bd commit 6166c87
Show file tree
Hide file tree
Showing 5 changed files with 19 additions and 5 deletions.
3 changes: 2 additions & 1 deletion doc/internals/chunked-arrays.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Expand Down
3 changes: 2 additions & 1 deletion doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,8 @@ v2024.07.1 (unreleased)

New Features
~~~~~~~~~~~~

- Make chunk manager an option in ``set_options`` (:pull:`9362`).
By `Tom White <https://github.com/tomwhite>`_.

Breaking changes
~~~~~~~~~~~~~~~~
Expand Down
7 changes: 6 additions & 1 deletion xarray/core/options.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@

Options = Literal[
"arithmetic_join",
"chunk_manager",
"cmap_divergent",
"cmap_sequential",
"display_max_rows",
Expand All @@ -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
Expand All @@ -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,
Expand Down Expand Up @@ -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
Expand Down
5 changes: 3 additions & 2 deletions xarray/namedarray/parallelcompat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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:
Expand Down
6 changes: 6 additions & 0 deletions xarray/tests/test_parallelcompat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")
Expand Down

0 comments on commit 6166c87

Please sign in to comment.