From 0d3c3ec56dc8d9dfacf058f0d18f93ca2255f6e3 Mon Sep 17 00:00:00 2001 From: Wes McKinney Date: Sat, 2 Jun 2012 15:41:26 -0400 Subject: [PATCH] BUG: honor column selection in GroupBy.transform on DataFrame, close #1365 --- RELEASE.rst | 1 + pandas/core/groupby.py | 4 +++- pandas/tests/test_groupby.py | 10 ++++++++++ 3 files changed, 14 insertions(+), 1 deletion(-) diff --git a/RELEASE.rst b/RELEASE.rst index 70bcfa1aa79e3..610dc390898df 100644 --- a/RELEASE.rst +++ b/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 ============ diff --git a/pandas/core/groupby.py b/pandas/core/groupby.py index d03e74ee90cbe..c83d5f7831fca 100644 --- a/pandas/core/groupby.py +++ b/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: diff --git a/pandas/tests/test_groupby.py b/pandas/tests/test_groupby.py index 368edee6d8b98..340ccfadb61b1 100644 --- a/pandas/tests/test_groupby.py +++ b/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)