BUG: honor column selection in GroupBy.transform on DataFrame, close #…

…1365

RELEASE.rst

@@ -166,6 +166,7 @@ pandas 0.8.0
     error (#1090)
   - Consistently set name on groupby pieces (#184)
   - Treat dict return values as Series in GroupBy.apply (#823)
+  - Respect column selection for DataFrame in in GroupBy.transform (#1365)
 
 pandas 0.7.3
 ============

pandas/core/groupby.py

@@ -1624,7 +1624,9 @@ def transform(self, func, *args, **kwargs):
         applied = []
 
         obj = self._obj_with_exclusions
-        for name, group in self:
+        gen = self.grouper.get_iterator(obj, axis=self.axis)
+
+        for name, group in gen:
             object.__setattr__(group, 'name', name)
 
             try:

pandas/tests/test_groupby.py

@@ -410,6 +410,16 @@ def test_dispatch_transform(self):
         expected = df.groupby(lambda x: x.month).transform(fillit)
         assert_frame_equal(filled, expected)
 
+    def test_transform_select_columns(self):
+        f = lambda x: x.mean()
+        result = self.df.groupby('A')['C', 'D'].transform(f)
+
+        selection = self.df[['C', 'D']]
+        expected = selection.groupby(self.df['A']).transform(f)
+
+        assert_frame_equal(result, expected)
+
+
     def test_with_na(self):
         index = Index(np.arange(10))
         values = Series(np.ones(10), index)