Use the input length for the sum dtype guesses and keep pure python integer sums exact

Author: OCCAM Coder

numpy_groupies/aggregate_numpy.py

@@ -669,7 +669,10 @@ def _aggregate_base(
                         a = a[good]
                     group_idx = group_idx[good]
 
-        dtype = check_dtype(dtype, funcname, a, flat_size)
+        # the sum overflow guess needs the maximum number of values that can
+        # be summed into a single group, i.e. the (nan-filtered) input length
+        # - flat_size counts output groups and can be arbitrarily smaller
+        dtype = check_dtype(dtype, funcname, a, len(group_idx))
         fill_value = resolve_fill_value(funcname, fill_value, dtype)
         check_fill_value(fill_value, dtype, func=funcname)
         ret = func(group_idx, a, flat_size, fill_value=fill_value, dtype=dtype, **kwargs)  # type: ignore[operator]

numpy_groupies/aggregate_pandas.py

@@ -67,7 +67,9 @@ def _wrapper(
     else:
         result = getattr(grouped, funcname)(**kwargs)
 
-    dtype = check_dtype(dtype, funcname, a, size)
+    # the sum overflow guess needs the maximum number of values that can be
+    # summed into a single group, i.e. the input length - not the output size
+    dtype = check_dtype(dtype, funcname, a, len(group_idx))
     if funcname.startswith("cum"):
         ret = result.to_numpy()
     else:

numpy_groupies/aggregate_purepy.py

@@ -20,9 +20,20 @@
 if TYPE_CHECKING:
     import numpy.typing as npt
 
+
 # min, max, sum, all, any - builtin
 
 
+def _sum(x: Sequence[Any]) -> Any:
+    # builtin sum over numpy integers accumulates in the input dtype without
+    # promotion (NEP 50), so large groups silently wrap - python ints keep
+    # the per-group accumulation exact, floats keep the builtin behaviour
+    kind = getattr(getattr(x[0], "dtype", None), "kind", None)
+    if kind in ("i", "u"):
+        return sum(int(val) for val in x)
+    return sum(x)
+
+
 def _last(x: Sequence[Any]) -> Any:
     return x[-1]
 
@@ -136,7 +147,7 @@ def reduction_with_nan(values):
 _impl_dict = {
     "min": _propagate_nan(min),
     "max": _propagate_nan(max),
-    "sum": sum,
+    "sum": _sum,
     "prod": _prod,
     "last": _last,
     "first": _first,

numpy_groupies/tests/test_generic.py

@@ -7,7 +7,7 @@
 import pytest
 
 from .. import aggregate_np
-from ..utils import default_fill_value
+from ..utils import default_fill_value, minimum_dtype
 from . import (
     _deselect_not_implemented,
     _deselect_purepy,
@@ -31,6 +31,35 @@ def test_preserve_missing(aggregate_all):
         assert "int" in res.dtype.name
 
 
+@pytest.mark.parametrize("dtype", [np.int8, np.uint8, np.int16, np.uint16, np.int32, np.uint32])
+def test_sum_dtype_guess_survives_single_group(aggregate_all, dtype):
+    # a single group receives every value - the sum overflow guess has to be
+    # based on the input length, not on the (here smaller) number of output
+    # groups, or the result silently wraps in the input dtype
+    maximum = int(np.iinfo(dtype).max)
+    group_idx = np.array([0, 0])
+    a = np.array([maximum, maximum], dtype=dtype)
+    # detected output size: one single group - the worst case for the guess
+    res = aggregate_all(group_idx, a, func="sum")
+    np.testing.assert_array_equal(res, [2 * maximum])
+    # an explicitly larger output size with fill values must not change this
+    res = aggregate_all(group_idx, a, func="sum", size=3, fill_value=-1)
+    np.testing.assert_array_equal(res, [2 * maximum, -1, -1])
+    if not isinstance(res, list):
+        assert res.dtype == minimum_dtype(2 * maximum, np.dtype(dtype))
+
+
+def test_sum_bool_dtype_guess_ignores_empty_groups(aggregate_all):
+    # empty groups inflate the output size, but never the number of values
+    # that get summed into one group - the bool sum guess stays narrow
+    res = aggregate_all(np.array([0, 0]), np.array([True, True]), func="sum", size=300)
+    values = np.asarray(res)
+    assert values[0] == 2
+    assert not values[1:].any()
+    if not isinstance(res, list):
+        assert res.dtype == np.int8
+
+
 @pytest.mark.parametrize("group_idx_type", [int, "uint32", "uint64"])
 def test_uint_group_idx(aggregate_all, group_idx_type):
     group_idx = np.array([1, 1, 2, 2, 2, 2, 4, 4], dtype=group_idx_type)