clean up broadcast tests

Author: jreback

doc/source/basics.rst

@@ -793,8 +793,12 @@ The :meth:`~DataFrame.apply` method will also dispatch on a string method name.
    df.apply('mean')
    df.apply('mean', axis=1)
 
-Depending on the return type of the function passed to :meth:`~DataFrame.apply`,
-the result will either be of lower dimension or the same dimension.
+The return type of the function passed to :meth:`~DataFrame.apply` affects the
+type of the ultimate output from DataFrame.apply
+
+* If the applied function returns a ``Series``, the ultimate output is a ``DataFrame``.
+  The columns match the index ``Series`` returned by the applied function.
+* If the applied function returns any other type, the ultimate output is a ``Series``.
 
 :meth:`~DataFrame.apply` combined with some cleverness can be used to answer many questions
 about a data set. For example, suppose we wanted to extract the date where the

doc/source/whatsnew/v0.23.0.txt

@@ -347,7 +347,8 @@ where a list-like (e.g. ``tuple`` or ``list`` is returned), (:issue:`16353`, :is
     df = pd.DataFrame(np.tile(np.arange(3), 6).reshape(6, -1) + 1, columns=['A', 'B', 'C'])
     df
 
-Previous Behavior. If the returned shape happened to match the index, this would return a list-like.
+Previous Behavior. If the returned shape happened to match the original columns, this would return a ``DataFrame``.
+If the return shape did not match, a ``Series`` with lists was returned.
 
 .. code-block:: python
 
@@ -379,12 +380,25 @@ New Behavior. The behavior is consistent. These will *always* return a ``Series`
     df.apply(lambda x: [1, 2, 3], axis=1)
     df.apply(lambda x: [1, 2], axis=1)
 
-To have automatic inference, you can use ``result_type='infer'``
+To have expanded columns, you can use ``result_type='infer'``
 
 .. ipython:: python
 
     df.apply(lambda x: [1, 2, 3], axis=1, result_type='infer')
 
+To have broadcast the result across, you can use ``result_type='broadcast'``. The shape
+must match the original columns.
+
+.. ipython:: python
+
+    df.apply(lambda x: [1, 2, 3], axis=1, result_type='broadcast')
+
+Returning a ``Series`` allows one to control the exact return structure and column names:
+
+.. ipython:: python
+
+    df.apply(lambda x: Series([1, 2, 3], index=x.index), axis=1)
+
 
 .. _whatsnew_0230.api_breaking.build_changes:
 

pandas/core/apply.py

@@ -43,7 +43,8 @@ def __init__(self, obj, func, broadcast, raw, reduce, result_type,
         if broadcast is not None:
             warnings.warn("The broadcast argument is deprecated and will "
                           "be removed in a future version. You can specify "
-                          "result_type='broadcast' broadcast a scalar result",
+                          "result_type='broadcast' to broadcast the result "
+                          "to the original dimensions",
                           FutureWarning, stacklevel=4)
             if broadcast:
                 result_type = 'broadcast'
@@ -160,11 +161,32 @@ def apply_raw(self):
 
     def apply_broadcast(self, target):
         result_values = np.empty_like(target.values)
-        columns = target.columns
-        for i, col in enumerate(columns):
-            result_values[:, i] = self.f(target[col])
 
-        result = self.obj._constructor(result_values, index=target.index,
+        # axis which we want to compare compliance
+        result_compare = target.shape[0]
+
+        index = target.index
+        for i, col in enumerate(target.columns):
+            res = self.f(target[col])
+            ares = np. asarray(res).ndim
+
+            # must be a scalar or 1d
+            if ares > 1:
+                raise ValueError("too many dims to broadcast")
+            elif ares == 1:
+
+                # must match return dim
+                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
+
+        result = self.obj._constructor(result_values, index=index,
                                        columns=target.columns)
         return result
 

pandas/core/frame.py

@@ -4870,10 +4870,11 @@ def apply(self, func, axis=0, broadcast=None, raw=False, reduce=None,
             while guessing, exceptions raised by func will be ignored). If
             reduce is True a Series will always be returned, and if False a
             DataFrame will always be returned.
+
         result_type : {'infer', 'broadcast, None}
             These only act when axis=1 {columns}
-            * infer : list-like results will be turned into columns
-            * broadcast : scalar results will be broadcast to all rows
+            * 'infer' : list-like results will be turned into columns
+            * 'broadcast' : scalar results will be broadcast to all columns
             * 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.

pandas/core/sparse/frame.py

@@ -855,8 +855,8 @@ def apply(self, func, axis=0, broadcast=None, reduce=False,
 
         result_type : {'infer', 'broadcast, None}
             These only act when axis=1 {columns}
-            * infer : list-like results will be turned into columns
-            * broadcast : scalar results will be broadcast to all rows
+            * 'infer' : list-like results will be turned into columns
+            * 'broadcast' : scalar results will be broadcast to all columns
             * 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.

pandas/tests/frame/test_apply.py

@@ -121,22 +121,58 @@ def test_with_string_args(self):
             expected = getattr(self.frame, arg)(axis=1)
             tm.assert_series_equal(result, expected)
 
-    def test_apply_broadcast(self):
+    def test_apply_broadcast_deprecated(self):
         with tm.assert_produces_warning(FutureWarning):
-            broadcasted = self.frame.apply(np.mean, broadcast=True)
-        agged = self.frame.apply(np.mean)
+            self.frame.apply(np.mean, broadcast=True)
 
-        for col, ts in compat.iteritems(broadcasted):
-            assert (ts == agged[col]).all()
+    def test_apply_broadcast(self):
 
-        with tm.assert_produces_warning(FutureWarning):
-            broadcasted = self.frame.apply(np.mean, axis=1, broadcast=True)
-        agged = self.frame.apply(np.mean, axis=1)
-        for idx in broadcasted.index:
-            assert (broadcasted.xs(idx) == agged[idx]).all()
+        # scalars
+        result = self.frame.apply(np.mean, result_type='broadcast')
+        expected = DataFrame([self.frame.mean()], index=self.frame.index)
+        tm.assert_frame_equal(result, expected)
 
-        with tm.assert_produces_warning(FutureWarning):
-            self.frame.apply(np.mean, axis=1, broadcast=False)
+        result = self.frame.apply(np.mean, axis=1, result_type='broadcast')
+        m = self.frame.mean(axis=1)
+        expected = DataFrame({c: m for c in self.frame.columns})
+        tm.assert_frame_equal(result, expected)
+
+        # lists
+        result = self.frame.apply(
+            lambda x: list(range(len(self.frame.columns))),
+            axis=1,
+            result_type='broadcast')
+        m = list(range(len(self.frame.columns)))
+        expected = DataFrame([m] * len(self.frame.index),
+                             dtype='float64',
+                             index=self.frame.index,
+                             columns=self.frame.columns)
+        tm.assert_frame_equal(result, expected)
+
+        result = self.frame.apply(lambda x: list(range(len(self.frame.index))),
+                                  result_type='broadcast')
+        m = list(range(len(self.frame.index)))
+        expected = DataFrame({c: m for c in self.frame.columns},
+                             dtype='float64',
+                             index=self.frame.index)
+        tm.assert_frame_equal(result, expected)
+
+    def test_apply_broadcast_error(self):
+        df = DataFrame(
+            np.tile(np.arange(3, dtype='int64'), 6).reshape(6, -1) + 1,
+            columns=['A', 'B', 'C'])
+
+        # > 1 ndim
+        with pytest.raises(ValueError):
+            df.apply(lambda x: np.array([1, 2]).reshape(-1, 2),
+                     axis=1,
+                     result_type='broadcast')
+
+        # cannot broadcast
+        with pytest.raises(ValueError):
+            df.apply(lambda x: [1, 2],
+                     axis=1,
+                     result_type='broadcast')
 
     def test_apply_raw(self):
         result0 = self.frame.apply(np.mean, raw=True)
@@ -213,8 +249,7 @@ def _checkit(axis=0, raw=False):
             _check(no_index, lambda x: x)
             _check(no_index, lambda x: x.mean())
 
-        with tm.assert_produces_warning(FutureWarning):
-            result = no_cols.apply(lambda x: x.mean(), broadcast=True)
+        result = no_cols.apply(lambda x: x.mean(), result_type='broadcast')
         assert isinstance(result, DataFrame)
 
     def test_apply_with_args_kwds(self):
@@ -680,6 +715,35 @@ def test_result_type(self):
         expected.columns = [0, 1]
         assert_frame_equal(result, expected)
 
+        # broadcast result
+        result = df.apply(lambda x: [1, 2, 3], axis=1, result_type='broadcast')
+        expected = df.copy()
+        assert_frame_equal(result, expected)
+
+        columns = ['other', 'col', 'names']
+        result = df.apply(
+            lambda x: pd.Series([1, 2, 3],
+                                index=columns),
+            axis=1,
+            result_type='broadcast')
+        expected = df.copy()
+        expected.columns = columns
+        assert_frame_equal(result, expected)
+
+        # series result
+        result = df.apply(lambda x: Series([1, 2, 3], index=x.index), axis=1)
+        expected = df.copy()
+        assert_frame_equal(result, expected)
+
+        # series result with other index
+        columns = ['other', 'col', 'names']
+        result = df.apply(
+            lambda x: pd.Series([1, 2, 3], index=columns),
+            axis=1)
+        expected = df.copy()
+        expected.columns = columns
+        assert_frame_equal(result, expected)
+
     @pytest.mark.parametrize(
         "box",
         [lambda x: list(x),