Repository: dask/dask
Pull request #11501 (2024-11-08 14:29:33)

Add push for ffill like operations
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`

Moving this over from xarray since it makes sene that this lives in Dask.

see https://github.com/pydata/xarray/pull/9712


Patch:
diff --git a/continuous_integration/environment-3.10.yaml b/continuous_integration/environment-3.10.yaml
index 25e1a5078..7e7f8a514 100644
--- a/continuous_integration/environment-3.10.yaml
+++ b/continuous_integration/environment-3.10.yaml
@@ -48,6 +48,7 @@ dependencies:
   - boto3
   - botocore
   - bokeh
+  - bottleneck
   - httpretty
   - aiohttp
   - s3fs>=2021.9.0
diff --git a/continuous_integration/environment-3.11.yaml b/continuous_integration/environment-3.11.yaml
index dfa1de685..f248258e2 100644
--- a/continuous_integration/environment-3.11.yaml
+++ b/continuous_integration/environment-3.11.yaml
@@ -46,6 +46,7 @@ dependencies:
   - boto3
   - botocore
   - bokeh
+  - bottleneck
   - httpretty
   - aiohttp
   - s3fs>=2021.9.0
diff --git a/continuous_integration/environment-3.12.yaml b/continuous_integration/environment-3.12.yaml
index 9c9cf358d..6db5ef8c9 100644
--- a/continuous_integration/environment-3.12.yaml
+++ b/continuous_integration/environment-3.12.yaml
@@ -46,6 +46,7 @@ dependencies:
   - boto3
   - botocore
   - bokeh
+  - bottleneck
   - httpretty
   - aiohttp
   - s3fs>=2021.9.0
diff --git a/continuous_integration/environment-mindeps-optional.yaml b/continuous_integration/environment-mindeps-optional.yaml
index 6681cc765..9765d0020 100644
--- a/continuous_integration/environment-mindeps-optional.yaml
+++ b/continuous_integration/environment-mindeps-optional.yaml
@@ -47,6 +47,7 @@ dependencies:
   - tblib=1.6.0
   - tiledb-py=0.12
   - zarr=2.12.0
+  - bottleneck=1.3.7
   - pip
   - pip:
       # optional dependencies pulled in by pip install dask[distributed]
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
index 7a8d2e767..02a5fb3fb 100644
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -99,7 +99,7 @@ try:
     from dask.array.gufunc import apply_gufunc, as_gufunc, gufunc
     from dask.array.numpy_compat import moveaxis, rollaxis
     from dask.array.optimization import optimize
-    from dask.array.overlap import map_overlap
+    from dask.array.overlap import map_overlap, push
     from dask.array.percentile import percentile
     from dask.array.rechunk import rechunk
     from dask.array.reductions import (
diff --git a/dask/array/overlap.py b/dask/array/overlap.py
index 43c4ba7a2..bd009f685 100644
--- a/dask/array/overlap.py
+++ b/dask/array/overlap.py
@@ -9,11 +9,14 @@ import numpy as np
 from tlz import concat, get, partial
 from tlz.curried import map
 
+from dask._compatibility import import_optional_dependency
 from dask.array import chunk
 from dask.array._shuffle import _calculate_new_chunksizes
-from dask.array.core import Array, concatenate, map_blocks, unify_chunks
-from dask.array.creation import empty_like, full_like, repeat
+from dask.array.core import Array, broadcast_to, concatenate, map_blocks, unify_chunks
+from dask.array.creation import arange, empty_like, full_like, repeat
 from dask.array.numpy_compat import normalize_axis_tuple
+from dask.array.reductions import cumreduction
+from dask.array.routines import notnull, where
 from dask.base import tokenize
 from dask.highlevelgraph import HighLevelGraph
 from dask.layers import ArrayOverlapLayer
@@ -880,3 +883,53 @@ def sliding_window_view(x, window_shape, axis=None, automatic_rechunk=True):
         window_shape=window_shape,
         axis=axis,
     )
+
+
+def push(array, n, axis):
+    """
+    Dask-version of bottleneck.push
+
+    .. note::
+
+        Requires bottleneck to be installed.
+    """
+    import_optional_dependency("bottleneck", min_version="1.3.7")
+
+    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
+        return np.where(~np.isnan(b), b, a)
+
+    if n is not None and 0 < n < array.shape[axis] - 1:
+        arr = broadcast_to(
+            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,
+        )
+        valid_arange = where(notnull(array), arr, np.nan)
+        valid_limits = (arr - push(valid_arange, None, axis)) <= n
+        # omit the forward fill that violate the limit
+        return where(valid_limits, push(array, None, axis), np.nan)
+
+    # The method parameter makes that the tests for python 3.7 fails.
+    return cumreduction(
+        func=_push,
+        binop=_fill_with_last_one,
+        ident=np.nan,
+        x=array,
+        axis=axis,
+        dtype=array.dtype,
+    )
+
+
+def _push(array, n: int | None = None, axis: int = -1):
+    # work around for bottleneck 178
+    limit = n if n is not None else array.shape[axis]
+
+    import bottleneck as bn
+
+    return bn.push(array, limit, axis)
diff --git a/docs/source/array-api.rst b/docs/source/array-api.rst
index 0f5b579bb..30c279ac9 100644
--- a/docs/source/array-api.rst
+++ b/docs/source/array-api.rst
@@ -180,6 +180,7 @@ Top level functions
    outer
    pad
    percentile
+   push
    ~core.PerformanceWarning
    piecewise
    positive
diff --git a/docs/source/install.rst b/docs/source/install.rst
index da138b5e6..97ad1c44b 100644
--- a/docs/source/install.rst
+++ b/docs/source/install.rst
@@ -106,6 +106,8 @@ These optional dependencies and their minimum supported versions are listed belo
 +==================+=================+=========================================================================================================+
 | `bokeh`_         | ``>=3.1.0``     | Generate profiles of Dask execution (required for ``dask.diagnostics``)                                 |
 +------------------+-----------------+---------------------------------------------------------------------------------------------------------+
+| `bottleneck`_    | ``>=1.3.7``     | Used for dask arrays ``push`` implementation                                                            |
++------------------+-----------------+---------------------------------------------------------------------------------------------------------+
 | `cachey`_        | ``>=0.1.1``     | Use caching for computation                                                                             |
 +------------------+-----------------+---------------------------------------------------------------------------------------------------------+
 | `cityhash`_      | ``>=0.2.4``     | Use CityHash and FarmHash hash functions for array hashing (~2x faster than MurmurHash)                 |
@@ -193,6 +195,7 @@ See the :ref:`section on testing <develop-test>` in the Development Guidelines f
 .. _Anaconda distribution: https://www.anaconda.com/download/
 .. _s3fs: https://s3fs.readthedocs.io/en/latest/
 .. _bokeh: https://bokeh.org/
+.. _bottleneck:: https://bottleneck.readthedocs.io/en/latest/index.html
 .. _cachey: https://github.com/dask/cachey
 .. _cityhash: https://github.com/escherba/python-cityhash
 .. _crick: https://github.com/dask/crick


Test patch:
diff --git a/dask/array/tests/test_overlap.py b/dask/array/tests/test_overlap.py
index 2c1b2a7ee..3592ccae3 100644
--- a/dask/array/tests/test_overlap.py
+++ b/dask/array/tests/test_overlap.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import pytest
 
+from dask.array import from_array
 from dask.base import collections_to_dsk
 
 pytest.importorskip("numpy")
@@ -19,6 +20,7 @@ from dask.array.overlap import (
     overlap,
     overlap_internal,
     periodic,
+    push,
     reflect,
     trim_internal,
 )
@@ -725,6 +727,22 @@ def test_overlap_few_dimensions():
     assert len(c.dask) < 10 * len(a.dask)
 
 
+def test_push():
+    bottleneck = pytest.importorskip("bottleneck")
+
+    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):
+            actual = push(from_array(array, chunks=c), axis=0, n=n)
+            np.testing.assert_equal(actual, expected)
+
+        # some chunks of size-1 with NaN
+        actual = push(from_array(array, chunks=(1, 2, 3, 2, 2, 1, 1)), axis=0, n=n)
+        np.testing.assert_equal(actual, expected)
+
+
 @pytest.mark.parametrize("boundary", ["reflect", "periodic", "nearest", "none"])
 def test_trim_boundary(boundary):
     x = da.from_array(np.arange(24).reshape(4, 6), chunks=(2, 3))