use da.pad when dealing with dask arrays

pydata · Jan 20, 2020 · 1e07dce · 1e07dce
1 parent 88320ef
commit 1e07dce
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/xarray/core/variable.py b/xarray/core/variable.py
@@ -1107,7 +1107,12 @@ def _shift_one_dim(self, dim, count, fill_value=dtypes.NA):
         dim_pad = (width, 0) if count >= 0 else (0, width)
         pads = [(0, 0) if d != dim else dim_pad for d in self.dims]
 
-        data = np.pad(
+        if isinstance(trimmed_data, dask_array_type):
+            pad_func = da.pad
+        else:
+            pad_func = np.pad
+
+        data = pad_func(
             trimmed_data.astype(dtype),
             pads,
             mode="constant",