New algorithm for forward filling

xarray/core/dask_array_ops.py

@@ -57,24 +57,41 @@ def push(array, n, axis):
     """
     Dask-aware bottleneck.push
     """
-    from bottleneck import push
+    import bottleneck
+    import dask.array as da
+    import numpy as np
 
-    if len(array.chunks[axis]) > 1 and n is not None and n < array.shape[axis]:
-        raise NotImplementedError(
-            "Cannot fill along a chunked axis when limit is not None."
-            "Either rechunk to a single chunk along this axis or call .compute() or .load() first."
-        )
-    if all(c == 1 for c in array.chunks[axis]):
-        array = array.rechunk({axis: 2})
-    pushed = array.map_blocks(push, axis=axis, n=n, dtype=array.dtype, meta=array._meta)
-    if len(array.chunks[axis]) > 1:
-        pushed = pushed.map_overlap(
-            push,
-            axis=axis,
-            n=n,
-            depth={axis: (1, 0)},
-            boundary="none",
-            dtype=array.dtype,
-            meta=array._meta,
+    def _fill_with_last_one(a, b):
+        # cumreduction apply the push func over all the blocks first so, the only missing part is filling
+        # the missing values using the last data of the previous chunk
+        if isinstance(a, np.ma.masked_array) or isinstance(b, np.ma.masked_array):
+            a, b = np.ma.getdata(a), np.ma.getdata(b)
+            values = np.where(~np.isnan(b), b, a)
+            return np.ma.masked_array(values, mask=np.ma.getmaskarray(b))
+
+        return np.where(~np.isnan(b), b, a)
+
+    if n is not None and n > 0:
+        arange = da.broadcast_to(
+            da.arange(
+                array.shape[axis], chunks=array.chunks[axis], dtype=array.dtype
+            ).reshape(
+                tuple(size if i == axis else 1 for i, size in enumerate(array.shape))
+            ),
+            array.shape,
+            array.chunks,
         )
-    return pushed
+        valid_arange = da.where(da.notnull(array), arange, np.nan)
+        valid_limits = (arange - push(valid_arange, None, axis)) <= n
+        # omit the forward fill that violate the limit
+        return da.where(valid_limits, push(array, None, axis), np.nan)
+
+    # The method parameter makes that the tests for python 3.7 fails.
+    return da.reductions.cumreduction(
+        func=bottleneck.push,
+        binop=_fill_with_last_one,
+        ident=np.nan,
+        x=array,
+        axis=axis,
+        dtype=array.dtype,
+    )

xarray/tests/test_duck_array_ops.py

@@ -884,16 +884,18 @@ def test_push_dask():
     import bottleneck
     import dask.array
 
-    array = np.array([np.nan, np.nan, np.nan, 1, 2, 3, np.nan, np.nan, 4, 5, np.nan, 6])
-    expected = bottleneck.push(array, axis=0)
-    for c in range(1, 11):
+    array = np.array([np.nan, 1, 2, 3, np.nan, np.nan, np.nan, np.nan, 4, 5, np.nan, 6])
+
+    for n in [None, 1, 2, 3, 4, 5, 11]:
+        expected = bottleneck.push(array, axis=0, n=n)
+        for c in range(1, 11):
+            with raise_if_dask_computes():
+                actual = push(dask.array.from_array(array, chunks=c), axis=0, n=n)
+            np.testing.assert_equal(actual, expected)
+
+        # some chunks of size-1 with NaN
         with raise_if_dask_computes():
-            actual = push(dask.array.from_array(array, chunks=c), axis=0, n=None)
+            actual = push(
+                dask.array.from_array(array, chunks=(1, 2, 3, 2, 2, 1, 1)), axis=0, n=n
+            )
         np.testing.assert_equal(actual, expected)
-
-    # some chunks of size-1 with NaN
-    with raise_if_dask_computes():
-        actual = push(
-            dask.array.from_array(array, chunks=(1, 2, 3, 2, 2, 1, 1)), axis=0, n=None
-        )
-    np.testing.assert_equal(actual, expected)

xarray/tests/test_missing.py

@@ -451,10 +451,6 @@ def test_ffill_bfill_dask(method):
     expected = numpy_method("time", limit=3)
     assert_equal(actual, expected)
 
-    # limit < axis size
-    with pytest.raises(NotImplementedError):
-        actual = dask_method("x", limit=2)
-
     # limit > axis size
     with raise_if_dask_computes():
         actual = dask_method("x", limit=41)