Added tests for EAs

Co-authored-by: André Correia <[email protected]>

Author: tiago-firmino

pandas/_libs/groupby.pyx

@@ -735,6 +735,9 @@ def group_sum(
             for j in range(K):
                 val = values[i, j]
 
+                if _treat_as_na(sumx[lab, j], is_datetimelike):
+                    continue
+
                 if uses_mask:
                     isna_entry = mask[i, j]
                 else:
@@ -1107,7 +1110,7 @@ def group_mean(
                     isna_entry = _treat_as_na(val, is_datetimelike)
 
                 if not skipna and isna_entry:
-                    sumx[lab, j] = nan_val
+                    sumx[lab, j] = val
                     nobs[lab, j] = 0
                     continue
 

pandas/core/arrays/arrow/array.py

@@ -2215,7 +2215,9 @@ def _replace_with_mask(
     def _to_masked(self):
         pa_dtype = self._pa_array.type
 
-        if pa.types.is_floating(pa_dtype) or pa.types.is_integer(pa_dtype):
+        if pa.types.is_floating(pa_dtype):
+            na_value = np.nan
+        elif pa.types.is_integer(pa_dtype):
             na_value = 1
         elif pa.types.is_boolean(pa_dtype):
             na_value = True

pandas/core/groupby/ops.py

@@ -392,6 +392,13 @@ def _call_cython_op(
                         values[mask] = True
             values = values.astype(bool, copy=False).view(np.int8)
             is_numeric = True
+        elif (
+            self.how in ["median", "sem", "std", "var"]
+            and "skipna" in kwargs
+            and not kwargs["skipna"]
+        ):
+            # if skipna=False we don't want to use masks created for Nullable dtypes
+            mask = None
 
         values = values.T
         if mask is not None:
@@ -1257,4 +1264,4 @@ def _get_splitter(
         # i.e. DataFrame
         klass = FrameSplitter
 
-    return klass(data, labels, ngroups, sort_idx=sort_idx, sorted_ids=sorted_ids)
+    return klass(data, labels, ngroups, sort_idx=sort_idx, sorted_ids=sorted_ids)

pandas/tests/groupby/test_reductions.py

@@ -3,6 +3,7 @@
 from string import ascii_lowercase
 
 import numpy as np
+import pyarrow as pa
 import pytest
 
 from pandas._libs.tslibs import iNaT
@@ -1052,7 +1053,31 @@ def scipy_sem(*args, **kwargs):
     [
         ("sum", [-1.0, 1.2, -1.1, 1.5, np.nan, 1.0]),
         ("sum", ["foo", "bar", "baz", "foo", pd.NA, "foo"]),
+        (
+            "sum",
+            Series(pd.array([-1.0, 1.2, -1.1, 1.5, np.nan, 1.0], dtype="Float64")),
+        ),
+        (
+            "sum",
+            Series(
+                pd.array(
+                    [1.0, 2.0, 3.0, np.nan, 4.0, 5.0], dtype=pd.ArrowDtype(pa.float64())
+                )
+            ),
+        ),
         ("min", [-1.0, 1.2, -1.1, 1.5, np.nan, 1.0]),
+        (
+            "min",
+            Series(pd.array([-1.0, 1.2, -1.1, 1.5, np.nan, 1.0], dtype="Float64")),
+        ),
+        (
+            "min",
+            Series(
+                pd.array(
+                    [1.0, 2.0, 3.0, np.nan, 4.0, 5.0], dtype=pd.ArrowDtype(pa.float64())
+                )
+            ),
+        ),
         (
             "min",
             [
@@ -1076,6 +1101,18 @@ def scipy_sem(*args, **kwargs):
             ],
         ),
         ("max", [-1.0, 1.2, -1.1, 1.5, np.nan, 1.0]),
+        (
+            "max",
+            Series(pd.array([-1.0, 1.2, -1.1, 1.5, np.nan, 1.0], dtype="Float64")),
+        ),
+        (
+            "max",
+            Series(
+                pd.array(
+                    [1.0, 2.0, 3.0, np.nan, 4.0, 5.0], dtype=pd.ArrowDtype(pa.float64())
+                )
+            ),
+        ),
         (
             "max",
             [
@@ -1099,6 +1136,18 @@ def scipy_sem(*args, **kwargs):
             ],
         ),
         ("mean", [-1.0, 1.2, -1.1, 1.5, np.nan, 1.0]),
+        (
+            "mean",
+            Series(pd.array([-1.0, 1.2, -1.1, 1.5, np.nan, 1.0], dtype="Float64")),
+        ),
+        (
+            "mean",
+            Series(
+                pd.array(
+                    [1.0, 2.0, 3.0, np.nan, 4.0, 5.0], dtype=pd.ArrowDtype(pa.float64())
+                )
+            ),
+        ),
         (
             "mean",
             [
@@ -1122,6 +1171,18 @@ def scipy_sem(*args, **kwargs):
             ],
         ),
         ("median", [-1.0, 1.2, -1.1, 1.5, np.nan, 1.0]),
+        (
+            "median",
+            Series(pd.array([-1.0, 1.2, -1.1, 1.5, np.nan, 1.0], dtype="Float64")),
+        ),
+        (
+            "median",
+            Series(
+                pd.array(
+                    [1.0, 2.0, 3.0, np.nan, 4.0, 5.0], dtype=pd.ArrowDtype(pa.float64())
+                )
+            ),
+        ),
         (
             "median",
             [
@@ -1145,9 +1206,57 @@ def scipy_sem(*args, **kwargs):
             ],
         ),
         ("prod", [-1.0, 1.2, -1.1, 1.5, np.nan, 1.0]),
+        (
+            "prod",
+            Series(pd.array([-1.0, 1.2, -1.1, 1.5, np.nan, 1.0], dtype="Float64")),
+        ),
+        (
+            "prod",
+            Series(
+                pd.array(
+                    [1.0, 2.0, 3.0, np.nan, 4.0, 5.0], dtype=pd.ArrowDtype(pa.float64())
+                )
+            ),
+        ),
         ("sem", [-1.0, 1.2, -1.1, 1.5, np.nan, 1.0]),
+        (
+            "sem",
+            Series(pd.array([-1.0, 1.2, -1.1, 1.5, np.nan, 1.0], dtype="Float64")),
+        ),
+        (
+            "sem",
+            Series(
+                pd.array(
+                    [1.0, 2.0, 3.0, np.nan, 4.0, 5.0], dtype=pd.ArrowDtype(pa.float64())
+                )
+            ),
+        ),
         ("std", [-1.0, 1.2, -1.1, 1.5, np.nan, 1.0]),
+        (
+            "std",
+            Series(pd.array([-1.0, 1.2, -1.1, 1.5, np.nan, 1.0], dtype="Float64")),
+        ),
+        (
+            "std",
+            Series(
+                pd.array(
+                    [1.0, 2.0, 3.0, np.nan, 4.0, 5.0], dtype=pd.ArrowDtype(pa.float64())
+                )
+            ),
+        ),
         ("var", [-1.0, 1.2, -1.1, 1.5, np.nan, 1.0]),
+        (
+            "var",
+            Series(pd.array([-1.0, 1.2, -1.1, 1.5, np.nan, 1.0], dtype="Float64")),
+        ),
+        (
+            "var",
+            Series(
+                pd.array(
+                    [1.0, 2.0, 3.0, np.nan, 4.0, 5.0], dtype=pd.ArrowDtype(pa.float64())
+                )
+            ),
+        ),
         ("any", [-1.0, 1.2, -1.1, 1.5, np.nan, 1.0]),
         ("all", [-1.0, 1.2, -1.1, 1.5, np.nan, 1.0]),
         ("skew", [-1.0, 1.2, -1.1, 1.5, np.nan, 1.0]),
@@ -1163,7 +1272,7 @@ def test_skipna_reduction_ops_cython(reduction_method, values):
     expected = gb.apply(
         lambda x: getattr(x, reduction_method)(skipna=False), include_groups=False
     )
-    tm.assert_frame_equal(result_cython, expected, check_exact=False)
+    tm.assert_frame_equal(result_cython, expected, check_exact=False, check_dtype=False)
 
 
 @pytest.mark.parametrize(
@@ -1310,4 +1419,4 @@ def test_groupby_std_datetimelike():
     td4 = pd.Timedelta("2886 days 00:42:34.664668096")
     exp_ser = Series([td1 * 2, td1, td1, td1, td4], index=np.arange(5))
     expected = DataFrame({"A": exp_ser, "B": exp_ser, "C": exp_ser})
-    tm.assert_frame_equal(result, expected)
+    tm.assert_frame_equal(result, expected)