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),