Repository: pandas-dev/pandas
Pull request: BUG: GroupBy return EA dtype
- [x] closes #23227
- [x] tests added / passed
- [x] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`
- [x] whatsnew entry


Problem statement:
BUG/EA: groupby on an EA should return the EA type
```
In [1]: df = pd.DataFrame({'Int': pd.Series([1, 2, 3], dtype='Int64'), 'A': [1, 2, 1]})
   ...: df
   ...: 
Out[1]: 
  Int  A
0   1  1
1   2  2
2   3  1

In [2]: df.groupby('A').Int.sum()
Out[2]: 
A
1    4
2    2
Name: Int, dtype: int64
```

[2] should be of type Int64; IOW this aggregation needs to be passed thru to the EA _from_sequence

Hints from issue discussion:
@jreback Hi working on this issue; did some code walkthrough, learned that here https://github.com/pandas-dev/pandas/blob/master/pandas/core/groupby/groupby.py#L834 in the method `_python_agg_general` at `result, counts = self.grouper.agg_series(obj, f)` the `obj` type is `Int64` but the `result` type is `int64`. So, I checked the `agg_series` method, seems like a Cython converter converts it from object to int64 after calculating the aggregation result. So, I am kind of lost as to where I should put `_from_squence` after checking if it's of type `ExtensionArray`
so look at _wrap_aggregated_output

there are a number of paths here though 
@jreback 
1) Tests under `pytest pandas/tests/arrays/` like `test_integer.py/test_preserve_dtypes`, `test_integer.py/test_reduce_to_float` are failing because:
```
AssertionError: Attributes are different
Attribute "dtype" are different
     [left]:  Int64      <-- result
     [right]: float64  <-- expected
```
there is custom datatype `Int64` in the test input data of the above tests and test expects it to be converted to `int64` or `float64`. Does pandas consider `Int64` as `int64`, cause `Int64` satisfies both the following conditions: `if is_extension_array_dtype(dtype):` and `if numeric_only and is_numeric_dtype(dtype) or not numeric_only` ?

2)  Tests under `pandas/tests/sparse/` like `test_groupby.py/test_groupby_includes_fill_value` are failing because:
```
AssertionError: Attributes are different           
Attribute "dtype" are different
    [left]:  Sparse[float64, nan]   <-- result
    [right]: float64                       <-- expected
```
@5hirish probably best to keep the conversation in the PR now that it's open.

Commit message: Fixed test case to re-sample categorical data for timedelta

Patch to review:
diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py
index 5acccbf688e30..d24e5a76dfc70 100644
--- a/pandas/core/groupby/groupby.py
+++ b/pandas/core/groupby/groupby.py
@@ -27,7 +27,9 @@ class providing the base-class of operations.
 )
 from pandas.core.config import option_context
 from pandas.core.dtypes.cast import maybe_downcast_to_dtype
-from pandas.core.dtypes.common import ensure_float, is_numeric_dtype, is_scalar
+from pandas.core.dtypes.common import (
+    ensure_float, is_numeric_dtype, is_scalar, is_extension_array_dtype
+)
 from pandas.core.dtypes.missing import isna, notna
 from pandas.core.frame import DataFrame
 from pandas.core.generic import NDFrame
@@ -753,7 +755,18 @@ def _try_cast(self, result, obj, numeric_only=False):
             dtype = obj.dtype
 
         if not is_scalar(result):
-            if numeric_only and is_numeric_dtype(dtype) or not numeric_only:
+            if is_extension_array_dtype(dtype):
+                # The function can return something of any type, so check
+                # if the type is compatible with the calling EA.
+                try:
+                    result = obj.values._from_sequence(result)
+                except Exception:
+                    # https://github.com/pandas-dev/pandas/issues/22850
+                    # pandas has no control over what 3rd-party ExtensionArrays
+                    # do in _values_from_sequence. We still want ops to work
+                    # though, so we catch any regular Exception.
+                    pass
+            elif numeric_only and is_numeric_dtype(dtype) or not numeric_only:
                 result = maybe_downcast_to_dtype(result, dtype)
 
         return result
diff --git a/pandas/tests/arrays/test_integer.py b/pandas/tests/arrays/test_integer.py
index 10961173d4b6b..42d0fc2844c47 100644
--- a/pandas/tests/arrays/test_integer.py
+++ b/pandas/tests/arrays/test_integer.py
@@ -652,11 +652,9 @@ def test_preserve_dtypes(op):
 
     # groupby
     result = getattr(df.groupby("A"), op)()
-    expected = pd.DataFrame({
-        "B": np.array([1.0, 3.0]),
-        "C": np.array([1, 3], dtype="int64")
-    }, index=pd.Index(['a', 'b'], name='A'))
-    tm.assert_frame_equal(result, expected)
+
+    assert result.dtypes['B'].name == 'float64'
+    assert result.dtypes['C'].name == 'Int64'
 
 
 @pytest.mark.parametrize('op', ['mean'])
@@ -675,11 +673,23 @@ def test_reduce_to_float(op):
 
     # groupby
     result = getattr(df.groupby("A"), op)()
-    expected = pd.DataFrame({
-        "B": np.array([1.0, 3.0]),
-        "C": np.array([1, 3], dtype="float64")
-    }, index=pd.Index(['a', 'b'], name='A'))
-    tm.assert_frame_equal(result, expected)
+
+    assert result.dtypes['B'].name == 'float64'
+    assert result.dtypes['C'].name == 'Int64'
+
+
+@pytest.mark.parametrize('op', ['sum'])
+def test_groupby_extension_array(op):
+    # GH23227
+    # groupby on an extension array should return the extension array type
+    df = pd.DataFrame({
+        'Int': pd.Series([1, 2, 3], dtype='Int64'),
+        'A': [1, 2, 1]
+    })
+
+    result = getattr(df.groupby('A').Int, op)()
+    assert result is not None
+    assert result.dtype.name == 'Int64'
 
 
 def test_astype_nansafe():
diff --git a/pandas/tests/sparse/test_groupby.py b/pandas/tests/sparse/test_groupby.py
index 1d2129312fb1b..ecc001702ddf7 100644
--- a/pandas/tests/sparse/test_groupby.py
+++ b/pandas/tests/sparse/test_groupby.py
@@ -26,9 +26,9 @@ def test_first_last_nth(self):
 
         # TODO: shouldn't these all be spares or not?
         tm.assert_frame_equal(sparse_grouped.first(),
-                              dense_grouped.first())
+                              dense_grouped.first().to_sparse())
         tm.assert_frame_equal(sparse_grouped.last(),
-                              dense_grouped.last())
+                              dense_grouped.last().to_sparse())
         tm.assert_frame_equal(sparse_grouped.nth(1),
                               dense_grouped.nth(1).to_sparse())
 
@@ -36,15 +36,15 @@ def test_aggfuncs(self):
         sparse_grouped = self.sparse.groupby('A')
         dense_grouped = self.dense.groupby('A')
 
-        tm.assert_frame_equal(sparse_grouped.mean(),
-                              dense_grouped.mean())
+        tm.assert_frame_equal(sparse_grouped.mean().to_sparse(),
+                              dense_grouped.mean().to_sparse())
 
         # ToDo: sparse sum includes str column
         # tm.assert_frame_equal(sparse_grouped.sum(),
         #                       dense_grouped.sum())
 
-        tm.assert_frame_equal(sparse_grouped.count(),
-                              dense_grouped.count())
+        tm.assert_frame_equal(sparse_grouped.count().to_sparse(),
+                              dense_grouped.count().to_sparse())
 
 
 @pytest.mark.parametrize("fill_value", [0, np.nan])
@@ -55,5 +55,5 @@ def test_groupby_includes_fill_value(fill_value):
     sdf = df.to_sparse(fill_value=fill_value)
     result = sdf.groupby('a').sum()
     expected = df.groupby('a').sum()
-    tm.assert_frame_equal(result, expected,
+    tm.assert_frame_equal(result, expected.to_sparse(fill_value=fill_value),
                           check_index_type=False)
diff --git a/pandas/tests/test_resample.py b/pandas/tests/test_resample.py
index 5cd31e08e0a9b..7e8aaa41a378b 100644
--- a/pandas/tests/test_resample.py
+++ b/pandas/tests/test_resample.py
@@ -1576,6 +1576,7 @@ def test_resample_categorical_data_with_timedeltaindex(self):
                               'Group': ['A', 'A']},
                              index=pd.to_timedelta([0, 10], unit='s'))
         expected = expected.reindex(['Group_obj', 'Group'], axis=1)
+        expected['Group'] = expected['Group_obj'].astype('category')
         tm.assert_frame_equal(result, expected)
 
     def test_resample_daily_anchored(self):


Code review comments:
diff_hunk: @@ -652,11 +652,9 @@ def test_preserve_dtypes(op):
 
     # groupby
     result = getattr(df.groupby("A"), op)()
-    expected = pd.DataFrame({
-        "B": np.array([1.0, 3.0]),
-        "C": np.array([1, 3], dtype="int64")
-    }, index=pd.Index(['a', 'b'], name='A'))
-    tm.assert_frame_equal(result, expected)
path: pandas/tests/arrays/test_integer.py
text: @user1:
fully construct the return frame as in the existing example

@author:
Do you mean this 
a) 
```
result = df.groupby("A").op()
assert result.dtypes['C'].name == 'Int64'
``` 
b) or do you mean that I should construct the expected as:
```
expected = pd.DataFrame({	
        "B": np.array([1.0, 3.0]),
        "C": np.array([1, 3], dtype="Int64")
    }, index=pd.Index(['a', 'b'], name='A'))
tm.assert_frame_equal(result, expected)
```
Because in b) the expected constructed dataframe converts `Int64` to `int64`.

@user2:
> Because in b) the expected constructed dataframe converts Int64 to int64

The "Int64" (capital "I") is a pandas thing. You need

```
"C": integer_array([1, 3], dtype="Int64")
```

@user1:
the 2nd. we want to construct the resulting frame and compare

@user2:
Why was this marked resolved? Am I missing the `expected` and the `tm.assert_frame_equal`?

@author:
No, you are correct, I missed that one...

diff_hunk: @@ -675,11 +673,23 @@ def test_reduce_to_float(op):
 
     # groupby
     result = getattr(df.groupby("A"), op)()
-    expected = pd.DataFrame({
-        "B": np.array([1.0, 3.0]),
-        "C": np.array([1, 3], dtype="float64")
-    }, index=pd.Index(['a', 'b'], name='A'))
-    tm.assert_frame_equal(result, expected)
+
+    assert result.dtypes['B'].name == 'float64'
+    assert result.dtypes['C'].name == 'Int64'
+
+
+@pytest.mark.parametrize('op', ['sum'])
path: pandas/tests/arrays/test_integer.py
text: @user1:
this is a duplicate of the above, no?

@author:
The intention of both the test cases is to preserve the `dtype`. The second one was added by me using the example in the reported bug. I can remove it if you say.

@author:
correct me if I am wrong

@user2:
They seem to be doing the same thing (grouping by another column and reducing the Integer column). I think this one can be removed.