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

use the DaskIndexingAdapter for duck dask arrays #6414

Merged
merged 9 commits into from
Mar 27, 2022
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
15 changes: 10 additions & 5 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,8 @@ New Features

- Add a ``create_index=True`` parameter to :py:meth:`Dataset.stack` and
:py:meth:`DataArray.stack` so that the creation of multi-indexes is optional
(:pull:`5692`). By `Benoît Bovy <https://github.com/benbovy>`_.
(:pull:`5692`).
By `Benoît Bovy <https://github.com/benbovy>`_.
- Multi-index levels are now accessible through their own, regular coordinates
instead of virtual coordinates (:pull:`5692`).
By `Benoît Bovy <https://github.com/benbovy>`_.
Expand All @@ -33,7 +34,8 @@ Breaking changes
~~~~~~~~~~~~~~~~

- The Dataset and DataArray ``rename*`` methods do not implicitly add or drop
indexes. (:pull:`5692`). By `Benoît Bovy <https://github.com/benbovy>`_.
indexes. (:pull:`5692`).
By `Benoît Bovy <https://github.com/benbovy>`_.
- Many arguments like ``keep_attrs``, ``axis``, and ``skipna`` are now keyword
only for all reduction operations like ``.mean``.
By `Deepak Cherian <https://github.com/dcherian>`_, `Jimmy Westling <https://github.com/illviljan>`_.
Expand All @@ -47,12 +49,16 @@ Bug fixes

- Set ``skipna=None`` for all ``quantile`` methods (e.g. :py:meth:`Dataset.quantile`) and
ensure it skips missing values for float dtypes (consistent with other methods). This should
not change the behavior (:pull:`6303`). By `Mathias Hauser <https://github.com/mathause>`_.
not change the behavior (:pull:`6303`).
By `Mathias Hauser <https://github.com/mathause>`_.
- Many bugs fixed by the explicit indexes refactor, mainly related to multi-index (virtual)
coordinates. See the corresponding pull-request on GitHub for more details. (:pull:`5692`).
By `Benoît Bovy <https://github.com/benbovy>`_.
- Fixed "unhashable type" error trying to read NetCDF file with variable having its 'units'
attribute not ``str`` (e.g. ``numpy.ndarray``) (:issue:`6368`). By `Oleh Khoma <https://github.com/okhoma>`_.
attribute not ``str`` (e.g. ``numpy.ndarray``) (:issue:`6368`).
By `Oleh Khoma <https://github.com/okhoma>`_.
- Allow fancy indexing of duck dask arrays along multiple dimensions. (:pull:`6414`)
By `Justus Magin <https://github.com/keewis>`_.

Documentation
~~~~~~~~~~~~~
Expand All @@ -65,7 +71,6 @@ Internal Changes
corresponding pull-request on GitHub for more details. (:pull:`5692`).
By `Benoît Bovy <https://github.com/benbovy>`_.

.. _whats-new.2022.02.0:
.. _whats-new.2022.03.0:

v2022.03.0 (2 March 2022)
Expand Down
10 changes: 2 additions & 8 deletions xarray/core/indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,13 +25,7 @@

from . import duck_array_ops, nputils, utils
from .npcompat import DTypeLike
from .pycompat import (
dask_array_type,
dask_version,
integer_types,
is_duck_dask_array,
sparse_array_type,
)
from .pycompat import dask_version, integer_types, is_duck_dask_array, sparse_array_type
from .types import T_Xarray
from .utils import either_dict_or_kwargs, get_valid_numpy_dtype

Expand Down Expand Up @@ -682,7 +676,7 @@ def as_indexable(array):
return NumpyIndexingAdapter(array)
if isinstance(array, pd.Index):
return PandasIndexingAdapter(array)
if isinstance(array, dask_array_type):
if is_duck_dask_array(array):
return DaskIndexingAdapter(array)
if hasattr(array, "__array_function__"):
return NdArrayLikeIndexingAdapter(array)
Expand Down
40 changes: 31 additions & 9 deletions xarray/tests/test_units.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@

# make sure scalars are converted to 0d arrays so quantities can
# always be treated like ndarrays
unit_registry = pint.UnitRegistry(force_ndarray=True)
unit_registry = pint.UnitRegistry(force_ndarray_like=True)
Quantity = unit_registry.Quantity


Expand Down Expand Up @@ -1837,21 +1837,43 @@ def test_broadcast_equals(self, unit, dtype):

assert expected == actual

@pytest.mark.parametrize("dask", [False, pytest.param(True, marks=[requires_dask])])
@pytest.mark.parametrize(
"indices",
["variable", "indexers"],
(
pytest.param(4, id="single index"),
pytest.param([5, 2, 9, 1], id="multiple indices"),
pytest.param(
xr.Variable("x", np.linspace(0, 5, 10)),
{"x": 4},
id="single value-single indexer",
),
pytest.param(
xr.Variable("x", np.linspace(0, 5, 10)),
{"x": [5, 2, 9, 1]},
id="multiple values-single indexer",
),
pytest.param(
xr.Variable(("x", "y"), np.linspace(0, 5, 20).reshape(4, 5)),
{"x": 1, "y": 4},
id="single value-multiple indexers",
),
pytest.param(
xr.Variable(("x", "y"), np.linspace(0, 5, 20).reshape(4, 5)),
{"x": [0, 1, 2], "y": [0, 2, 4]},
id="multiple values-multiple indexers",
),
),
)
def test_isel(self, indices, dtype):
array = np.linspace(0, 5, 10).astype(dtype) * unit_registry.s
variable = xr.Variable("x", array)
def test_isel(self, variable, indexers, dask, dtype):
if dask:
variable = variable.chunk({dim: 2 for dim in variable.dims})
quantified = xr.Variable(
variable.dims, variable.data.astype(dtype) * unit_registry.s
)

expected = attach_units(
strip_units(variable).isel(x=indices), extract_units(variable)
strip_units(quantified).isel(indexers), extract_units(quantified)
)
actual = variable.isel(x=indices)
actual = quantified.isel(indexers)

assert_units_equal(expected, actual)
assert_identical(expected, actual)
Expand Down