broadcast always returns original column names

pandas-dev · Feb 7, 2018 · 1d93380 · 1d93380
1 parent 8a1837c
commit 1d93380
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 18 deletions.
diff --git a/pandas/core/apply.py b/pandas/core/apply.py
@@ -187,7 +187,6 @@ def apply_broadcast(self, target):
         # axis which we want to compare compliance
         result_compare = target.shape[0]
 
-        index = None
         for i, col in enumerate(target.columns):
             res = self.f(target[col])
             ares = np. asarray(res).ndim
@@ -201,19 +200,11 @@ def apply_broadcast(self, target):
                 if result_compare != len(res):
                     raise ValueError("cannot broadcast result")
 
-                # if we have a Series result, then then index
-                # is our result
-                if isinstance(res, ABCSeries):
-                    index = res.index
-
             result_values[:, i] = res
 
-        # if we are returning a list-like
-        # then preserve the original index
-        if index is None:
-            index = target.index
-
-        result = self.obj._constructor(result_values, index=index,
+        # we *always* preserve the original index / columns
+        result = self.obj._constructor(result_values,
+                                       index=target.index,
                                        columns=target.columns)
         return result
 

diff --git a/pandas/core/frame.py b/pandas/core/frame.py
@@ -4877,10 +4877,11 @@ def apply(self, func, axis=0, broadcast=None, raw=False, reduce=None,
 
         result_type : {'expand', 'reduce', 'broadcast, None}
             These only act when axis=1 {columns}
-            * 'expand' : list-like results will be turned into columns
+            * 'expand' : list-like results will be turned into columns.
             * 'reduce' : return a Series if possible rather than expanding
-              list-like results. This is the opposite to 'expand'
-            * 'broadcast' : scalar results will be broadcast to all columns
+              list-like results. This is the opposite to 'expand'.
+            * 'broadcast' : results will be broadcast to the original shape
+              of the frame, the original index & columns will be retained.
             * None : list-like results will be returned as a list
               in a single column. However if the apply function
               returns a Series these are expanded to columns.

diff --git a/pandas/tests/frame/test_apply.py b/pandas/tests/frame/test_apply.py
@@ -171,14 +171,12 @@ def test_apply_broadcast(self):
                           result_type='broadcast')
         tm.assert_frame_equal(result, df)
 
-        # columms come from the returned Series
         df = DataFrame(np.tile(np.arange(3), 6).reshape(6, -1) + 1,
                        columns=list('ABC'))
         result = df.apply(lambda x: Series([1, 2, 3], index=list('abc')),
                           axis=1,
                           result_type='broadcast')
         expected = df.copy()
-        expected.columns = list('abc')
         tm.assert_frame_equal(result, expected)
 
     def test_apply_broadcast_error(self):
@@ -756,7 +754,6 @@ def test_result_type(self):
             axis=1,
             result_type='broadcast')
         expected = df.copy()
-        expected.columns = columns
         assert_frame_equal(result, expected)
 
         # series result