diff --git a/flox/core.py b/flox/core.py index d7fa5f6a..bb84ab80 100644 --- a/flox/core.py +++ b/flox/core.py @@ -2663,10 +2663,18 @@ def groupby_reduce( groups = (groups[0][sorted_idx],) if factorize_early: + assert len(groups) == 1 + (groups_,) = groups # nan group labels are factorized to -1, and preserved # now we get rid of them by reindexing - # This also handles bins with no data - result = reindex_(result, from_=groups[0], to=expected_, fill_value=fill_value).reshape( + # First, for "blockwise", we can have -1 repeated in different blocks + # This breaks the reindexing so remove those first. + if method == "blockwise" and (mask := groups_ == -1).sum(axis=-1) > 1: + result = result[..., ~mask] + groups_ = groups_[..., ~mask] + + # This reindex also handles bins with no data + result = reindex_(result, from_=groups_, to=expected_, fill_value=fill_value).reshape( result.shape[:-1] + grp_shape ) groups = final_groups diff --git a/tests/test_core.py b/tests/test_core.py index 5d4e7ec3..2d225206 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -1929,3 +1929,17 @@ def test_ffill_bfill(chunks, size, add_nan_by, func): expected = flox.groupby_scan(array.compute(), by, func=func) actual = flox.groupby_scan(array, by, func=func) assert_equal(expected, actual) + + +@requires_dask +def test_blockwise_nans(): + array = dask.array.ones((1, 10), chunks=2) + by = np.array([-1, 0, -1, 1, -1, 2, -1, 3, 4, 4]) + actual, actual_groups = flox.groupby_reduce( + array, by, func="sum", expected_groups=pd.RangeIndex(0, 5) + ) + expected, expected_groups = flox.groupby_reduce( + array.compute(), by, func="sum", expected_groups=pd.RangeIndex(0, 5) + ) + assert_equal(expected_groups, actual_groups) + assert_equal(expected, actual)