backend → engine

Author: dcherian

dask_groupby/core.py

@@ -32,14 +32,14 @@
 FinalResultsDict = Dict[str, Union["DaskArray", np.ndarray]]
 
 
-def _get_aggregate(backend):
-    if backend == "numba":
+def _get_aggregate(engine):
+    if engine == "numba":
         return npg.aggregate_numba.aggregate
-    elif backend == "numpy":
+    elif engine == "numpy":
         return npg.aggregate_numpy.aggregate
     else:
         raise ValueError(
-            "Expected backend to be one of ['numpy', 'numba']. Received {backend} instead."
+            "Expected engine to be one of ['numpy', 'numba']. Received {engine} instead."
         )
 
 
@@ -373,7 +373,7 @@ def chunk_argreduce(
     dtype=None,
     reindex: bool = False,
     isbin: bool = False,
-    backend: str = "numpy",
+    engine: str = "numpy",
 ) -> IntermediateDict:
     """
     Per-chunk arg reduction.
@@ -392,7 +392,7 @@ def chunk_argreduce(
         fill_value=fill_value,
         isbin=isbin,
         dtype=dtype,
-        backend=backend,
+        engine=engine,
     )
     if not np.isnan(results["groups"]).all():
         # will not work for empty groups...
@@ -422,7 +422,7 @@ def chunk_reduce(
     dtype=None,
     reindex: bool = False,
     isbin: bool = False,
-    backend: str = "numpy",
+    engine: str = "numpy",
     kwargs=None,
 ) -> IntermediateDict:
     """
@@ -533,7 +533,7 @@ def chunk_reduce(
                     **kw,
                 )
             else:
-                result = _get_aggregate(backend)(
+                result = _get_aggregate(engine)(
                     group_idx,
                     array,
                     axis=-1,
@@ -648,11 +648,11 @@ def _npg_aggregate(
     group_ndim: int,
     fill_value: Any = None,
     min_count: Optional[int] = None,
-    backend: str = "numpy",
+    engine: str = "numpy",
     finalize_kwargs: Optional[Mapping] = None,
 ) -> FinalResultsDict:
     """Final aggregation step of tree reduction"""
-    results = _npg_combine(x_chunk, agg, axis, keepdims, group_ndim, backend)
+    results = _npg_combine(x_chunk, agg, axis, keepdims, group_ndim, engine)
     return _finalize_results(
         results, agg, axis, expected_groups, fill_value, min_count, finalize_kwargs
     )
@@ -664,7 +664,7 @@ def _npg_combine(
     axis: Sequence,
     keepdims: bool,
     group_ndim: int,
-    backend: str,
+    engine: str,
 ) -> IntermediateDict:
     """Combine intermediates step of tree reduction."""
     from dask.array.core import _concatenate2
@@ -722,7 +722,7 @@ def _conc2(key1, key2=None, axis=None) -> np.ndarray:
             expected_groups=None,
             fill_value=agg.fill_value["intermediate"][slicer],
             dtype=agg.dtype,
-            backend=backend,
+            engine=engine,
         )
 
         if agg.chunk[-1] == "nanlen":
@@ -737,7 +737,7 @@ def _conc2(key1, key2=None, axis=None) -> np.ndarray:
                     expected_groups=None,
                     fill_value=(0,),
                     dtype=np.intp,
-                    backend=backend,
+                    engine=engine,
                 )["intermediates"][0]
             )
 
@@ -762,7 +762,7 @@ def _conc2(key1, key2=None, axis=None) -> np.ndarray:
                     axis=axis,
                     expected_groups=None,
                     fill_value=fv,
-                    backend=backend,
+                    engine=engine,
                 )
                 results["intermediates"].append(*_results["intermediates"])
                 results["groups"] = _results["groups"]
@@ -812,7 +812,7 @@ def groupby_agg(
     method: str = "mapreduce",
     min_count: Optional[int] = None,
     isbin: bool = False,
-    backend: str = "numpy",
+    engine: str = "numpy",
     finalize_kwargs: Optional[Mapping] = None,
 ) -> Tuple["DaskArray", Union[np.ndarray, "DaskArray"]]:
 
@@ -851,7 +851,7 @@ def groupby_agg(
             fill_value=agg.fill_value["intermediate"],
             isbin=isbin,
             reindex=split_out > 1,
-            backend=backend,
+            engine=engine,
         ),
         inds,
         array,
@@ -887,7 +887,7 @@ def groupby_agg(
         group_ndim=by.ndim,
         fill_value=fill_value,
         min_count=min_count,
-        backend=backend,
+        engine=engine,
         finalize_kwargs=finalize_kwargs,
     )
 
@@ -904,7 +904,7 @@ def groupby_agg(
                 expected_groups=expected_agg,
                 **agg_kwargs,
             ),
-            combine=partial(_npg_combine, agg=agg, group_ndim=by.ndim, backend=backend),
+            combine=partial(_npg_combine, agg=agg, group_ndim=by.ndim, engine=engine),
             name=f"{name}-reduce",
             dtype=array.dtype,
             axis=axis,
@@ -1015,7 +1015,7 @@ def groupby_reduce(
     min_count: Optional[int] = None,
     split_out: int = 1,
     method: str = "mapreduce",
-    backend: str = "numpy",
+    engine: str = "numpy",
     finalize_kwargs: Optional[Mapping] = None,
 ) -> Tuple["DaskArray", Union[np.ndarray, "DaskArray"]]:
     """
@@ -1069,8 +1069,8 @@ def groupby_reduce(
             where the group labels repeat at regular intervals like 'hour',
             'month', dayofyear' etc. Optimize chunking ``array`` for this
             method by first rechunking using ``rechunk_for_cohorts``.
-    backend : {"numpy", "numba"}, optional, default: ``"numpy"``
-        Backend for ``numpy_groupies``.
+    engine : {"numpy", "numba"}, optional, default: ``"numpy"``
+        Engine for ``numpy_groupies``.
     finalize_kwargs : dict, optional
         Kwargs passed to finalize the reduction such as ``ddof`` for var, std.
 
@@ -1252,7 +1252,7 @@ def groupby_reduce(
             fill_value=fill_value,
             min_count=min_count,
             isbin=isbin,
-            backend=backend,
+            engine=engine,
             finalize_kwargs=finalize_kwargs,
         )
         if method == "cohorts":

dask_groupby/xarray.py

@@ -58,7 +58,7 @@ def xarray_reduce(
     split_out: int = 1,
     fill_value=None,
     method: str = "mapreduce",
-    backend: str = "numpy",
+    engine: str = "numpy",
     keep_attrs: bool = True,
     skipna: Optional[bool] = None,
     min_count: Optional[int] = None,
@@ -110,8 +110,8 @@ def xarray_reduce(
             'month', dayofyear' etc. Optimize chunking ``array`` for this
             method by first rechunking using ``rechunk_for_cohorts``.
 
-    backend : {"numpy", "numba"}, optional
-        Backend for numpy_groupies
+    engine : {"numpy", "numba"}, optional
+        Engine for numpy_groupies
     keep_attrs : bool, optional
         Preserve attrs?
     skipna : bool, optional
@@ -316,7 +316,7 @@ def wrapper(array, to_group, *, func, skipna, **kwargs):
             "method": method,
             "min_count": min_count,
             "skipna": skipna,
-            "backend": backend,
+            "engine": engine,
             # The following mess exists because for multiple `by`s I factorize eagerly
             # here before passing it on; this means I have to handle the
             # "binning by single by variable" case explicitly where the factorization

tests/test_core.py

@@ -56,7 +56,7 @@ def test_alignment_error():
         groupby_reduce(da, labels, func="mean")
 
 
-@pytest.mark.parametrize("backend", ["numpy", "numba"])
+@pytest.mark.parametrize("engine", ["numpy", "numba"])
 @pytest.mark.parametrize("dtype", (float, int))
 @pytest.mark.parametrize("chunk, split_out", [(False, 1), (True, 1), (True, 2), (True, 3)])
 @pytest.mark.parametrize("expected_groups", [None, [0, 1, 2], np.array([0, 1, 2])])
@@ -86,7 +86,7 @@ def test_alignment_error():
     ],
 )
 def test_groupby_reduce(
-    array, by, expected, func, expected_groups, chunk, split_out, dtype, backend
+    array, by, expected, func, expected_groups, chunk, split_out, dtype, engine
 ):
     array = array.astype(dtype)
     if chunk:
@@ -109,18 +109,18 @@ def test_groupby_reduce(
         expected_groups=expected_groups,
         fill_value=123,
         split_out=split_out,
-        backend=backend,
+        engine=engine,
     )
     assert_equal(expected, result)
 
 
-@pytest.mark.parametrize("backend", ["numpy", "numba"])
+@pytest.mark.parametrize("engine", ["numpy", "numba"])
 @pytest.mark.parametrize("size", ((12,), (12, 5)))
 @pytest.mark.parametrize(
     "func",
     ALL_FUNCS,
 )
-def test_groupby_reduce_all(size, func, backend):
+def test_groupby_reduce_all(size, func, engine):
 
     by = np.ones(size[-1])
     array = np.random.randn(*size)
@@ -141,7 +141,7 @@ def test_groupby_reduce_all(size, func, backend):
             expected = getattr(np, func)(array, axis=-1, **kwargs)
         expected = np.expand_dims(expected, -1)
 
-        actual, _ = groupby_reduce(array, by, func=func, backend=backend, finalize_kwargs=kwargs)
+        actual, _ = groupby_reduce(array, by, func=func, engine=engine, finalize_kwargs=kwargs)
         if "arg" in func:
             assert actual.dtype.kind == "i"
         assert_equal(actual, expected)
@@ -152,7 +152,7 @@ def test_groupby_reduce_all(size, func, backend):
                 by,
                 func=func,
                 method=method,
-                backend=backend,
+                engine=engine,
                 finalize_kwargs=kwargs,
             )
             if "arg" in func:
@@ -362,11 +362,11 @@ def test_dask_reduce_axis_subset():
 
 
 @pytest.mark.parametrize("func", ALL_FUNCS)
-@pytest.mark.parametrize("backend", ["numpy", "numba"])
+@pytest.mark.parametrize("engine", ["numpy", "numba"])
 @pytest.mark.parametrize(
     "axis", [None, (0, 1, 2), (0, 1), (0, 2), (1, 2), 0, 1, 2, (0,), (1,), (2,)]
 )
-def test_groupby_reduce_axis_subset_against_numpy(func, axis, backend):
+def test_groupby_reduce_axis_subset_against_numpy(func, axis, engine):
     if not isinstance(axis, int) and "arg" in func and (axis is None or len(axis) > 1):
         pytest.skip()
     if func in ["all", "any"]:
@@ -377,7 +377,7 @@ def test_groupby_reduce_axis_subset_against_numpy(func, axis, backend):
     by = np.broadcast_to(labels2d, (3, *labels2d.shape))
     array = np.ones_like(by)
     kwargs = dict(
-        func=func, axis=axis, expected_groups=[0, 2], fill_value=fill_value, backend=backend
+        func=func, axis=axis, expected_groups=[0, 2], fill_value=fill_value, engine=engine
     )
     with raise_if_dask_computes():
         actual, _ = groupby_reduce(
@@ -389,7 +389,7 @@ def test_groupby_reduce_axis_subset_against_numpy(func, axis, backend):
     assert_equal(actual, expected)
 
 
-@pytest.mark.parametrize("backend", ["numpy", "numba"])
+@pytest.mark.parametrize("engine", ["numpy", "numba"])
 @pytest.mark.parametrize("chunks", [None, (2, 2, 3)])
 @pytest.mark.parametrize(
     "axis, groups, expected_shape",
@@ -399,7 +399,7 @@ def test_groupby_reduce_axis_subset_against_numpy(func, axis, backend):
         (None, [0], (1,)),  # global reduction; 0 shaped group axis; 1 group
     ],
 )
-def test_groupby_reduce_nans(chunks, axis, groups, expected_shape, backend):
+def test_groupby_reduce_nans(chunks, axis, groups, expected_shape, engine):
     def _maybe_chunk(arr):
         if chunks:
             return da.from_array(arr, chunks=chunks)
@@ -419,7 +419,7 @@ def _maybe_chunk(arr):
         expected_groups=groups,
         axis=axis,
         fill_value=0,
-        backend=backend,
+        engine=engine,
     )
     assert_equal(result, np.zeros(expected_shape, dtype=np.int64))
 
@@ -431,8 +431,8 @@ def _maybe_chunk(arr):
     # by = np.broadcast_to(labels2d, (3, *labels2d.shape))
 
 
-@pytest.mark.parametrize("backend", ["numpy", "numba"])
-def test_groupby_all_nan_blocks(backend):
+@pytest.mark.parametrize("engine", ["numpy", "numba"])
+def test_groupby_all_nan_blocks(engine):
     labels = np.array([0, 0, 2, 2, 2, 1, 1, 2, 2, 1, 1, 0])
     nan_labels = labels.astype(float)  # copy
     nan_labels[:5] = np.nan
@@ -448,7 +448,7 @@ def test_groupby_all_nan_blocks(backend):
         da.from_array(by, chunks=(1, 3)),
         func="sum",
         expected_groups=None,
-        backend=backend,
+        engine=engine,
     )
     assert_equal(actual, expected)
 

tests/test_xarray.py

@@ -22,11 +22,11 @@
     pass
 
 
-@pytest.mark.parametrize("backend", ["numpy", "numba"])
+@pytest.mark.parametrize("engine", ["numpy", "numba"])
 @pytest.mark.parametrize("min_count", [None, 1, 3])
 @pytest.mark.parametrize("add_nan", [True, False])
 @pytest.mark.parametrize("skipna", [True, False])
-def test_xarray_reduce(skipna, add_nan, min_count, backend):
+def test_xarray_reduce(skipna, add_nan, min_count, engine):
     arr = np.ones((4, 12))
 
     if add_nan:
@@ -46,7 +46,7 @@ def test_xarray_reduce(skipna, add_nan, min_count, backend):
 
     expected = da.groupby("labels").sum(skipna=skipna, min_count=min_count)
     actual = xarray_reduce(
-        da, "labels", func="sum", skipna=skipna, min_count=min_count, backend=backend
+        da, "labels", func="sum", skipna=skipna, min_count=min_count, engine=engine
     )
     assert_equal(expected, actual)