pandas-dev/pandas · #68933

BUG: DataFrame.agg with a list of funcs casts uint64 and bool results (GH#65031)

jbrockmendel · merged Oct 8, 20262 files · 102 + / 5 −
pandas/core/apply.py29 + / 5 −
@@ -23,7 +23,10 @@     set_module, ) -from pandas.core.dtypes.cast import is_nested_object+from pandas.core.dtypes.cast import (+    find_common_type,+    is_nested_object,+) from pandas.core.dtypes.common import (     is_dict_like,     is_extension_array_dtype,@@ -68,6 +71,7 @@         AggObjType,         Axis,         AxisInt,+        DtypeObj,         NDFrameT,         npt,     )@@ -118,6 +122,18 @@ )  +def _stack_as_object(dtypes: list[DtypeObj]) -> bool:+    """+    Whether results of these dtypes must be stacked as object to keep their+    values: bools mixed with other dtypes can be cast to numbers, and signed+    with unsigned integers to a float, which rounds values above 2**53.+    """+    kinds = {dtype.kind for dtype in dtypes}+    if "b" in kinds and len(kinds) > 1:+        return True+    return {"i", "u"} <= kinds and find_common_type(dtypes).kind == "f"++ @set_module("pandas.api.executors") class BaseExecutionEngine(abc.ABC):     """@@ -1157,7 +1173,7 @@ def _agg_list_like_frame_reductions(self) -> DataFrame | None:         for dtype in groups:             cols = groups[dtype]             sub = obj[cols]-            group_pieces: list[DataFrame] = []+            rows: list[Series] = []             for func_name in func_names:                 try:                     row = getattr(sub, func_name)(*self.args, **self.kwargs)@@ -1172,9 +1188,17 @@ def _agg_list_like_frame_reductions(self) -> DataFrame | None:                     # by the columns; anything else would silently misalign                     # in the concat below.                     return None-                # to_frame().T avoids the slow DataFrame(list-of-Series) path-                group_pieces.append(row.to_frame(func_name).T)-            pieces.append(concat(group_pieces))+                rows.append(row)++            # to_frame().T avoids the slow DataFrame(list-of-Series) path+            frames = [+                row.to_frame(name).T for name, row in zip(func_names, rows, strict=True)+            ]+            if _stack_as_object([row.dtype for row in rows]):+                # GH#65031 see test_agg_list_like_unsigned_and_signed_is_object+                # and test_agg_list_like_bool_with_numeric_is_object+                frames = [frame.astype(object) for frame in frames]+            pieces.append(concat(frames))          result = concat(pieces, axis=1)         result = result.reindex(columns=obj.columns)
pandas/tests/apply/test_frame_apply.py73 + / 0 −
@@ -2037,6 +2037,79 @@ def test_agg_empty_dict():     tm.assert_frame_equal(result, expected)  +@pytest.mark.parametrize(+    "dtype", ["uint64", "UInt64", "uint64[pyarrow]", "Sparse[uint64]"]+)+@pytest.mark.parametrize("funcs", [["max", "count"], ["max", "count", "mean"]])+@pytest.mark.parametrize("big", [True, False])+def test_agg_list_like_unsigned_and_signed_is_object(dtype, funcs, big):+    # GH#65031 max returns unsigned and count signed, and their common float+    # dtype rounds values above 2**53, so the result is object whatever the values+    if dtype == "uint64[pyarrow]":+        pytest.importorskip("pyarrow")+    value = 2**64 - 3 if big else 5+    df = pd.DataFrame(+        {"a": pd.Series([value, 1], dtype=dtype), "b": pd.Series([5, 1], dtype=dtype)}+    )++    result = df.agg(funcs)++    # asserted entry-wise because mean's own result is not under test+    assert (result.dtypes == object).all()+    assert result.loc["max", "a"] == value+    assert result.loc["count", "a"] == 2+++@pytest.mark.parametrize("dtype", ["int64", "float64", "bool"])+@pytest.mark.parametrize("funcs", [["sum", "any"], ["max", "all", "sum"]])+def test_agg_list_like_bool_with_numeric_is_object(dtype, funcs):+    # GH#65031 bool results are not cast to numbers, as in Series.agg+    df = pd.DataFrame({"a": pd.Series([1, 1], dtype=dtype)})++    result = df.agg(funcs)++    tm.assert_series_equal(result["a"], df["a"].agg(funcs))+    assert result["a"].dtype == object+    assert result.loc[funcs[1], "a"] is True+++@pytest.mark.parametrize("funcs", [["sum", "mean"], ["min", "max", "mean"]])+@pytest.mark.parametrize("values", [[10**17, 10**17 + 16], [10**17, 10**17 + 1]])+def test_agg_list_like_signed_int_and_float_is_float(funcs, values):+    # GH#65031 int64 with float64 stays float64 as in Series.agg, whatever the+    # values+    df = pd.DataFrame({"a": values})++    result = df.agg(funcs)++    tm.assert_series_equal(result["a"], df["a"].agg(funcs))+    assert result["a"].dtype == np.float64+++def test_agg_list_like_unsigned_and_signed_common_int():+    # GH#65031 uint32 and int64 have a common integer dtype, so nothing is lost+    df = pd.DataFrame({"a": pd.Series([5, 1], dtype="uint32")})++    result = df.agg(["max", "count"])++    expected = pd.DataFrame({"a": [5, 2]}, index=["max", "count"])+    tm.assert_frame_equal(result, expected)+++def test_agg_list_like_axis_1_not_cast_to_float():+    # GH#65031 axis=1 reduces the transpose, so the object fallback applies there+    # too+    big = 2**64 - 3+    df = pd.DataFrame({"a": pd.Series([big, 5], dtype="uint64")})++    result = df.agg(["max", "count"], axis=1)++    # the dtype assert is what fires: a rounded np.float64 still compares equal to+    # big, because numpy widens the int rather than comparing it exactly+    assert result["max"].dtype == object+    assert result.loc[0, "max"] == big++ def test_agg_dist_like_and_nonunique_columns():     # GH#51099     df = pd.DataFrame(