BUG,TST: Remove case where vectorization fails in pct_change groupby method. Incorporate CR suggestion.

Author: simonariddell

pandas/core/groupby/groupby.py

@@ -2070,10 +2070,24 @@ def shift(self, periods=1, freq=None, axis=0):
     def pct_change(self, periods=1, fill_method='pad', limit=None, freq=None,
                    axis=0):
         """Calcuate pct_change of each value to previous entry in group"""
-        return self.apply(lambda x: x.pct_change(periods=periods,
-                                                 fill_method=fill_method,
-                                                 limit=limit, freq=freq,
-                                                 axis=axis))
+        if freq is not None or axis != 0:
+            return self.apply(lambda x: x.pct_change(periods=periods,
+                                                     fill_method=fill_method,
+                                                     limit=limit, freq=freq,
+                                                     axis=axis))
+        if fill_method:
+            new = DataFrameGroupBy(self._obj_with_exclusions,
+                                   grouper=self.grouper)
+            new.obj = getattr(new, fill_method)(limit=limit)
+            new._reset_cache()
+        else:
+            new = self
+
+        obj = new.obj.drop(self.grouper.names, axis=1)
+        shifted = new.shift(periods=periods, freq=freq).\
+            drop(self.grouper.names, axis=1)
+        return (obj / shifted) - 1
+
 
     @Substitution(name='groupby')
     @Appender(_doc_template)
@@ -3936,9 +3950,15 @@ def _apply_to_column_groupbys(self, func):
 
     def pct_change(self, periods=1, fill_method='pad', limit=None, freq=None):
         """Calcuate pct_change of each value to previous entry in group"""
-        return self.apply(lambda x: x.pct_change(periods=periods,
-                                                 fill_method=fill_method,
-                                                 limit=limit, freq=freq))
+        if fill_method:
+            new = SeriesGroupBy(self.obj, grouper=self.grouper)
+            new.obj = getattr(new, fill_method)(limit=limit)
+            new._reset_cache()
+        else:
+            new = self
+
+        shifted = new.shift(periods=periods, freq=freq)
+        return (new.obj / shifted) - 1
 
 
 class NDFrameGroupBy(GroupBy):

pandas/tests/groupby/test_transform.py

@@ -730,7 +730,7 @@ def interweave(list_obj):
     (-1, 'bfill', None), (-1, 'bfill', 1),
 ])
 def test_pct_change(test_series, shuffle, periods, fill_method, limit):
-    vals = [3, np.nan, 1, 2, 4, 10, np.nan, 9]
+    vals = [3, np.nan, 1, 2, 4, 10, np.nan, 4]
     keys = ['a', 'b']
     key_v = np.repeat(keys, len(vals))
     df = DataFrame({'key': key_v, 'vals': vals * 2})
@@ -765,7 +765,10 @@ def get_result(grp_obj):
     else:
         result = get_result(grp)
         result = result.reset_index(drop=True)
-        result.columns = ['A']
+        df.insert(0, 'A', result)
+        result = df.sort_values(by='key')
+        result = result.loc[:, ['A']]
+        result = result.reset_index(drop=True)
         tm.assert_frame_equal(result, exp)