Fix ffill, bfill bugs

Author: dcherian

flox/aggregations.py

@@ -654,8 +654,8 @@ def __post_init__(self):
 
 
 def reverse(a: AlignedArrays) -> AlignedArrays:
-    a.group_idx = a.group_idx[::-1]
-    a.array = a.array[::-1]
+    a.group_idx = a.group_idx[..., ::-1]
+    a.array = a.array[..., ::-1]
     return a
 
 

flox/core.py

@@ -2788,20 +2788,29 @@ def groupby_scan(
     if by_.shape[-1] == 1 or by_.shape == grp_shape:
         return array.astype(agg.dtype)
 
+    # Made a design choice here to have `preprocess` handle both array and group_idx
+    # Example: for reversing, we need to reverse the whole array, not just reverse
+    #          each block independently
+    inp = AlignedArrays(array=array, group_idx=by_)
+    if agg.preprocess:
+        inp = agg.preprocess(inp)
+
     if not has_dask:
-        final_state = chunk_scan(
-            AlignedArrays(array=array, group_idx=by_), axis=single_axis, agg=agg, dtype=agg.dtype
-        )
-        return extract_array(final_state)
+        final_state = chunk_scan(inp, axis=single_axis, agg=agg, dtype=agg.dtype)
+        result = _finalize_scan(final_state)
     else:
-        return dask_groupby_scan(array, by_, axes=axis_, agg=agg)
+        result = dask_groupby_scan(inp.array, inp.group_idx, axes=axis_, agg=agg)
+
+    # Made a design choice here to have `postprocess` handle both array and group_idx
+    out = AlignedArrays(array=result, group_idx=by_)
+    if agg.finalize:
+        out = agg.finalize(out)
+    return out.array
 
 
 def chunk_scan(inp: AlignedArrays, *, axis: int, agg: Scan, dtype=None, keepdims=None) -> ScanState:
     assert axis == inp.array.ndim - 1
 
-    if agg.preprocess:
-        inp = agg.preprocess(inp)
     # I don't think we need to re-factorize here unless we are grouping by a dask array
     accumulated = generic_aggregate(
         inp.group_idx,
@@ -2813,8 +2822,6 @@ def chunk_scan(inp: AlignedArrays, *, axis: int, agg: Scan, dtype=None, keepdims
         fill_value=agg.identity,
     )
     result = AlignedArrays(array=accumulated, group_idx=inp.group_idx)
-    if agg.finalize:
-        result = agg.finalize(result)
     return ScanState(result=result, state=None)
 
 
@@ -2840,10 +2847,9 @@ def _zip(group_idx: np.ndarray, array: np.ndarray) -> AlignedArrays:
     return AlignedArrays(group_idx=group_idx, array=array)
 
 
-def extract_array(block: ScanState, finalize: Callable | None = None) -> np.ndarray:
+def _finalize_scan(block: ScanState) -> np.ndarray:
     assert block.result is not None
-    result = finalize(block.result) if finalize is not None else block.result
-    return result.array
+    return block.result.array
 
 
 def dask_groupby_scan(array, by, axes: T_Axes, agg: Scan) -> DaskArray:
@@ -2859,9 +2865,8 @@ def dask_groupby_scan(array, by, axes: T_Axes, agg: Scan) -> DaskArray:
     array, by = _unify_chunks(array, by)
 
     # 1. zip together group indices & array
-    to_map = _zip if agg.preprocess is None else tlz.compose(agg.preprocess, _zip)
     zipped = map_blocks(
-        to_map, by, array, dtype=array.dtype, meta=array._meta, name="groupby-scan-preprocess"
+        _zip, by, array, dtype=array.dtype, meta=array._meta, name="groupby-scan-preprocess"
     )
 
     scan_ = partial(chunk_scan, agg=agg)
@@ -2882,7 +2887,7 @@ def dask_groupby_scan(array, by, axes: T_Axes, agg: Scan) -> DaskArray:
     )
 
     # 3. Unzip and extract the final result array, discard groups
-    result = map_blocks(extract_array, accumulated, dtype=agg.dtype, finalize=agg.finalize)
+    result = map_blocks(_finalize_scan, accumulated, dtype=agg.dtype)
 
     assert result.chunks == array.chunks
 

tests/test_core.py

@@ -187,11 +187,11 @@ def test_groupby_reduce(
     assert_equal(expected_result, result)
 
 
-def gen_array_by(size, func):
+def gen_array_by(size, func: str):
     by = np.ones(size[-1])
     rng = np.random.default_rng(12345)
     array = rng.random(tuple(6 if s == 1 else s for s in size))
-    if "nan" in func and "nanarg" not in func:
+    if ("nan" in func or "fill" in func) and "nanarg" not in func:
         array[[1, 4, 5], ...] = np.nan
     elif "nanarg" in func and len(size) > 1:
         array[[1, 4, 5], 1] = np.nan
@@ -1810,7 +1810,7 @@ def test_nanlen_string(dtype, engine):
     assert_equal(expected, actual)
 
 
-def test_scans():
+def test_cumusm():
     array = np.array([1, 1, 1], dtype=np.uint64)
     by = np.array([0] * array.shape[-1])
     kwargs = {"func": "nancumsum", "axis": -1}
@@ -1823,3 +1823,27 @@ def test_scans():
         da = dask.array.from_array(array, chunks=2)
         actual = groupby_scan(da, by, **kwargs)
         assert_equal(expected, actual)
+
+
+@pytest.mark.parametrize(
+    "chunks",
+    [
+        pytest.param(-1, marks=requires_dask),
+        pytest.param(3, marks=requires_dask),
+        pytest.param(4, marks=requires_dask),
+    ],
+)
+@pytest.mark.parametrize("size", ((1, 12), (12,), (12, 9)))
+@pytest.mark.parametrize("add_nan_by", [True, False])
+@pytest.mark.parametrize("func", ["ffill", "bfill"])
+def test_ffill_bfill(chunks, size, add_nan_by, func):
+    array, by = gen_array_by(size, func)
+    if chunks:
+        array = dask.array.from_array(array, chunks=chunks)
+    if add_nan_by:
+        by[0:3] = np.nan
+    by = tuple(by)
+
+    expected = flox.groupby_scan(array.compute(), by, func=func)
+    actual = flox.groupby_scan(array, by, func=func)
+    assert_equal(expected, actual)

tests/test_properties.py

@@ -159,14 +159,19 @@ def test_ffill_bfill_reverse(data, array: dask.array.Array) -> None:
     def reverse(arr):
         return arr[..., ::-1]
 
-    for a in (array, array.compute()):
-        forward = groupby_scan(a, by, func="ffill")
-        backward_reversed = reverse(groupby_scan(reverse(a), reverse(by), func="bfill"))
-        assert_equal(forward, backward_reversed)
-
-        backward = groupby_scan(a, by, func="bfill")
-        forward_reversed = reverse(groupby_scan(reverse(a), reverse(by), func="ffill"))
-        assert_equal(forward_reversed, backward)
+    forward = groupby_scan(array, by, func="ffill")
+    as_numpy = groupby_scan(array.compute(), by, func="ffill")
+    assert_equal(forward, as_numpy)
+
+    backward = groupby_scan(array, by, func="bfill")
+    as_numpy = groupby_scan(array.compute(), by, func="bfill")
+    assert_equal(backward, as_numpy)
+
+    backward_reversed = reverse(groupby_scan(reverse(array), reverse(by), func="bfill"))
+    assert_equal(forward, backward_reversed)
+
+    forward_reversed = reverse(groupby_scan(reverse(array), reverse(by), func="ffill"))
+    assert_equal(forward_reversed, backward)
 
 
 @given(