Ensure explicitly indexed arrays are preserved (#3027)

* Ensure indexing explicitly indexed arrays don't leak out.

Previously, indexing an ImplicitToExplicitIndexingAdapter object could directly
return an ExplicitlyIndexed object, which could not be indexed normally. This
resulted in broken behavior with dask's new `_meta` attribute.

This change almost but not entirely fixes xarray on dask master. There are
still errors raised inside two tests from dask's `blockwise_meta` helper
function:

>       return meta.astype(dtype)
E       AttributeError: 'ImplicitToExplicitIndexingAdapter' object has no attribute 'astype'

* Set meta in dask.array.from_array

Author: shoyer

.travis.yml

@@ -16,10 +16,10 @@ matrix:
   - env:
     - CONDA_ENV=py36
     - EXTRA_FLAGS="--run-flaky --run-network-tests"
+  - env: CONDA_ENV=py36-dask-dev
   - env: CONDA_ENV=py36-pandas-dev
   - env: CONDA_ENV=py36-rasterio
   - env: CONDA_ENV=py36-zarr-dev
-  - env: CONDA_ENV=py36-dask-dev
   - env: CONDA_ENV=docs
   - env: CONDA_ENV=lint
   - env: CONDA_ENV=typing

xarray/core/indexing.py

@@ -453,7 +453,13 @@ def __array__(self, dtype=None):
 
     def __getitem__(self, key):
         key = expanded_indexer(key, self.ndim)
-        return self.array[self.indexer_cls(key)]
+        result = self.array[self.indexer_cls(key)]
+        if isinstance(result, ExplicitlyIndexed):
+            return type(self)(result, self.indexer_cls)
+        else:
+            # Sometimes explicitly indexed arrays return NumPy arrays or
+            # scalars.
+            return result
 
 
 class LazilyOuterIndexedArray(ExplicitlyIndexedNDArrayMixin):

xarray/core/variable.py

@@ -3,6 +3,7 @@
 import typing
 from collections import OrderedDict, defaultdict
 from datetime import timedelta
+from distutils.version import LooseVersion
 
 import numpy as np
 import pandas as pd
@@ -870,6 +871,7 @@ def chunk(self, chunks=None, name=None, lock=False):
         -------
         chunked : xarray.Variable
         """
+        import dask
         import dask.array as da
 
         if utils.is_dict_like(chunks):
@@ -892,7 +894,17 @@ def chunk(self, chunks=None, name=None, lock=False):
             # https://github.com/dask/dask/issues/2883
             data = indexing.ImplicitToExplicitIndexingAdapter(
                 data, indexing.OuterIndexer)
-            data = da.from_array(data, chunks, name=name, lock=lock)
+
+            # For now, assume that all arrays that we wrap with dask (including
+            # our lazily loaded backend array classes) should use NumPy array
+            # operations.
+            if LooseVersion(dask.__version__) > '1.2.2':
+                kwargs = dict(meta=np.ndarray)
+            else:
+                kwargs = dict()
+
+            data = da.from_array(
+                data, chunks, name=name, lock=lock, **kwargs)
 
         return type(self)(self.dims, data, self._attrs, self._encoding,
                           fastpath=True)

xarray/tests/test_indexing.py

@@ -505,13 +505,20 @@ def test_decompose_indexers(shape, indexer_mode, indexing_support):
 
 
 def test_implicit_indexing_adapter():
-    array = np.arange(10)
+    array = np.arange(10, dtype=np.int64)
     implicit = indexing.ImplicitToExplicitIndexingAdapter(
         indexing.NumpyIndexingAdapter(array), indexing.BasicIndexer)
     np.testing.assert_array_equal(array, np.asarray(implicit))
     np.testing.assert_array_equal(array, implicit[:])
 
 
+def test_implicit_indexing_adapter_copy_on_write():
+    array = np.arange(10, dtype=np.int64)
+    implicit = indexing.ImplicitToExplicitIndexingAdapter(
+        indexing.CopyOnWriteArray(array))
+    assert isinstance(implicit[:], indexing.ImplicitToExplicitIndexingAdapter)
+
+
 def test_outer_indexer_consistency_with_broadcast_indexes_vectorized():
     def nonzero(x):
         if isinstance(x, np.ndarray) and x.dtype.kind == 'b':