Use numbagg for ffill

Author: max-sixty

xarray/backends/zarr.py

@@ -177,8 +177,8 @@ def _determine_zarr_chunks(enc_chunks, var_chunks, ndim, name, safe_chunks):
     # DESIGN CHOICE: do not allow multiple dask chunks on a single zarr chunk
     # this avoids the need to get involved in zarr synchronization / locking
     # From zarr docs:
-    #  "If each worker in a parallel computation is writing to a separate
-    #   region of the array, and if region boundaries are perfectly aligned
+    #  "If each worker in a parallel computation is writing to a
+    #   separate region of the array, and if region boundaries are perfectly aligned
     #   with chunk boundaries, then no synchronization is required."
     # TODO: incorporate synchronizer to allow writes from multiple dask
     # threads

xarray/core/dask_array_ops.py

@@ -59,10 +59,11 @@ def push(array, n, axis):
     """
     Dask-aware bottleneck.push
     """
-    import bottleneck
     import dask.array as da
     import numpy as np
 
+    from xarray.core.duck_array_ops import _push
+
     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
@@ -85,7 +86,8 @@ def _fill_with_last_one(a, b):
 
     # The method parameter makes that the tests for python 3.7 fails.
     return da.reductions.cumreduction(
-        func=bottleneck.push,
+        func=_push,
+        # func=bottleneck.push,
         binop=_fill_with_last_one,
         ident=np.nan,
         x=array,

xarray/core/duck_array_ops.py

@@ -31,8 +31,10 @@
 from numpy import concatenate as _concatenate
 from numpy.core.multiarray import normalize_axis_index  # type: ignore[attr-defined]
 from numpy.lib.stride_tricks import sliding_window_view  # noqa
+from packaging.version import Version
 
 from xarray.core import dask_array_ops, dtypes, nputils
+from xarray.core.options import OPTIONS
 from xarray.core.parallelcompat import get_chunked_array_type, is_chunked_array
 from xarray.core.pycompat import array_type, is_duck_dask_array
 from xarray.core.utils import is_duck_array, module_available
@@ -688,13 +690,43 @@ def least_squares(lhs, rhs, rcond=None, skipna=False):
         return nputils.least_squares(lhs, rhs, rcond=rcond, skipna=skipna)
 
 
-def push(array, n, axis):
-    from bottleneck import push
+def _push(array, n: int | None = None, axis: int = -1):
+    """
+    Use either bottleneck or numbagg depending on options & what's available
+    """
+    from xarray.core.nputils import NUMBAGG_VERSION, numbagg
+
+    if not OPTIONS["use_bottleneck"] and not OPTIONS["use_numbagg"]:
+        raise RuntimeError(
+            "ffill & bfill requires bottleneck or numbagg to be enabled."
+            " Call `xr.set_options(use_bottleneck=True)` or `xr.set_options(use_numbagg=True)` to enable one."
+        )
+    if OPTIONS["use_numbagg"] and NUMBAGG_VERSION is not None:
+        if NUMBAGG_VERSION < Version("0.6.2"):
+            warnings.warn(
+                f"numbagg >= 0.6.2 is required for bfill & ffill; {NUMBAGG_VERSION} is installed. We'll attempt with bottleneck instead."
+            )
+        else:
+            return numbagg.ffill(array, limit=n, axis=axis)
 
+    # 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)
+
+
+def push(array, n, axis):
+    if not OPTIONS["use_bottleneck"] and not OPTIONS["use_numbagg"]:
+        raise RuntimeError(
+            "ffill & bfill requires bottleneck or numbagg to be enabled."
+            " Call `xr.set_options(use_bottleneck=True)` or `xr.set_options(use_numbagg=True)` to enable one."
+        )
     if is_duck_dask_array(array):
         return dask_array_ops.push(array, n, axis)
     else:
-        return push(array, n, axis)
+        return _push(array, n, axis)
 
 
 def _first_last_wrapper(array, *, axis, op, keepdims):

xarray/core/missing.py

@@ -413,11 +413,6 @@ def _bfill(arr, n=None, axis=-1):
 
 def ffill(arr, dim=None, limit=None):
     """forward fill missing values"""
-    if not OPTIONS["use_bottleneck"]:
-        raise RuntimeError(
-            "ffill requires bottleneck to be enabled."
-            " Call `xr.set_options(use_bottleneck=True)` to enable it."
-        )
 
     axis = arr.get_axis_num(dim)
 
@@ -436,11 +431,6 @@ def ffill(arr, dim=None, limit=None):
 
 def bfill(arr, dim=None, limit=None):
     """backfill missing values"""
-    if not OPTIONS["use_bottleneck"]:
-        raise RuntimeError(
-            "bfill requires bottleneck to be enabled."
-            " Call `xr.set_options(use_bottleneck=True)` to enable it."
-        )
 
     axis = arr.get_axis_num(dim)
 

xarray/core/nputils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import warnings
+from typing import Callable
 
 import numpy as np
 import pandas as pd
@@ -25,14 +26,17 @@
     bn = np
     _BOTTLENECK_AVAILABLE = False
 
+NUMBAGG_VERSION: Version | None
+
 try:
     import numbagg
 
-    _HAS_NUMBAGG = Version(numbagg.__version__) >= Version("0.5.0")
+    v = getattr(numbagg, "__version__", "999")
+    NUMBAGG_VERSION = Version(v)
 except ImportError:
     # use numpy methods instead
     numbagg = np
-    _HAS_NUMBAGG = False
+    NUMBAGG_VERSION = None
 
 
 def _select_along_axis(values, idx, axis):
@@ -171,14 +175,15 @@ def __setitem__(self, key, value):
         self._array[key] = np.moveaxis(value, vindex_positions, mixed_positions)
 
 
-def _create_method(name, npmodule=np):
+def _create_method(name, npmodule=np) -> Callable:
     def f(values, axis=None, **kwargs):
         dtype = kwargs.get("dtype", None)
         bn_func = getattr(bn, name, None)
         nba_func = getattr(numbagg, name, None)
 
         if (
-            _HAS_NUMBAGG
+            NUMBAGG_VERSION is not None
+            and NUMBAGG_VERSION >= Version("0.5.0")
             and OPTIONS["use_numbagg"]
             and isinstance(values, np.ndarray)
             and nba_func is not None

xarray/core/rolling_exp.py

@@ -15,9 +15,9 @@
     import numbagg
     from numbagg import move_exp_nanmean, move_exp_nansum
 
-    _NUMBAGG_VERSION: Version | None = Version(numbagg.__version__)
+    NUMBAGG_VERSION: Version | None = Version(numbagg.__version__)
 except ImportError:
-    _NUMBAGG_VERSION = None
+    NUMBAGG_VERSION = None
 
 
 def _get_alpha(
@@ -100,17 +100,17 @@ def __init__(
         window_type: str = "span",
         min_weight: float = 0.0,
     ):
-        if _NUMBAGG_VERSION is None:
+        if NUMBAGG_VERSION is None:
             raise ImportError(
                 "numbagg >= 0.2.1 is required for rolling_exp but currently numbagg is not installed"
             )
-        elif _NUMBAGG_VERSION < Version("0.2.1"):
+        elif NUMBAGG_VERSION < Version("0.2.1"):
             raise ImportError(
-                f"numbagg >= 0.2.1 is required for rolling_exp but currently version {_NUMBAGG_VERSION} is installed"
+                f"numbagg >= 0.2.1 is required for rolling_exp but currently version {NUMBAGG_VERSION} is installed"
             )
-        elif _NUMBAGG_VERSION < Version("0.3.1") and min_weight > 0:
+        elif NUMBAGG_VERSION < Version("0.3.1") and min_weight > 0:
             raise ImportError(
-                f"numbagg >= 0.3.1 is required for `min_weight > 0` within `.rolling_exp` but currently version {_NUMBAGG_VERSION} is installed"
+                f"numbagg >= 0.3.1 is required for `min_weight > 0` within `.rolling_exp` but currently version {NUMBAGG_VERSION} is installed"
             )
 
         self.obj: T_DataWithCoords = obj
@@ -211,9 +211,9 @@ def std(self) -> T_DataWithCoords:
         Dimensions without coordinates: x
         """
 
-        if _NUMBAGG_VERSION is None or _NUMBAGG_VERSION < Version("0.4.0"):
+        if NUMBAGG_VERSION is None or NUMBAGG_VERSION < Version("0.4.0"):
             raise ImportError(
-                f"numbagg >= 0.4.0 is required for rolling_exp().std(), currently {_NUMBAGG_VERSION} is installed"
+                f"numbagg >= 0.4.0 is required for rolling_exp().std(), currently {NUMBAGG_VERSION} is installed"
             )
         dim_order = self.obj.dims
 
@@ -243,9 +243,9 @@ def var(self) -> T_DataWithCoords:
         Dimensions without coordinates: x
         """
 
-        if _NUMBAGG_VERSION is None or _NUMBAGG_VERSION < Version("0.4.0"):
+        if NUMBAGG_VERSION is None or NUMBAGG_VERSION < Version("0.4.0"):
             raise ImportError(
-                f"numbagg >= 0.4.0 is required for rolling_exp().var(), currently {_NUMBAGG_VERSION} is installed"
+                f"numbagg >= 0.4.0 is required for rolling_exp().var(), currently {NUMBAGG_VERSION} is installed"
             )
         dim_order = self.obj.dims
 
@@ -275,9 +275,9 @@ def cov(self, other: T_DataWithCoords) -> T_DataWithCoords:
         Dimensions without coordinates: x
         """
 
-        if _NUMBAGG_VERSION is None or _NUMBAGG_VERSION < Version("0.4.0"):
+        if NUMBAGG_VERSION is None or NUMBAGG_VERSION < Version("0.4.0"):
             raise ImportError(
-                f"numbagg >= 0.4.0 is required for rolling_exp().cov(), currently {_NUMBAGG_VERSION} is installed"
+                f"numbagg >= 0.4.0 is required for rolling_exp().cov(), currently {NUMBAGG_VERSION} is installed"
             )
         dim_order = self.obj.dims
 
@@ -308,9 +308,9 @@ def corr(self, other: T_DataWithCoords) -> T_DataWithCoords:
         Dimensions without coordinates: x
         """
 
-        if _NUMBAGG_VERSION is None or _NUMBAGG_VERSION < Version("0.4.0"):
+        if NUMBAGG_VERSION is None or NUMBAGG_VERSION < Version("0.4.0"):
             raise ImportError(
-                f"numbagg >= 0.4.0 is required for rolling_exp().cov(), currently {_NUMBAGG_VERSION} is installed"
+                f"numbagg >= 0.4.0 is required for rolling_exp().cov(), currently {NUMBAGG_VERSION} is installed"
             )
         dim_order = self.obj.dims
 

xarray/tests/__init__.py

@@ -53,7 +53,8 @@ def _importorskip(
         mod = importlib.import_module(modname)
         has = True
         if minversion is not None:
-            if Version(mod.__version__) < Version(minversion):
+            v = getattr(mod, "__version__", "999")
+            if Version(v) < Version(minversion):
                 raise ImportError("Minimum version not satisfied")
     except ImportError:
         has = False
@@ -89,6 +90,10 @@ def _importorskip(
 requires_scipy_or_netCDF4 = pytest.mark.skipif(
     not has_scipy_or_netCDF4, reason="requires scipy or netCDF4"
 )
+has_numbagg_or_bottleneck = has_numbagg or has_bottleneck
+requires_numbagg_or_bottleneck = pytest.mark.skipif(
+    not has_scipy_or_netCDF4, reason="requires scipy or netCDF4"
+)
 # _importorskip does not work for development versions
 has_pandas_version_two = Version(pd.__version__).major >= 2
 requires_pandas_version_two = pytest.mark.skipif(

xarray/tests/test_missing.py

@@ -24,6 +24,8 @@
     requires_bottleneck,
     requires_cftime,
     requires_dask,
+    requires_numbagg,
+    requires_numbagg_or_bottleneck,
     requires_scipy,
 )
 
@@ -407,33 +409,44 @@ def test_interpolate_dask_expected_dtype(dtype, method):
     assert da.dtype == da.compute().dtype
 
 
-@requires_bottleneck
+@requires_numbagg_or_bottleneck
 def test_ffill():
     da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
     expected = xr.DataArray(np.array([4, 5, 5], dtype=np.float64), dims="x")
     actual = da.ffill("x")
     assert_equal(actual, expected)
 
 
-def test_ffill_use_bottleneck():
+def test_ffill_use_bottleneck_numbagg():
     da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
     with xr.set_options(use_bottleneck=False):
-        with pytest.raises(RuntimeError):
-            da.ffill("x")
+        with xr.set_options(use_numbagg=False):
+            with pytest.raises(RuntimeError):
+                da.ffill("x")
 
 
 @requires_dask
 def test_ffill_use_bottleneck_dask():
     da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
     da = da.chunk({"x": 1})
-    with xr.set_options(use_bottleneck=False):
+    with xr.set_options(use_bottleneck=False, use_numbagg=False):
         with pytest.raises(RuntimeError):
             da.ffill("x")
 
 
+@requires_numbagg
+@requires_dask
+def test_ffill_use_numbagg_dask():
+    with xr.set_options(use_bottleneck=False):
+        da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
+        da = da.chunk(x=-1)
+        # Succeeds with a single chunk:
+        _ = da.ffill("x").compute()
+
+
 def test_bfill_use_bottleneck():
     da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
-    with xr.set_options(use_bottleneck=False):
+    with xr.set_options(use_bottleneck=False, use_numbagg=False):
         with pytest.raises(RuntimeError):
             da.bfill("x")
 
@@ -442,7 +455,7 @@ def test_bfill_use_bottleneck():
 def test_bfill_use_bottleneck_dask():
     da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x")
     da = da.chunk({"x": 1})
-    with xr.set_options(use_bottleneck=False):
+    with xr.set_options(use_bottleneck=False, use_numbagg=False):
         with pytest.raises(RuntimeError):
             da.bfill("x")