wip

Author: jreback

pandas/core/groupby.py

@@ -341,10 +341,14 @@ class _GroupBy(PandasObject, SelectionMixin):
 
     def __init__(self, obj, keys=None, axis=0, level=None,
                  grouper=None, exclusions=None, selection=None, as_index=True,
-                 sort=True, group_keys=True, squeeze=False, **kwargs):
+                 sort=True, group_keys=True, squeeze=False, ref_obj=None, **kwargs):
 
         self._selection = selection
 
+        if ref_obj is None:
+            ref_obj = obj
+        self.ref_obj = ref_obj
+
         if isinstance(obj, NDFrame):
             obj._consolidate_inplace()
 
@@ -796,7 +800,9 @@ def _cython_agg_general(self, how, weights=None, numeric_only=True):
         if weights is not None:
 
             # TODO, need to integrate this with the exclusions
-            _, weights = weightby.weightby(self.obj, weights=weights, axis=axis)
+            _, weights = weightby.weightby(self.ref_obj,
+                                           weights=weights,
+                                           axis=self.axis)
 
         output = {}
         for name, obj in self._iterate_slices():
@@ -3189,13 +3195,14 @@ def _wrap_agged_blocks(self, items, blocks):
 
     _block_agg_axis = 0
 
-    def _cython_agg_blocks(self, how, weights=None, numeric_only=True):
+    def _cython_agg_blocks(self, how, weights=None, numeric_only=True,
+                           **kwargs):
         data, agg_axis = self._get_data_to_aggregate()
 
         if weights is not None:
 
             # TODO, need to integrate this with the exclusions
-            _, weights = weightby.weightby(self.obj,
+            _, weights = weightby.weightby(self.ref_obj,
                                            weights=weights,
                                            axis=self.axis)
 
@@ -3765,19 +3772,20 @@ def _gotitem(self, key, ndim, subset=None):
         subset : object, default None
             subset to act on
         """
-
         if ndim == 2:
             if subset is None:
                 subset = self.obj
             return DataFrameGroupBy(subset, self.grouper, selection=key,
                                     grouper=self.grouper,
                                     exclusions=self.exclusions,
-                                    as_index=self.as_index)
+                                    as_index=self.as_index,
+                                    ref_obj=self.obj)
         elif ndim == 1:
             if subset is None:
                 subset = self.obj[key]
             return SeriesGroupBy(subset, selection=key,
-                                 grouper=self.grouper)
+                                 grouper=self.grouper,
+                                 ref_obj=self.obj)
 
         raise AssertionError("invalid ndim for _gotitem")
 

pandas/tools/tests/test_weightby.py

@@ -61,13 +61,20 @@ def test_basic(self):
     def test_groupby(self):
 
         for f in ['mean', 'sum']:
+
             weights = (self.df3['A'] / self.df3.A.sum()).values
             result = getattr(self.df3.groupby('C'), f)(weights='A')
             adj = self.df3.assign(A=self.df3.A * weights,
                                   B=self.df3.B * weights)
             expected = getattr(adj.groupby('C'), f)()
             tm.assert_frame_equal(result, expected)
 
+            weights = (self.df3['A'] / self.df3.A.sum()).values
+            result = getattr(self.df3.groupby('C').B, f)(weights='A')
+            adj = self.df3.assign(B=self.df3.B * weights)
+            expected = getattr(adj.groupby('C').B, f)()
+            tm.assert_series_equal(result, expected)
+
     def test_unsupported(self):
         for f in ['first', 'median', 'min', 'max', 'prod']: