Adjust groupby tests for string option

pandas/tests/groupby/aggregate/test_aggregate.py

@@ -337,7 +337,7 @@ def test_wrap_agg_out(three_group):
     grouped = three_group.groupby(["A", "B"])
 
     def func(ser):
-        if ser.dtype == object:
+        if ser.dtype in [object, pd.StringDtype("pyarrow_numpy")]:
             raise TypeError("Test error message")
         return ser.sum()
 
@@ -1098,18 +1098,19 @@ def test_lambda_named_agg(func):
     tm.assert_frame_equal(result, expected)
 
 
-def test_aggregate_mixed_types():
+def test_aggregate_mixed_types(using_infer_string):
     # GH 16916
     df = DataFrame(
         data=np.array([0] * 9).reshape(3, 3), columns=list("XYZ"), index=list("abc")
     )
     df["grouping"] = ["group 1", "group 1", 2]
     result = df.groupby("grouping").aggregate(lambda x: x.tolist())
     expected_data = [[[0], [0], [0]], [[0, 0], [0, 0], [0, 0]]]
+    dtype = "string[pyarrow_numpy]" if using_infer_string else object
     expected = DataFrame(
         expected_data,
         index=Index([2, "group 1"], dtype="object", name="grouping"),
-        columns=Index(["X", "Y", "Z"], dtype="object"),
+        columns=Index(["X", "Y", "Z"], dtype=dtype),
     )
     tm.assert_frame_equal(result, expected)
 

pandas/tests/groupby/aggregate/test_cython.py

@@ -93,7 +93,7 @@ def test_cython_agg_boolean():
     tm.assert_series_equal(result, expected)
 
 
-def test_cython_agg_nothing_to_agg():
+def test_cython_agg_nothing_to_agg(using_infer_string):
     frame = DataFrame(
         {"a": np.random.default_rng(2).integers(0, 5, 50), "b": ["foo", "bar"] * 25}
     )
@@ -107,8 +107,12 @@ def test_cython_agg_nothing_to_agg():
     )
 
     result = frame[["b"]].groupby(frame["a"]).mean(numeric_only=True)
+    dtype = "string[pyarrow_numpy]" if using_infer_string else object
+
     expected = DataFrame(
-        [], index=frame["a"].sort_values().drop_duplicates(), columns=[]
+        [],
+        index=frame["a"].sort_values().drop_duplicates(),
+        columns=Index([], dtype=dtype),
     )
     tm.assert_frame_equal(result, expected)
 

pandas/tests/groupby/aggregate/test_other.py

@@ -355,7 +355,8 @@ def test_series_agg_multi_pure_python():
     )
 
     def bad(x):
-        assert len(x.values.base) > 0
+        if x.dtype == object:
+            assert len(x.values.base) > 0
         return "foo"
 
     result = data.groupby(["A", "B"]).agg(bad)

pandas/tests/groupby/methods/test_nth.py

@@ -1,3 +1,5 @@
+from decimal import Decimal
+
 import numpy as np
 import pytest
 
@@ -707,7 +709,15 @@ def test_first_multi_key_groupby_categorical():
 @pytest.mark.parametrize("method", ["first", "last", "nth"])
 def test_groupby_last_first_nth_with_none(method, nulls_fixture):
     # GH29645
-    expected = Series(["y"])
+    if nulls_fixture is not pd.NA and (
+        nulls_fixture is pd.NaT
+        or isinstance(nulls_fixture, Decimal)
+        and Decimal.is_nan(nulls_fixture)
+    ):
+        dtype = object
+    else:
+        dtype = None
+    expected = Series(["y"], dtype=dtype)
     data = Series(
         [nulls_fixture, nulls_fixture, nulls_fixture, "y", nulls_fixture],
         index=[0, 0, 0, 0, 0],

pandas/tests/groupby/methods/test_quantile.py

@@ -171,7 +171,9 @@ def test_groupby_quantile_with_arraylike_q_and_int_columns(frame_size, groupby,
 def test_quantile_raises():
     df = DataFrame([["foo", "a"], ["foo", "b"], ["foo", "c"]], columns=["key", "val"])
 
-    with pytest.raises(TypeError, match="cannot be performed against 'object' dtypes"):
+    with pytest.raises(
+        TypeError, match="cannot be performed against 'object' dtypes|No matching"
+    ):
         df.groupby("key").quantile()
 
 
@@ -260,7 +262,8 @@ def test_groupby_quantile_raises_on_invalid_dtype(q, numeric_only):
         tm.assert_frame_equal(result, expected)
     else:
         with pytest.raises(
-            TypeError, match="'quantile' cannot be performed against 'object' dtypes!"
+            TypeError,
+            match="'quantile' cannot be performed against 'object' dtypes!|No matching",
         ):
             df.groupby("a").quantile(q, numeric_only=numeric_only)
 

pandas/tests/groupby/test_apply.py

@@ -37,7 +37,7 @@ def store(group):
     tm.assert_frame_equal(groups[0], expected_value)
 
 
-def test_apply_index_date():
+def test_apply_index_date(using_infer_string):
     # GH 5788
     ts = [
         "2011-05-16 00:00",
@@ -77,7 +77,7 @@ def test_apply_index_date():
     tm.assert_frame_equal(result, expected)
 
 
-def test_apply_index_date_object():
+def test_apply_index_date_object(using_infer_string):
     # GH 5789
     # don't auto coerce dates
     ts = [
@@ -109,8 +109,9 @@ def test_apply_index_date_object():
         1.40750,
         1.40649,
     ]
+    dtype = "string[pyarrow_numpy]" if using_infer_string else None
     exp_idx = Index(
-        ["2011-05-16", "2011-05-17", "2011-05-18"], dtype=object, name="date"
+        ["2011-05-16", "2011-05-17", "2011-05-18"], dtype=dtype, name="date"
     )
     expected = Series(["00:00", "02:00", "02:00"], index=exp_idx)
     msg = "DataFrameGroupBy.apply operated on the grouping columns"
@@ -121,14 +122,15 @@ def test_apply_index_date_object():
     tm.assert_series_equal(result, expected)
 
 
-def test_apply_trivial():
+def test_apply_trivial(using_infer_string):
     # GH 20066
     # trivial apply: ignore input and return a constant dataframe.
     df = DataFrame(
         {"key": ["a", "a", "b", "b", "a"], "data": [1.0, 2.0, 3.0, 4.0, 5.0]},
         columns=["key", "data"],
     )
-    expected = pd.concat([df.iloc[1:], df.iloc[1:]], axis=1, keys=["float64", "object"])
+    dtype = "string" if using_infer_string else "object"
+    expected = pd.concat([df.iloc[1:], df.iloc[1:]], axis=1, keys=["float64", dtype])
 
     msg = "DataFrame.groupby with axis=1 is deprecated"
     with tm.assert_produces_warning(FutureWarning, match=msg):
@@ -138,13 +140,14 @@ def test_apply_trivial():
     tm.assert_frame_equal(result, expected)
 
 
-def test_apply_trivial_fail():
+def test_apply_trivial_fail(using_infer_string):
     # GH 20066
     df = DataFrame(
         {"key": ["a", "a", "b", "b", "a"], "data": [1.0, 2.0, 3.0, 4.0, 5.0]},
         columns=["key", "data"],
     )
-    expected = pd.concat([df, df], axis=1, keys=["float64", "object"])
+    dtype = "string" if using_infer_string else "object"
+    expected = pd.concat([df, df], axis=1, keys=["float64", dtype])
     msg = "DataFrame.groupby with axis=1 is deprecated"
     with tm.assert_produces_warning(FutureWarning, match=msg):
         gb = df.groupby([str(x) for x in df.dtypes], axis=1, group_keys=True)
@@ -941,7 +944,7 @@ def test_func_returns_object():
     "group_column_dtlike",
     [datetime.today(), datetime.today().date(), datetime.today().time()],
 )
-def test_apply_datetime_issue(group_column_dtlike):
+def test_apply_datetime_issue(group_column_dtlike, using_infer_string):
     # GH-28247
     # groupby-apply throws an error if one of the columns in the DataFrame
     #   is a datetime object and the column labels are different from
@@ -952,9 +955,8 @@ def test_apply_datetime_issue(group_column_dtlike):
     with tm.assert_produces_warning(FutureWarning, match=msg):
         result = df.groupby("a").apply(lambda x: Series(["spam"], index=[42]))
 
-    expected = DataFrame(
-        ["spam"], Index(["foo"], dtype="object", name="a"), columns=[42]
-    )
+    dtype = "string" if using_infer_string else "object"
+    expected = DataFrame(["spam"], Index(["foo"], dtype=dtype, name="a"), columns=[42])
     tm.assert_frame_equal(result, expected)
 
 
@@ -1021,7 +1023,7 @@ def test_apply_multi_level_name(category):
     assert df.index.names == ["A", "B"]
 
 
-def test_groupby_apply_datetime_result_dtypes():
+def test_groupby_apply_datetime_result_dtypes(using_infer_string):
     # GH 14849
     data = DataFrame.from_records(
         [
@@ -1035,8 +1037,9 @@ def test_groupby_apply_datetime_result_dtypes():
     msg = "DataFrameGroupBy.apply operated on the grouping columns"
     with tm.assert_produces_warning(FutureWarning, match=msg):
         result = data.groupby("color").apply(lambda g: g.iloc[0]).dtypes
+    dtype = "string" if using_infer_string else object
     expected = Series(
-        [np.dtype("datetime64[ns]"), object, object, np.int64, object],
+        [np.dtype("datetime64[ns]"), dtype, dtype, np.int64, dtype],
         index=["observation", "color", "mood", "intensity", "score"],
     )
     tm.assert_series_equal(result, expected)
@@ -1302,9 +1305,7 @@ def test_apply_dropna_with_indexed_same(dropna):
     [
         [
             False,
-            DataFrame(
-                [[1, 1, 1], [2, 2, 1]], columns=Index(["a", "b", None], dtype=object)
-            ),
+            DataFrame([[1, 1, 1], [2, 2, 1]], columns=Index(["a", "b", None])),
         ],
         [
             True,

pandas/tests/groupby/test_categorical.py

@@ -82,7 +82,7 @@ def get_stats(group):
     assert result.index.names[0] == "C"
 
 
-def test_basic():  # TODO: split this test
+def test_basic(using_infer_string):  # TODO: split this test
     cats = Categorical(
         ["a", "a", "a", "b", "b", "b", "c", "c", "c"],
         categories=["a", "b", "c", "d"],
@@ -129,7 +129,8 @@ def f(x):
         result = g.apply(f)
     expected = x.iloc[[0, 1]].copy()
     expected.index = Index([1, 2], name="person_id")
-    expected["person_name"] = expected["person_name"].astype("object")
+    dtype = "string[pyarrow_numpy]" if using_infer_string else object
+    expected["person_name"] = expected["person_name"].astype(dtype)
     tm.assert_frame_equal(result, expected)
 
     # GH 9921
@@ -337,14 +338,18 @@ def test_apply(ordered):
     tm.assert_series_equal(result, expected)
 
 
-def test_observed(observed):
+def test_observed(observed, using_infer_string, request):
     # multiple groupers, don't re-expand the output space
     # of the grouper
     # gh-14942 (implement)
     # gh-10132 (back-compat)
     # gh-8138 (back-compat)
     # gh-8869
 
+    if not observed and using_infer_string:
+        mark = pytest.mark.xfail(reason="fill_value=0 invalid for string dtype")
+        request.applymarker(mark)
+
     cat1 = Categorical(["a", "a", "b", "b"], categories=["a", "b", "z"], ordered=True)
     cat2 = Categorical(["c", "d", "c", "d"], categories=["c", "d", "y"], ordered=True)
     df = DataFrame({"A": cat1, "B": cat2, "values": [1, 2, 3, 4]})