ENH: add groupby & reduce support to EA

pandas/tests/arrays/test_integer.py

@@ -114,6 +114,13 @@ def _check_op(self, s, op_name, other, exc=None):
         # compute expected
         mask = s.isna()
 
+        # if s is a DataFrame, squeeze to a Series
+        # for comparison
+        if isinstance(s, pd.DataFrame):
+            result = result.squeeze()
+            s = s.squeeze()
+            mask = mask.squeeze()
+
         # other array is an Integer
         if isinstance(other, IntegerArray):
             omask = getattr(other, 'mask', None)
@@ -215,7 +222,6 @@ def test_arith_series_with_scalar(self, data, all_arithmetic_operators):
         s = pd.Series(data)
         self._check_op(s, op, 1, exc=TypeError)
 
-    @pytest.mark.xfail(run=False, reason="_reduce needs implementation")
     def test_arith_frame_with_scalar(self, data, all_arithmetic_operators):
         # frame & scalar
         op = all_arithmetic_operators
@@ -588,7 +594,7 @@ def test_cross_type_arithmetic():
 
 
 @pytest.mark.parametrize('op', ['sum', 'min', 'max'])
-def test_preserve_groupby_dtypes(op):
+def test_preserve_dtypes(op):
     # TODO(#22346): preserve Int64 dtype
     # for ops that enable (mean would actually work here
     # but generally it is a float return value)
@@ -598,6 +604,11 @@ def test_preserve_groupby_dtypes(op):
         "C": integer_array([1, None, 3], dtype='Int64'),
     })
 
+    # op
+    result = getattr(df.C, op)()
+    assert isinstance(result, int)
+
+    # groupby
     result = getattr(df.groupby("A"), op)()
     expected = pd.DataFrame({
         "B": np.array([1.0, 3.0]),

pandas/tests/extension/test_integer.py

@@ -207,35 +207,7 @@ class TestCasting(base.BaseCastingTests):
 
 
 class TestGroupby(base.BaseGroupbyTests):
-
-    @pytest.mark.parametrize('as_index', [True, False])
-    def test_groupby_extension_agg(self, as_index, data_for_grouping):
-        df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1, 4],
-                           "B": data_for_grouping})
-        result = df.groupby("B", as_index=as_index).A.mean()
-        _, index = pd.factorize(data_for_grouping, sort=True)
-
-        # TODO(ExtensionIndex): remove coercion to object
-        # we don't have an easy way to represent an EA as an Index object
-        index = pd.Index(index, name="B", dtype=object)
-        expected = pd.Series([3, 1, 4], index=index, name="A")
-        if as_index:
-            self.assert_series_equal(result, expected)
-        else:
-            expected = expected.reset_index()
-            self.assert_frame_equal(result, expected)
-
-    def test_groupby_extension_no_sort(self, data_for_grouping):
-        df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1, 4],
-                           "B": data_for_grouping})
-        result = df.groupby("B", sort=False).A.mean()
-        _, index = pd.factorize(data_for_grouping, sort=False)
-
-        # TODO(ExtensionIndex): remove coercion to object
-        # we don't have an easy way to represent an EA as an Index object
-        index = pd.Index(index, name="B", dtype=object)
-        expected = pd.Series([1, 3, 4], index=index, name="A")
-        self.assert_series_equal(result, expected)
+    pass
 
 
 class TestNumericReduce(base.BaseNumericReduceTests):