Backport PR pandas-dev#55620 on branch 2.1.x (BUG: Groupby not keepin…

…g object dtype when infer string is set) (pandas-dev#55629) Backport PR pandas-dev#55620: BUG: Groupby not keeping object dtype when infer string is set Co-authored-by: Patrick Hoefler <[email protected]>
meeseeksmachine · Oct 22, 2023 · c26935c · c26935c
1 parent 3c1fe5c
commit c26935c
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 4 deletions.
diff --git a/doc/source/whatsnew/v2.1.2.rst b/doc/source/whatsnew/v2.1.2.rst
@@ -23,6 +23,7 @@ Fixed regressions
 
 Bug fixes
 ~~~~~~~~~
+- Fixed bug in :class:`.DataFrameGroupBy` reductions not preserving object dtype when ``infer_string`` is set (:issue:`55620`)
 - Fixed bug in :meth:`Categorical.equals` if other has arrow backed string dtype (:issue:`55364`)
 - Fixed bug in :meth:`DataFrame.__setitem__` not inferring string dtype for zero-dimensional array with ``infer_string=True`` (:issue:`55366`)
 - Fixed bug in :meth:`DataFrame.idxmin` and :meth:`DataFrame.idxmax` raising for arrow dtypes (:issue:`55368`)

diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py
@@ -1855,8 +1855,7 @@ def _agg_py_fallback(
             ser = Series(values, copy=False)
         else:
             # We only get here with values.dtype == object
-            # TODO: special case not needed with ArrayManager
-            df = DataFrame(values.T)
+            df = DataFrame(values.T, dtype=values.dtype)
             # bc we split object blocks in grouped_reduce, we have only 1 col
             # otherwise we'd have to worry about block-splitting GH#39329
             assert df.shape[1] == 1

diff --git a/pandas/tests/groupby/test_groupby.py b/pandas/tests/groupby/test_groupby.py
@@ -5,6 +5,7 @@
 import numpy as np
 import pytest
 
+from pandas.compat import pa_version_under7p0
 from pandas.errors import (
     PerformanceWarning,
     SpecificationError,
@@ -2513,13 +2514,24 @@ def test_groupby_column_index_name_lost(func):
     tm.assert_index_equal(result, expected)
 
 
-def test_groupby_duplicate_columns():
+@pytest.mark.parametrize(
+    "infer_string",
+    [
+        False,
+        pytest.param(
+            True,
+            marks=pytest.mark.skipif(pa_version_under7p0, reason="arrow not installed"),
+        ),
+    ],
+)
+def test_groupby_duplicate_columns(infer_string):
     # GH: 31735
     df = DataFrame(
         {"A": ["f", "e", "g", "h"], "B": ["a", "b", "c", "d"], "C": [1, 2, 3, 4]}
     ).astype(object)
     df.columns = ["A", "B", "B"]
-    result = df.groupby([0, 0, 0, 0]).min()
+    with pd.option_context("future.infer_string", infer_string):
+        result = df.groupby([0, 0, 0, 0]).min()
     expected = DataFrame(
         [["e", "a", 1]], index=np.array([0]), columns=["A", "B", "B"], dtype=object
     )