Fixed issue when filling Series after GroupBy

Author: WillAyd

pandas/core/groupby.py

@@ -2041,31 +2041,37 @@ def _get_group_keys(self):
                                           self.levels,
                                           self.labels)
 
-    def _cython_apply(self, f, data, axis, **kwargs):
-        output = collections.OrderedDict()
-        for col in data.columns:
-            if col in self.names:
-                output[col] = data[col].values
-            else:
-                # duplicative of _get_cython_function; needs refactor
-                dtype_str = data[col].dtype.name
-                values = data[col].values[:, None]
-                func = afunc = self._get_func(f['name'], dtype_str)
-                f = f.get('f')
+    def _cython_apply(self, ftype, data, axis, **kwargs):
+        def _generate_output(ser):
+            # duplicative of _get_cython_function; needs refactor
+            dtype_str = ser.dtype.name
+            values = ser.values[:, None]
+            func = afunc = self._get_func(ftype['name'], dtype_str)
+            f = ftype.get('f')
 
-                def wrapper(*args, **kwargs):
-                    return f(afunc, *args, **kwargs)
+            def wrapper(*args, **kwargs):
+                return f(afunc, *args, **kwargs)
 
-                func = wrapper
-                labels, _, _ = self.group_info
+            func = wrapper
+            labels, _, _ = self.group_info
+
+            result = _maybe_fill(np.empty_like(values, dtype=dtype_str),
+                                 fill_value=np.nan)
+            func(result, values, labels, **kwargs)
 
-                result = _maybe_fill(np.empty_like(values, dtype=dtype_str),
-                                     fill_value=np.nan)
-                func(result, values, labels, **kwargs)
-                output[col] = result[:, 0]
+            return result[:, 0]
 
-        # Ugh
-        return DataFrame(output, index=data.index)
+        # Using introspection to determine result; not ideal needs refactor
+        if type(data) is Series:
+            return Series(_generate_output(data), name=data.name)
+        else:
+            output = collections.OrderedDict()
+            for col in data.columns:
+                if col in self.names:
+                    output[col] = data[col].values
+                else:
+                    output[col] = _generate_output(data[col])
+            return DataFrame(output, index=data.index)
 
     def apply(self, f, data, axis=0):
         mutated = self.mutated

pandas/tests/groupby/test_groupby.py

@@ -2061,6 +2061,7 @@ def test_rank_object_raises(self, ties_method, ascending, na_option,
                                    ascending=ascending,
                                    na_option=na_option, pct=pct)
 
+    @pytest.mark.parametrize("as_series", [True, False])
     @pytest.mark.parametrize("fill_method,limit,exp_vals", [
         ("ffill", None,
          [np.nan, np.nan, 'foo', 'foo', 'foo', 'bar', 'bar', 'bar']),
@@ -2071,14 +2072,20 @@ def test_rank_object_raises(self, ties_method, ascending, na_option,
         ("bfill", 1,
          [np.nan, 'foo', 'foo', np.nan, 'bar', 'bar', np.nan, np.nan])
     ])
-    def test_group_fill_methods(self, fill_method, limit, exp_vals):
+    def test_group_fill_methods(self, as_series, fill_method, limit, exp_vals):
         vals = [np.nan, np.nan, 'foo', np.nan, np.nan, 'bar', np.nan, np.nan]
         keys = ['a'] * len(vals) + ['b'] * len(vals)
         df = DataFrame({'key': keys, 'val': vals * 2})
-        result = getattr(df.groupby('key'), fill_method)(limit=limit)
 
-        exp = DataFrame({'key': keys, 'val': exp_vals * 2})
-        assert_frame_equal(result, exp)
+        if as_series:
+            result = getattr(
+                df.groupby('key')['val'], fill_method)(limit=limit)
+            exp = Series(exp_vals * 2, name='val')
+            assert_series_equal(result, exp)
+        else:
+            result = getattr(df.groupby('key'), fill_method)(limit=limit)
+            exp = DataFrame({'key': keys, 'val': exp_vals * 2})
+            assert_frame_equal(result, exp)
 
     def test_dont_clobber_name_column(self):
         df = DataFrame({'key': ['a', 'a', 'a', 'b', 'b', 'b'],