Keep subclassing in apply (#19823)

Author: jaumebonet

doc/source/whatsnew/v0.23.0.txt

@@ -295,8 +295,10 @@ Other Enhancements
 - ``IntervalIndex.astype`` now supports conversions between subtypes when passed an ``IntervalDtype`` (:issue:`19197`)
 - :class:`IntervalIndex` and its associated constructor methods (``from_arrays``, ``from_breaks``, ``from_tuples``) have gained a ``dtype`` parameter (:issue:`19262`)
 - Added :func:`SeriesGroupBy.is_monotonic_increasing` and :func:`SeriesGroupBy.is_monotonic_decreasing` (:issue:`17015`)
+- For subclassed ``DataFrames``, :func:`DataFrame.apply` will now preserve the ``Series`` subclass (if defined) when passing the data to the applied function (:issue:`19822`)
 - :func:`DataFrame.from_dict` now accepts a ``columns`` argument that can be used to specify the column names when ``orient='index'`` is used (:issue:`18529`)
 
+
 .. _whatsnew_0230.api_breaking:
 
 Backwards incompatible API changes

pandas/core/apply.py

@@ -162,7 +162,7 @@ def apply_empty_result(self):
  pass
 
  if reduce:
- return Series(np.nan, index=self.agg_axis)
+ return self.obj._constructor_sliced(np.nan, index=self.agg_axis)
  else:
  return self.obj.copy()
 
@@ -175,11 +175,13 @@ def apply_raw(self):
  result = np.apply_along_axis(self.f, self.axis, self.values)
 
  # TODO: mixed type case
- from pandas import DataFrame, Series
  if result.ndim == 2:
- return DataFrame(result, index=self.index, columns=self.columns)
+ return self.obj._constructor(result,
+ index=self.index,
+ columns=self.columns)
  else:
- return Series(result, index=self.agg_axis)
+ return self.obj._constructor_sliced(result,
+ index=self.agg_axis)
 
  def apply_broadcast(self, target):
  result_values = np.empty_like(target.values)
@@ -232,7 +234,7 @@ def apply_standard(self):
  axis=self.axis,
  dummy=dummy,
  labels=labels)
- return Series(result, index=labels)
+ return self.obj._constructor_sliced(result, index=labels)
  except Exception:
  pass
 
@@ -291,8 +293,7 @@ def wrap_results(self):
  return self.wrap_results_for_axis()
 
  # dict of scalars
- from pandas import Series
- result = Series(results)
+ result = self.obj._constructor_sliced(results)
  result.index = self.res_index
 
  return result
@@ -379,7 +380,6 @@ def wrap_results_for_axis(self):
  # we have a non-series and don't want inference
  elif not isinstance(results[0], ABCSeries):
  from pandas import Series
-
  result = Series(results)
  result.index = self.res_index
 

pandas/tests/frame/test_subclass.py

@@ -514,3 +514,59 @@ def test_subclassed_wide_to_long(self):
  long_frame = pd.wide_to_long(df, ["A", "B"], i="id", j="year")
 
  tm.assert_frame_equal(long_frame, expected)
+
+ def test_subclassed_apply(self):
+ # GH 19822
+
+ def check_row_subclass(row):
+ assert isinstance(row, tm.SubclassedSeries)
+
+ def strech(row):
+ if row["variable"] == "height":
+ row["value"] += 0.5
+ return row
+
+ df = tm.SubclassedDataFrame([
+ ['John', 'Doe', 'height', 5.5],
+ ['Mary', 'Bo', 'height', 6.0],
+ ['John', 'Doe', 'weight', 130],
+ ['Mary', 'Bo', 'weight', 150]],
+ columns=['first', 'last', 'variable', 'value'])
+
+ df.apply(lambda x: check_row_subclass(x))
+ df.apply(lambda x: check_row_subclass(x), axis=1)
+
+ expected = tm.SubclassedDataFrame([
+ ['John', 'Doe', 'height', 6.0],
+ ['Mary', 'Bo', 'height', 6.5],
+ ['John', 'Doe', 'weight', 130],
+ ['Mary', 'Bo', 'weight', 150]],
+ columns=['first', 'last', 'variable', 'value'])
+
+ result = df.apply(lambda x: strech(x), axis=1)
+ assert isinstance(result, tm.SubclassedDataFrame)
+ tm.assert_frame_equal(result, expected)
+
+ expected = tm.SubclassedDataFrame([
+ [1, 2, 3],
+ [1, 2, 3],
+ [1, 2, 3],
+ [1, 2, 3]])
+
+ result = df.apply(lambda x: tm.SubclassedSeries([1, 2, 3]), axis=1)
+ assert isinstance(result, tm.SubclassedDataFrame)
+ tm.assert_frame_equal(result, expected)
+
+ result = df.apply(lambda x: [1, 2, 3], axis=1, result_type="expand")
+ assert isinstance(result, tm.SubclassedDataFrame)
+ tm.assert_frame_equal(result, expected)
+
+ expected = tm.SubclassedSeries([
+ [1, 2, 3],
+ [1, 2, 3],
+ [1, 2, 3],
+ [1, 2, 3]])
+
+ result = df.apply(lambda x: [1, 2, 3], axis=1)
+ assert not isinstance(result, tm.SubclassedDataFrame)
+ tm.assert_series_equal(result, expected)