Some fixes.

Author: dcherian

flox/aggregations.py

@@ -573,17 +573,21 @@ class Scan:
     # This dataclass is separate from aggregations since there's not much in common
     # between reductions and scans
     name: str
-    # binary ufunc name (e.g. add)
-    ufunc: np.ufunc
-    # in-memory scan function (e.g. cumsum)
+    # binary operation (e.g. add)
+    binary_op: Callable
+    # in-memory grouped scan function (e.g. cumsum)
     scan: str
-    # reduction that yields the last result of the scan (e.g. sum)
+    # Grouped reduction that yields the last result of the scan (e.g. sum)
     reduction: str
+    # Identity element
+    identity: Any
+    # dtype of result
+    dtype: Any = None
 
 
-cumsum = Scan("cumsum", ufunc=np.add, reduction="sum", scan="cumsum")
-nancumsum = Scan("nancumsum", ufunc=np.add, reduction="nansum", scan="nancumsum")
-# cumprod = Scan("cumprod", ufunc=np.multiply, preop="prod", scan="cumprod")
+cumsum = Scan("cumsum", binary_op=np.add, reduction="sum", scan="cumsum", identity=0)
+nancumsum = Scan("nancumsum", binary_op=np.add, reduction="nansum", scan="nancumsum", identity=0)
+# cumprod = Scan("cumprod", binary_op=np.multiply, preop="prod", scan="cumprod")
 
 
 aggregations = {

flox/core.py

@@ -2644,31 +2644,37 @@ def __post_init__(self):
         assert self.array.shape[-1] == self.group_idx.size
 
 
-def grouped_scan(inp: AlignedArrays, *, func, axis, dtype=None, keepdims=None) -> AlignedArrays:
+def grouped_scan(
+    inp: AlignedArrays, *, func: str, axis, fill_value=None, dtype=None, keepdims=None
+) -> AlignedArrays:
     assert axis == inp.array.ndim - 1
     accumulated = generic_aggregate(
-        inp.group_idx, inp.array, axis=axis, engine="numpy", func=func, dtype=dtype
+        inp.group_idx,
+        inp.array,
+        axis=axis,
+        engine="numpy",
+        func=func,
+        dtype=dtype,
+        fill_value=fill_value,
     )
     return AlignedArrays(array=accumulated, group_idx=inp.group_idx)
 
 
-def grouped_reduce(
-    inp: AlignedArrays, *, func, axis, fill_value=None, dtype=None, keepdims=None
-) -> AlignedArrays:
+def grouped_reduce(inp: AlignedArrays, *, agg: Scan, axis: int, keepdims=None) -> AlignedArrays:
     assert axis == inp.array.ndim - 1
     reduced = generic_aggregate(
         inp.group_idx,
         inp.array,
         axis=axis,
         engine="numpy",
-        func=func,
-        dtype=dtype,
-        fill_value=fill_value,
+        func=agg.reduction,
+        dtype=inp.array.dtype,
+        fill_value=agg.binary_op.identity,
     )
     return AlignedArrays(array=reduced, group_idx=np.arange(reduced.shape[-1]))
 
 
-def grouped_binop(left: AlignedArrays, right: AlignedArrays, op: np.ufunc) -> AlignedArrays:
+def grouped_binop(left: AlignedArrays, right: AlignedArrays, op: Callable) -> AlignedArrays:
     reindexed = reindex_(
         left.array,
         from_=pd.Index(left.group_idx),
@@ -2708,26 +2714,39 @@ def dask_groupby_scan(array, by, axes: T_Axes, agg: Scan):
         _zip, by, array, dtype=array.dtype, meta=array._meta, name="groupby-scan-preprocess"
     )
 
+    # TODO: move to aggregate_npg.py
+    if agg.name in ["cumsum", "nancumsum"]:
+        # https://numpy.org/doc/stable/reference/generated/numpy.cumsum.html
+        # it defaults to the dtype of a, unless a
+        # has an integer dtype with a precision less than that of the default platform integer.
+        if array.dtype.kind == "i":
+            agg.dtype = np.result_type(array.dtype, np.intp)
+        elif array.dtype.kind == "u":
+            agg.dtype = np.result_type(array.dtype, np.uintp)
+        else:
+            agg.dtype = array.dtype
+    else:
+        agg.dtype = array.dtype
+
+    scan_ = partial(grouped_scan, func=agg.scan, fill_value=agg.identity)
     # dask tokenizing error workaround
-    scan_ = partial(grouped_scan, func=agg.scan)
     scan_.__name__ = scan_.func.__name__
 
     # 2. Run the scan
     accumulated = scan(
         func=scan_,
-        binop=partial(grouped_binop, op=agg.ufunc),
-        ident=agg.ufunc.identity,
+        binop=partial(grouped_binop, op=agg.binary_op),
+        ident=agg.identity,
         x=zipped,
         axis=axis,
         method="blelloch",
-        preop=partial(grouped_reduce, func=agg.reduction, fill_value=agg.ufunc.identity),
-        dtype=array.dtype,
+        preop=partial(grouped_reduce, agg=agg),
+        dtype=agg.dtype,
     )
 
     # 3. Unzip and extract the final result array, discard groups
-    result = map_blocks(extract_array, accumulated, dtype=array.dtype)
+    result = map_blocks(extract_array, accumulated, dtype=agg.dtype)
 
-    assert result.dtype == array.dtype
     assert result.chunks == array.chunks
 
     return result

tests/test_properties.py

@@ -8,10 +8,12 @@
 import numpy as np
 from hypothesis import HealthCheck, assume, given, note, settings
 
+from flox.aggregations import cumsum
 from flox.core import dask_groupby_scan, groupby_reduce
 
 from . import ALL_FUNCS, SCIPY_STATS_FUNCS, assert_equal
 
+dask.config.set(scheduler="sync")
 NON_NUMPY_FUNCS = ["first", "last", "nanfirst", "nanlast", "count", "any", "all"] + list(
     SCIPY_STATS_FUNCS
 )
@@ -128,25 +130,30 @@ def chunked_arrays(
     return from_array(array, chunks=("auto",) * (array.ndim - 1) + (chunks,))
 
 
-from flox.aggregations import cumsum
-
-dask.config.set(scheduler="sync")
-
-
 def test():
-    array = np.array([0.0, 0.0, 0.0], dtype=np.float32)
+    # TODO: FIX
+    # array =np.array([[5592407., 5592407.],
+    #         [5592407., 5592407.]], dtype=np.float32)
+
+    array = np.array([1, 1, 1], dtype=np.uint64)
     da = dask.array.from_array(array, chunks=2)
     actual = dask_groupby_scan(
         da, np.array([0] * array.shape[-1]), agg=cumsum, axes=(array.ndim - 1,)
     )
     actual.compute()
+    expected = np.cumsum(array, axis=-1)
+    np.testing.assert_array_equal(expected, actual)
 
 
 @given(data=st.data(), array=chunked_arrays())
 def test_scans(data, array):
     note(np.array(array))
+    # overflow behaviour differs between bincount and sum (for example)
+    assume(not_overflowing_array(np.asarray(array)))
+
     actual = dask_groupby_scan(
-        array, np.array([0] * array.shape[-1]), agg=cumsum, axes=(array.ndim - 1,)
+        array, np.repeat(0, array.shape[-1]), agg=cumsum, axes=(array.ndim - 1,)
     )
     expected = np.cumsum(np.asarray(array), axis=-1)
-    np.testing.assert_array_equal(expected, actual)
+    tolerance = {"rtol": 1e-13, "atol": 1e-15}
+    assert_equal(actual, expected, tolerance)