ENH: groupby.apply for Categorical should preserve categories (closes #10138)

Author: mortada

doc/source/whatsnew/v0.17.0.txt

@@ -26,6 +26,8 @@ New features
 Other enhancements
 ^^^^^^^^^^^^^^^^^^
 
+- groupby.apply aggregation for Categorical now preserves categories (:issue:`10138`)
+
 .. _whatsnew_0170.api:
 
 Backwards incompatible API changes

pandas/core/groupby.py

@@ -2940,7 +2940,8 @@ def _wrap_applied_output(self, keys, values, not_indexed_same=False):
                     cd = 'coerce'
                 else:
                     cd = True
-                return result.convert_objects(convert_dates=cd)
+                result = result.convert_objects(convert_dates=cd)
+                return self._reindex_output(result)
 
             else:
                 # only coerce dates if we find at least 1 datetime

pandas/tests/test_groupby.py

@@ -2596,6 +2596,34 @@ def get_stats(group):
         result = self.df.groupby(cats).D.apply(get_stats)
         self.assertEqual(result.index.names[0], 'C')
 
+    def test_apply_categorical_data(self):
+        # GH 10138
+        dense = Categorical(list('abc'))
+        # 'b' is in the categories but not in the list
+        missing = Categorical(list('aaa'), categories=['a', 'b'])
+        values = np.arange(len(dense))
+        df = DataFrame({'missing': missing,
+                        'dense': dense,
+                        'values': values})
+        grouped = df.groupby(['missing', 'dense'])
+
+        # missing category 'b' should still exist in the output index
+        idx = MultiIndex.from_product([['a', 'b'], ['a', 'b', 'c']],
+                                      names=['missing', 'dense'])
+        expected = DataFrame([0, 1, 2, np.nan, np.nan, np.nan],
+                             index=idx,
+                             columns=['values'])
+
+        assert_frame_equal(grouped.apply(lambda x: np.mean(x)), expected)
+        assert_frame_equal(grouped.mean(), expected)
+        assert_frame_equal(grouped.agg(np.mean), expected)
+
+        # but for transform we should still get back the original index
+        idx = MultiIndex.from_product([['a'], ['a', 'b', 'c']],
+                                      names=['missing', 'dense'])
+        expected = Series(1, index=idx)
+        assert_series_equal(grouped.apply(lambda x: 1), expected)
+
     def test_apply_corner_cases(self):
         # #535, can't use sliding iterator