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

Make chunk manager an option in set_options #9362

Merged
merged 6 commits into from
Aug 16, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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: 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
Loading