Generalize cumulative reduction (scan) to non-dask types (pydata#8019)

* add scan to ChunkManager ABC * implement scan for dask using cumreduction * generalize push to work for non-dask chunked arrays * whatsnew * fix importerror * Allow arbitrary kwargs Co-authored-by: Deepak Cherian <[email protected]> * Type hint return value of T_ChunkedArray Co-authored-by: Illviljan <[email protected]> * Type hint return value of Dask array * ffill -> bfill in doc/whats-new.rst Co-authored-by: Deepak Cherian <[email protected]> * hopefully fix docs warning --------- Co-authored-by: Deepak Cherian <[email protected]> Co-authored-by: Illviljan <[email protected]>
headtr1ck · Dec 18, 2023 · 766da34 · 766da34
1 parent 2971994
commit 766da34
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 0 deletions.
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
@@ -589,6 +589,10 @@ Internal Changes
 
 - :py:func:`as_variable` now consistently includes the variable name in any exceptions
   raised. (:pull:`7995`). By `Peter Hill <https://github.com/ZedThree>`_
+- Redirect cumulative reduction functions internally through the :py:class:`ChunkManagerEntryPoint`,
+  potentially allowing :py:meth:`~xarray.DataArray.ffill` and :py:meth:`~xarray.DataArray.bfill` to
+  use non-dask chunked array types.
+  (:pull:`8019`) By `Tom Nicholas <https://github.com/TomNicholas>`_.
 - :py:func:`encode_dataset_coordinates` now sorts coordinates automatically assigned to
   `coordinates` attributes during serialization (:issue:`8026`, :pull:`8034`).
   `By Ian Carroll <https://github.com/itcarroll>`_.

diff --git a/xarray/core/daskmanager.py b/xarray/core/daskmanager.py
@@ -97,6 +97,28 @@ def reduction(
             keepdims=keepdims,
         )
 
+    def scan(
+        self,
+        func: Callable,
+        binop: Callable,
+        ident: float,
+        arr: T_ChunkedArray,
+        axis: int | None = None,
+        dtype: np.dtype | None = None,
+        **kwargs,
+    ) -> DaskArray:
+        from dask.array.reductions import cumreduction
+
+        return cumreduction(
+            func,
+            binop,
+            ident,
+            arr,
+            axis=axis,
+            dtype=dtype,
+            **kwargs,
+        )
+
     def apply_gufunc(
         self,
         func: Callable,

diff --git a/xarray/core/parallelcompat.py b/xarray/core/parallelcompat.py
@@ -403,6 +403,43 @@ def reduction(
         """
         raise NotImplementedError()
 
+    def scan(
+        self,
+        func: Callable,
+        binop: Callable,
+        ident: float,
+        arr: T_ChunkedArray,
+        axis: int | None = None,
+        dtype: np.dtype | None = None,
+        **kwargs,
+    ) -> T_ChunkedArray:
+        """
+        General version of a 1D scan, also known as a cumulative array reduction.
+
+        Used in ``ffill`` and ``bfill`` in xarray.
+
+        Parameters
+        ----------
+        func: callable
+            Cumulative function like np.cumsum or np.cumprod
+        binop: callable
+            Associated binary operator like ``np.cumsum->add`` or ``np.cumprod->mul``
+        ident: Number
+            Associated identity like ``np.cumsum->0`` or ``np.cumprod->1``
+        arr: dask Array
+        axis: int, optional
+        dtype: dtype
+
+        Returns
+        -------
+        Chunked array
+
+        See also
+        --------
+        dask.array.cumreduction
+        """
+        raise NotImplementedError()
+
     @abstractmethod
     def apply_gufunc(
         self,