Load nonindex coords ahead of concat()

doc/whats-new.rst

@@ -116,7 +116,13 @@ Bug fixes
   ``rtol`` arguments when called on ``DataArray`` objects.
   By `Stephan Hoyer <https://github.com/shoyer>`_.
 
-- Xarray ``quantile`` methods now properly raise a ``TypeError`` when applied to
+- :py:func:`~xarray.concat` was computing multiple times coordinates that are
+  not index and not in memory (e.g. dask-based); :py:func:`~xarray.open_mfdataset`
+  was loading them multiple times from disk. Now, both functions will instead
+  load them once and store them as numpy arrays (:issue:`1521`).
+  By `Guido Imperiale <https://github.com/crusaderky>`_.
+
+- xarray ``quantile`` methods now properly raise a ``TypeError`` when applied to
   objects with data stored as ``dask`` arrays (:issue:`1529`).
   By `Joe Hamman <https://github.com/jhamman>`_.
 
@@ -2032,4 +2038,4 @@ Miles.
 v0.1 (2 May 2014)
 -----------------
 
-Initial release.
+Initial release.

xarray/core/combine.py

@@ -195,6 +195,17 @@ def differs(vname):
     return concat_over
 
 
+def _load_coords(dataset):
+    """Load into memory any non-index coords. Preserve original.
+    """
+    if all(coord._in_memory for coord in dataset.coords.values()):
+        return dataset
+    dataset = dataset.copy()
+    for coord in dataset.coords.values():
+        coord.load()
+    return dataset
+
+
 def _dataset_concat(datasets, dim, data_vars, coords, compat, positions):
     """
     Concatenate a sequence of datasets along a new or existing dimension
@@ -208,6 +219,8 @@ def _dataset_concat(datasets, dim, data_vars, coords, compat, positions):
     dim, coord = _calc_concat_dim_coord(dim)
     datasets = [as_dataset(ds) for ds in datasets]
     datasets = align(*datasets, join='outer', copy=False, exclude=[dim])
+    # TODO: compute dask coords with a single invocation of dask.compute()
+    datasets = [_load_coords(ds) for ds in datasets]
 
     concat_over = _calc_concat_over(datasets, dim, data_vars, coords)
 

xarray/tests/test_dask.py

@@ -228,6 +228,22 @@ def test_lazy_array(self):
         actual = xr.concat([v[:2], v[2:]], 'x')
         self.assertLazyAndAllClose(u, actual)
 
+    def test_concat_loads_coords(self):
+        # Test that concat() computes dask-based, non-index
+        # coordinates exactly once and loads them in the output,
+        # while leaving the input unaltered.
+        y = build_dask_array()
+        ds1 = Dataset(coords={'x': [1], 'y': ('x', y)})
+        ds2 = Dataset(coords={'x': [1], 'y': ('x', [2.0])})
+        assert kernel_call_count == 0
+        ds3 = xr.concat([ds1, ds2], dim='z')
+        # BUG fixed in #1532 where getattr('to_dataset')
+        # will cause non-index coords to be computed.
+        assert kernel_call_count == 2
+        assert ds1['y'].data is y
+        assert isinstance(ds3['y'].data, np.ndarray)
+        assert ds3['y'].values.tolist() == [[1.0], [2.0]]
+
     def test_groupby(self):
         u = self.eager_array
         v = self.lazy_array