Skip to content

Commit

Permalink
Attempt to fix indexing for Dask
Browse files Browse the repository at this point in the history
This is a naive attempt to make `isel` work with Dask

Known limitation: it triggers the computation.
  • Loading branch information
bzah committed Oct 18, 2021
1 parent 07de257 commit 71d5ff6
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions xarray/core/indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from contextlib import suppress
from datetime import timedelta
from typing import Any, Callable, Iterable, List, Optional, Tuple, Union
import dask.array as da

import numpy as np
import pandas as pd
Expand Down Expand Up @@ -307,7 +308,7 @@ def __init__(self, key):
for k in key:
if isinstance(k, slice):
k = as_integer_slice(k)
elif isinstance(k, np.ndarray):
elif isinstance(k, np.ndarray) or isinstance(k, da.Array):
if not np.issubdtype(k.dtype, np.integer):
raise TypeError(
f"invalid indexer array, does not have integer dtype: {k!r}"
Expand All @@ -320,7 +321,10 @@ def __init__(self, key):
"invalid indexer key: ndarray arguments "
f"have different numbers of dimensions: {ndims}"
)
k = np.asarray(k, dtype=np.int64)
if isinstance(k, da.Array):
k = da.asarray(k, dtype=np.int64)
else:
k = np.asarray(k, dtype=np.int64)
else:
raise TypeError(
f"unexpected indexer type for {type(self).__name__}: {k!r}"
Expand Down Expand Up @@ -973,7 +977,6 @@ def _arrayize_vectorized_indexer(indexer, shape):

def _dask_array_with_chunks_hint(array, chunks):
"""Create a dask array using the chunks hint for dimensions of size > 1."""
import dask.array as da

if len(chunks) < array.ndim:
raise ValueError("not enough chunks in hint")
Expand Down

0 comments on commit 71d5ff6

Please sign in to comment.