PERF: DataFrame transform

Author: chris-b1

asv_bench/benchmarks/groupby.py

@@ -773,6 +773,21 @@ def setup(self):
     def time_groupby_transform_series2(self):
         self.df.groupby('id')['val'].transform(np.mean)
 
+
+class groupby_transform_dataframe(object):
+    # GH 12737
+    goal_time = 0.2
+
+    def setup(self):
+        self.df = pd.DataFrame({'group': np.repeat(np.arange(1000), 10),
+                                'B': np.nan,
+                                'C': np.nan})
+        self.df.ix[4::10, 'B':'C'] = 5
+
+    def time_groupby_transform_dataframe(self):
+        self.df.groupby('group').transform('first')
+
+
 class groupby_transform_cythonized(object):
     goal_time = 0.2
 

doc/source/whatsnew/v0.18.2.txt

@@ -104,7 +104,7 @@ Performance Improvements
 - increased performance of ``DataFrame.quantile()`` as it now operates per-block (:issue:`11623`)
 
 
-
+- Improved performance of ``DataFrameGroupBy.transform`` (:issue:`12737`)
 
 
 .. _whatsnew_0182.bug_fixes:
@@ -123,7 +123,7 @@ Bug Fixes
 
 - Regression in ``Series.quantile`` with nans (also shows up in ``.median()`` and ``.describe()``); furthermore now names the ``Series`` with the quantile (:issue:`13098`, :issue:`13146`)
 
-
+- Bug in ``SeriesGroupBy.transform`` with datetime values and missing groups (:issue:`13191`)
 
 - Bug in ``Series.str.extractall()`` with ``str`` index raises ``ValueError``  (:issue:`13156`)
 

pandas/core/groupby.py

@@ -2776,18 +2776,9 @@ def _transform_fast(self, func):
             func = getattr(self, func)
 
         ids, _, ngroup = self.grouper.group_info
-        mask = ids != -1
-
-        out = func().values[ids]
-        if not mask.all():
-            out = np.where(mask, out, np.nan)
-
-        obs = np.zeros(ngroup, dtype='bool')
-        obs[ids[mask]] = True
-        if not obs.all():
-            out = self._try_cast(out, self._selected_obj)
 
-        return Series(out, index=self.obj.index)
+        out = algos.take_1d(func().values, ids)
+        return Series(out, index=self.obj.index, name=self.obj.name)
 
     def filter(self, func, dropna=True, *args, **kwargs):  # noqa
         """
@@ -3465,19 +3456,13 @@ def transform(self, func, *args, **kwargs):
         if not result.columns.equals(obj.columns):
             return self._transform_general(func, *args, **kwargs)
 
-        results = np.empty_like(obj.values, result.values.dtype)
-        for (name, group), (i, row) in zip(self, result.iterrows()):
-            indexer = self._get_index(name)
-            if len(indexer) > 0:
-                results[indexer] = np.tile(row.values, len(
-                    indexer)).reshape(len(indexer), -1)
-
-        counts = self.size().fillna(0).values
-        if any(counts == 0):
-            results = self._try_cast(results, obj[result.columns])
+        # Fast transform
+        ids, _, ngroup = self.grouper.group_info
+        out = {}
+        for col in result:
+            out[col] = algos.take_nd(result[col].values, ids)
 
-        return (DataFrame(results, columns=result.columns, index=obj.index)
-                ._convert(datetime=True))
+        return DataFrame(out, columns=result.columns, index=obj.index)
 
     def _define_paths(self, func, *args, **kwargs):
         if isinstance(func, compat.string_types):

pandas/tests/test_groupby.py

@@ -1051,13 +1051,33 @@ def test_transform_fast(self):
 
         values = np.repeat(grp.mean().values,
                            com._ensure_platform_int(grp.count().values))
-        expected = pd.Series(values, index=df.index)
+        expected = pd.Series(values, index=df.index, name='val')
         result = grp.transform(np.mean)
         assert_series_equal(result, expected)
 
         result = grp.transform('mean')
         assert_series_equal(result, expected)
 
+        # GH 12737
+        df = pd.DataFrame({'grouping': [0, 1, 1, 3], 'f': [1.1, 2.1, 3.1, 4.5],
+                           'd': pd.date_range('2014-1-1', '2014-1-4'),
+                           'i': [1, 2, 3, 4]},
+                          columns=['grouping', 'f', 'i', 'd'])
+        result = df.groupby('grouping').transform('first')
+
+        dates = [pd.Timestamp('2014-1-1'), pd.Timestamp('2014-1-2'),
+                 pd.Timestamp('2014-1-2'), pd.Timestamp('2014-1-4')]
+        expected = pd.DataFrame({'f': [1.1, 2.1, 2.1, 4.5],
+                                 'd': dates,
+                                 'i': [1, 2, 2, 4]},
+                                columns=['f', 'i', 'd'])
+        assert_frame_equal(result, expected)
+
+        # selection
+        result = df.groupby('grouping')[['f', 'i']].transform('first')
+        expected = expected[['f', 'i']]
+        assert_frame_equal(result, expected)
+
     def test_transform_broadcast(self):
         grouped = self.ts.groupby(lambda x: x.month)
         result = grouped.transform(np.mean)
@@ -1191,6 +1211,16 @@ def test_transform_function_aliases(self):
         expected = self.df.groupby('A')['C'].transform(np.mean)
         assert_series_equal(result, expected)
 
+    def test_series_fast_transform_date(self):
+        # GH 13191
+        df = pd.DataFrame({'grouping': [np.nan, 1, 1, 3],
+                           'd': pd.date_range('2014-1-1', '2014-1-4')})
+        result = df.groupby('grouping')['d'].transform('first')
+        dates = [pd.NaT, pd.Timestamp('2014-1-2'), pd.Timestamp('2014-1-2'),
+                 pd.Timestamp('2014-1-4')]
+        expected = pd.Series(dates, name='d')
+        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]})
@@ -4406,7 +4436,7 @@ def test_groupby_datetime64_32_bit(self):
 
         df = DataFrame({"A": range(2), "B": [pd.Timestamp('2000-01-1')] * 2})
         result = df.groupby("A")["B"].transform(min)
-        expected = Series([pd.Timestamp('2000-01-1')] * 2)
+        expected = Series([pd.Timestamp('2000-01-1')] * 2, name='B')
         assert_series_equal(result, expected)
 
     def test_groupby_categorical_unequal_len(self):