Merge pull request #9699 from dsm054/fix-transform-length-mismatch

BUG: ensure we use group sizes, not group counts, in transform (GH9697)

Author: jreback

doc/source/whatsnew/v0.16.1.txt

@@ -46,3 +46,5 @@ Performance Improvements
 
 Bug Fixes
 ~~~~~~~~~
+
+- Bug in ``transform`` causing length mismatch when null entries were present and a fast aggregator was being used (:issue:`9697`)

pandas/core/groupby.py

@@ -2453,7 +2453,7 @@ def _transform_fast(self, func):
         if isinstance(func, compat.string_types):
             func = getattr(self,func)
         values = func().values
-        counts = self.count().values
+        counts = self.size().values
         values = np.repeat(values, com._ensure_platform_int(counts))
 
         return self._set_result_index_ordered(Series(values))

pandas/tests/test_groupby.py

@@ -1058,6 +1058,19 @@ def test_transform_function_aliases(self):
         expected = self.df.groupby('A')['C'].transform(np.mean)
         assert_series_equal(result, expected)
 
+    def test_transform_length(self):
+        # GH 9697
+        df = pd.DataFrame({'col1':[1,1,2,2], 'col2':[1,2,3,np.nan]})
+        expected = pd.Series([3.0]*4)
+        def nsum(x):
+            return np.nansum(x)
+        results = [df.groupby('col1').transform(sum)['col2'],
+                   df.groupby('col1')['col2'].transform(sum),
+                   df.groupby('col1').transform(nsum)['col2'],
+                   df.groupby('col1')['col2'].transform(nsum)]
+        for result in results:
+            assert_series_equal(result, expected)
+
     def test_with_na(self):
         index = Index(np.arange(10))