Repository: mars-project/mars

Issue:
[BUG] `output_type` does not take effect for `DataFrame.apply`
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

`output_type` does not take effect for `DataFrame.apply`.

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version
2. The version of Mars you use
3. Versions of crucial packages, such as numpy, scipy and protobuf
4. Full stack of the error.
5. Minimized code to reproduce the error.

```python
In [1]: import numpy as np                                                                                                                                                                                  

In [2]: import pandas as pd                                                                                                                                                                                 

In [3]: import mars.dataframe as md                                                                                                                                                                         

In [4]: df = pd.DataFrame({'a': [np.array([1, 2, 3]), np.array([4, 5, 6])]})                                                                                                                                

In [5]: mdf = md.DataFrame(df)                                                                                                                                                                              

In [6]: dtypes = pd.Series([np.dtype(float)] * 3)                                                                                                                                                           

In [13]: df2 = mdf.apply(lambda x: x['a'].tolist(), axis=1, output_type='dataframe', dtypes=dtypes, result_type='expand')                                                                                   

In [14]: df2                                                                                                                                                                                                
Out[14]: Series(op=ApplyOperand)  # Should be DataFrame
```




Discussion:

Resolving patch:
diff --git a/mars/dataframe/base/apply.py b/mars/dataframe/base/apply.py
--- a/mars/dataframe/base/apply.py
+++ b/mars/dataframe/base/apply.py
@@ -172,8 +172,12 @@ def _tile_series(cls, op):
             new_op = op.copy().reset_key()
             new_op._tileable_op_key = op.key
             kw = c.params.copy()
-            kw['dtype'] = out_series.dtype
-            if out_series.ndim == 2:
+            if out_series.ndim == 1:
+                kw['dtype'] = out_series.dtype
+            else:
+                kw['index'] = (c.index[0], 0)
+                kw['shape'] = (c.shape[0], out_series.shape[1])
+                kw['dtypes'] = out_series.dtypes
                 kw['columns_value'] = out_series.columns_value
             chunks.append(new_op.new_chunk([c], **kw))
 
@@ -181,6 +185,7 @@ def _tile_series(cls, op):
         kw = out_series.params.copy()
         kw.update(dict(chunks=chunks, nsplits=in_series.nsplits))
         if out_series.ndim == 2:
+            kw['nsplits'] = (in_series.nsplits[0], (out_series.shape[1],))
             kw['columns_value'] = out_series.columns_value
         return new_op.new_tileables(op.inputs, **kw)
 
@@ -225,18 +230,6 @@ def _infer_df_func_returns(self, df, dtypes, index):
         self._elementwise = new_elementwise if self._elementwise is None else self._elementwise
         return dtypes, index_value
 
-    def _infer_series_func_returns(self, df):
-        try:
-            empty_series = build_series(df, size=2, name=df.name)
-            with np.errstate(all='ignore'):
-                infer_series = empty_series.apply(self._func, args=self.args, **self.kwds)
-            new_dtype = infer_series.dtype
-            name = infer_series.name
-        except:  # noqa: E722  # nosec  # pylint: disable=bare-except
-            new_dtype = np.dtype('object')
-            name = None
-        return new_dtype, name
-
     def _call_dataframe(self, df, dtypes=None, index=None):
         dtypes, index_value = self._infer_df_func_returns(df, dtypes, index)
         if index_value is None:
@@ -268,22 +261,62 @@ def _call_dataframe(self, df, dtypes=None, index=None):
         else:
             return self.new_series([df], shape=shape, dtype=dtypes, index_value=index_value)
 
-    def _call_series(self, series):
+    def _call_series(self, series, dtype=None, index=None):
         if self._convert_dtype:
-            dtype, name = self._infer_series_func_returns(series)
+            try:
+                test_series = build_series(series, size=2, name=series.name)
+                with np.errstate(all='ignore'):
+                    infer_series = test_series.apply(self._func, args=self.args, **self.kwds)
+            except:  # noqa: E722  # nosec  # pylint: disable=bare-except
+                infer_series = None
+
+            output_type = self._output_types[0]
+
+            if index is not None:
+                index_value = parse_index(index)
+            elif infer_series is not None:
+                index_value = parse_index(infer_series.index)
+            else:
+                index_value = parse_index(None, series)
+
+            if output_type == OutputType.dataframe:
+                if dtype is not None:
+                    dtypes = dtype
+                elif infer_series is not None and infer_series.ndim == 2:
+                    dtypes = infer_series.dtypes
+                else:
+                    raise TypeError('Cannot determine dtypes, '
+                                    'please specify `dtypes` as argument')
+
+                columns_value = parse_index(dtypes.index, store_data=True)
+
+                return self.new_dataframe([series], shape=(series.shape[0], len(dtypes)),
+                                          index_value=index_value, columns_value=columns_value,
+                                          dtypes=dtypes)
+            else:
+                if dtype is None and infer_series is not None and infer_series.ndim == 1:
+                    dtype = infer_series.dtype
+                else:
+                    dtype = np.dtype(object)
+                if infer_series is not None and infer_series.ndim == 1:
+                    name = infer_series.name
+                else:
+                    name = None
+                return self.new_series([series], dtype=dtype, shape=series.shape,
+                                       index_value=index_value, name=name)
         else:
             dtype, name = np.dtype('object'), None
-        return self.new_series([series], dtype=dtype, shape=series.shape,
-                               index_value=series.index_value, name=name)
+            return self.new_series([series], dtype=dtype, shape=series.shape,
+                                   index_value=series.index_value, name=name)
 
-    def __call__(self, df, dtypes=None, index=None):
+    def __call__(self, df_or_series, dtypes=None, index=None):
         axis = getattr(self, 'axis', None) or 0
-        self._axis = validate_axis(axis, df)
+        self._axis = validate_axis(axis, df_or_series)
 
-        if df.op.output_types[0] == OutputType.dataframe:
-            return self._call_dataframe(df, dtypes=dtypes, index=index)
+        if df_or_series.op.output_types[0] == OutputType.dataframe:
+            return self._call_dataframe(df_or_series, dtypes=dtypes, index=index)
         else:
-            return self._call_series(df)
+            return self._call_series(df_or_series, dtype=dtypes, index=index)
 
 
 def df_apply(df, func, axis=0, raw=False, result_type=None, args=(), dtypes=None,
@@ -306,11 +339,12 @@ def df_apply(df, func, axis=0, raw=False, result_type=None, args=(), dtypes=None
         return func(*args, **kwds)
 
     op = ApplyOperand(func=func, axis=axis, raw=raw, result_type=result_type, args=args, kwds=kwds,
-                      output_types=output_type, elementwise=elementwise)
+                      output_type=output_type, elementwise=elementwise)
     return op(df, dtypes=dtypes, index=index)
 
 
-def series_apply(series, func, convert_dtype=True, args=(), **kwds):
+def series_apply(series, func, convert_dtype=True, output_type=None,
+                 args=(), index=None, **kwds):
     if isinstance(func, (list, dict)):
         return series.aggregate(func)
 
@@ -319,10 +353,19 @@ def series_apply(series, func, convert_dtype=True, args=(), **kwds):
         func_body = getattr(series, func, None)
         if func_body is not None:
             return func_body(*args, **kwds)
-        func = getattr(np, func, None)
+        func_str = func
+        func = getattr(np, func_str, None)
         if func is None:
-            raise AttributeError(f"'{func!r}' is not a valid function for '{type(series.__name__)}' object")
+            raise AttributeError(f"'{func_str!r}' is not a valid function "
+                                 f"for '{type(series).__name__}' object")
+
+    output_types = kwds.pop('output_types', None)
+    object_type = kwds.pop('object_type', None)
+    output_types = validate_output_types(
+        output_type=output_type, output_types=output_types, object_type=object_type)
+    output_type = output_types[0] if output_types else OutputType.series
+    dtypes = kwds.pop('dtypes', kwds.pop('dtype', None))
 
     op = ApplyOperand(func=func, convert_dtype=convert_dtype, args=args, kwds=kwds,
-                      output_type=OutputType.series)
-    return op(series)
+                      output_type=output_type)
+    return op(series, dtypes=dtypes, index=index)
diff --git a/mars/dataframe/base/map_chunk.py b/mars/dataframe/base/map_chunk.py
--- a/mars/dataframe/base/map_chunk.py
+++ b/mars/dataframe/base/map_chunk.py
@@ -101,18 +101,19 @@ def __call__(self, df_or_series, index=None, dtypes=None):
                                    shape=shape, index_value=index_value,
                                    name=obj.name)
         else:
+            dtypes = dtypes if dtypes is not None else obj.dtypes
             # dataframe
             if obj.shape == test_obj.shape:
-                shape = (df_or_series.shape[0], obj.shape[1])
+                shape = (df_or_series.shape[0], len(dtypes))
             else:
-                shape = (np.nan, obj.shape[1])
-            columns_value = parse_index(obj.dtypes.index, store_data=True)
+                shape = (np.nan, len(dtypes))
+            columns_value = parse_index(dtypes.index, store_data=True)
             if index is None:
                 index = obj.index
             index_value = parse_index(index, df_or_series,
                                       self._func, self._args, self._kwargs)
             return self.new_dataframe([df_or_series], shape=shape,
-                                      dtypes=obj.dtypes, index_value=index_value,
+                                      dtypes=dtypes, index_value=index_value,
                                       columns_value=columns_value)
 
     @classmethod


Test patch:
diff --git a/mars/dataframe/base/tests/test_base.py b/mars/dataframe/base/tests/test_base.py
--- a/mars/dataframe/base/tests/test_base.py
+++ b/mars/dataframe/base/tests/test_base.py
@@ -350,6 +350,12 @@ def testDataFrameApply(self):
         finally:
             options.chunk_store_limit = old_chunk_store_limit
 
+        raw = pd.DataFrame({'a': [np.array([1, 2, 3]), np.array([4, 5, 6])]})
+        df = from_pandas_df(raw)
+        df2 = df.apply(lambda x: x['a'].astype(pd.Series), axis=1,
+                       output_type='dataframe', dtypes=pd.Series([np.dtype(float)] * 3))
+        self.assertEqual(df2.ndim, 2)
+
     def testSeriesApply(self):
         idxes = [chr(ord('A') + i) for i in range(20)]
         s_raw = pd.Series([i ** 2 for i in range(20)], index=idxes)
@@ -383,6 +389,37 @@ def testSeriesApply(self):
         self.assertEqual(r.chunks[0].shape, (5,))
         self.assertEqual(r.chunks[0].inputs[0].shape, (5,))
 
+        s_raw2 = pd.Series([np.array([1, 2, 3]), np.array([4, 5, 6])])
+        series = from_pandas_series(s_raw2)
+
+        r = series.apply(np.sum)
+        self.assertEqual(r.dtype, np.dtype(object))
+
+        r = series.apply(lambda x: pd.Series([1]), output_type='dataframe')
+        expected = s_raw2.apply(lambda x: pd.Series([1]))
+        pd.testing.assert_series_equal(r.dtypes, expected.dtypes)
+
+        dtypes = pd.Series([np.dtype(float)] * 3)
+        r = series.apply(pd.Series, output_type='dataframe',
+                         dtypes=dtypes)
+        self.assertEqual(r.ndim, 2)
+        pd.testing.assert_series_equal(r.dtypes, dtypes)
+        self.assertEqual(r.shape, (2, 3))
+
+        r = series.apply(pd.Series, output_type='dataframe',
+                         dtypes=dtypes, index=pd.RangeIndex(2))
+        self.assertEqual(r.ndim, 2)
+        pd.testing.assert_series_equal(r.dtypes, dtypes)
+        self.assertEqual(r.shape, (2, 3))
+
+        with self.assertRaises(AttributeError) as cm:
+            series.apply('abc')
+        self.assertIn('abc', str(cm.exception))
+
+        with self.assertRaises(TypeError):
+            # dtypes not provided
+            series.apply(lambda x: x.tolist(), output_type='dataframe')
+
     def testTransform(self):
         cols = [chr(ord('A') + i) for i in range(10)]
         df_raw = pd.DataFrame(dict((c, [i ** 2 for i in range(20)]) for c in cols))
diff --git a/mars/dataframe/base/tests/test_base_execution.py b/mars/dataframe/base/tests/test_base_execution.py
--- a/mars/dataframe/base/tests/test_base_execution.py
+++ b/mars/dataframe/base/tests/test_base_execution.py
@@ -495,6 +495,16 @@ def testSeriesApplyExecute(self):
         expected = s_raw.apply(lambda x: [x, x + 1], convert_dtype=False)
         pd.testing.assert_series_equal(result, expected)
 
+        s_raw2 = pd.Series([np.array([1, 2, 3]), np.array([4, 5, 6])])
+        series = from_pandas_series(s_raw2)
+
+        dtypes = pd.Series([np.dtype(float)] * 3)
+        r = series.apply(pd.Series, output_type='dataframe',
+                         dtypes=dtypes)
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = s_raw2.apply(pd.Series)
+        pd.testing.assert_frame_equal(result, expected)
+
     @unittest.skipIf(pa is None, 'pyarrow not installed')
     def testApplyWithArrowDtypeExecution(self):
         df1 = pd.DataFrame({'a': [1, 2, 1],
@@ -1645,6 +1655,17 @@ def f4(pdf):
         expected = f4(raw)
         pd.testing.assert_frame_equal(result, expected)
 
+        raw2 = pd.DataFrame({'a': [np.array([1, 2, 3]), np.array([4, 5, 6])]})
+        df2 = from_pandas_df(raw2)
+        dtypes = pd.Series([np.dtype(float)] * 3)
+        r = df2.map_chunk(lambda x: x['a'].apply(pd.Series), output_type='dataframe',
+                          dtypes=dtypes)
+        self.assertEqual(r.shape, (2, 3))
+        pd.testing.assert_series_equal(r.dtypes, dtypes)
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = raw2.apply(lambda x: x['a'], axis=1, result_type='expand')
+        pd.testing.assert_frame_equal(result, expected)
+
     def testRebalanceExecution(self):
         raw = pd.DataFrame(np.random.rand(10, 3), columns=list('abc'))
         df = from_pandas_df(raw)