Repository: mars-project/mars

Issue:
[BUG] Failed when calling Series.transform
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
Failed to run Series.transform(lambda x: x + 's').

**To Reproduce**
``` Python
In [1]: import pandas as pd 
   ...: import mars.dataframe as md 
   ...: mdf = md.Series(pd.Series(list('abc'))) 
   ...: mdf.transform(lambda x: x + 's').execute()                                                  
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~/miniconda3/lib/python3.7/site-packages/pandas/core/series.py in aggregate(self, func, axis, *args, **kwargs)
   3704             try:
-> 3705                 result = self.apply(func, *args, **kwargs)
   3706             except (ValueError, AttributeError, TypeError):

~/miniconda3/lib/python3.7/site-packages/pandas/core/series.py in apply(self, func, convert_dtype, args, **kwds)
   3847                 values = self.astype(object).values
-> 3848                 mapped = lib.map_infer(values, f, convert=convert_dtype)
   3849 

pandas/_libs/lib.pyx in pandas._libs.lib.map_infer()

<ipython-input-1-13040ef89e14> in <lambda>(x)
      3 mdf = md.Series(pd.Series(list('abc')))
----> 4 mdf.transform(lambda x: x + 's').execute()

TypeError: unsupported operand type(s) for +: 'float' and 'str'

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
<ipython-input-1-13040ef89e14> in <module>
      2 import mars.dataframe as md
      3 mdf = md.Series(pd.Series(list('abc')))
----> 4 mdf.transform(lambda x: x + 's').execute()

~/Documents/mars_dev/mars/mars/dataframe/base/transform.py in series_transform(series, func, convert_dtype, axis, dtype, *args, **kwargs)
    246                           output_types=[OutputType.series], call_agg=kwargs.pop('_call_agg', False))
    247     dtypes = (series.name, dtype) if dtype is not None else None
--> 248     return op(series, dtypes=dtypes)

~/Documents/mars_dev/mars/mars/dataframe/base/transform.py in __call__(self, df, dtypes, index)
    203             dtypes = self._infer_df_func_returns(df.dtypes, dtypes)
    204         else:
--> 205             dtypes = self._infer_df_func_returns((df.name, df.dtype), dtypes)
    206 
    207         for arg, desc in zip((self.output_types, dtypes), ('output_types', 'dtypes')):

~/Documents/mars_dev/mars/mars/dataframe/base/transform.py in _infer_df_func_returns(self, in_dtypes, dtypes)
    185                 else:
    186                     infer_df = empty_df.transform(self._func, convert_dtype=self.convert_dtype,
--> 187                                                   args=self.args, **self.kwds)
    188 
    189         if isinstance(infer_df, pd.DataFrame):

~/miniconda3/lib/python3.7/site-packages/pandas/core/series.py in transform(self, func, axis, *args, **kwargs)
   3715         # Validate the axis parameter
   3716         self._get_axis_number(axis)
-> 3717         return super().transform(func, *args, **kwargs)
   3718 
   3719     def apply(self, func, convert_dtype=True, args=(), **kwds):

~/miniconda3/lib/python3.7/site-packages/pandas/core/generic.py in transform(self, func, *args, **kwargs)
  10423     @Appender(_shared_docs["transform"] % dict(axis="", **_shared_doc_kwargs))
  10424     def transform(self, func, *args, **kwargs):
> 10425         result = self.agg(func, *args, **kwargs)
  10426         if is_scalar(result) or len(result) != len(self):
  10427             raise ValueError("transforms cannot produce aggregated results")

~/miniconda3/lib/python3.7/site-packages/pandas/core/series.py in aggregate(self, func, axis, *args, **kwargs)
   3705                 result = self.apply(func, *args, **kwargs)
   3706             except (ValueError, AttributeError, TypeError):
-> 3707                 result = func(self, *args, **kwargs)
   3708 
   3709         return result

TypeError: <lambda>() got an unexpected keyword argument 'convert_dtype'
```


Discussion:

Resolving patch:
diff --git a/mars/dataframe/base/transform.py b/mars/dataframe/base/transform.py
--- a/mars/dataframe/base/transform.py
+++ b/mars/dataframe/base/transform.py
@@ -172,21 +172,35 @@ def tile(cls, op: "TransformOperand"):
     def _infer_df_func_returns(self, in_dtypes, dtypes):
         if self.output_types[0] == OutputType.dataframe:
             empty_df = build_empty_df(in_dtypes, index=pd.RangeIndex(2))
-            with np.errstate(all='ignore'):
-                if self.call_agg:
-                    infer_df = empty_df.agg(self._func, axis=self._axis, *self.args, **self.kwds)
-                else:
-                    infer_df = empty_df.transform(self._func, axis=self._axis, *self.args, **self.kwds)
+            try:
+                with np.errstate(all='ignore'):
+                    if self.call_agg:
+                        infer_df = empty_df.agg(self._func, axis=self._axis, *self.args, **self.kwds)
+                    else:
+                        infer_df = empty_df.transform(self._func, axis=self._axis, *self.args, **self.kwds)
+            except:  # noqa: E722
+                infer_df = None
         else:
             empty_df = build_empty_series(in_dtypes[1], index=pd.RangeIndex(2), name=in_dtypes[0])
-            with np.errstate(all='ignore'):
-                if self.call_agg:
-                    infer_df = empty_df.agg(self._func, args=self.args, **self.kwds)
-                else:
-                    infer_df = empty_df.transform(self._func, convert_dtype=self.convert_dtype,
-                                                  args=self.args, **self.kwds)
+            try:
+                with np.errstate(all='ignore'):
+                    if self.call_agg:
+                        infer_df = empty_df.agg(self._func, args=self.args, **self.kwds)
+                    else:
+                        infer_df = empty_df.transform(self._func, convert_dtype=self.convert_dtype,
+                                                      args=self.args, **self.kwds)
+            except:  # noqa: E722
+                infer_df = None
+
+        if infer_df is None and dtypes is None:
+            raise TypeError('Failed to infer dtype, please specify dtypes as arguments.')
+
+        if infer_df is None:
+            is_df = self.output_types[0] == OutputType.dataframe
+        else:
+            is_df = isinstance(infer_df, pd.DataFrame)
 
-        if isinstance(infer_df, pd.DataFrame):
+        if is_df:
             new_dtypes = dtypes or infer_df.dtypes
             self.output_types = [OutputType.dataframe]
         else:


Test patch:
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
@@ -581,6 +581,23 @@ def rename_fn(f, new_name):
             result = self.executor.execute_dataframe(r, concat=True)[0]
             expected = s_raw.transform(['cumsum', lambda x: x + 1])
             pd.testing.assert_frame_equal(result, expected)
+
+            # test transform on string dtype
+            df_raw = pd.DataFrame({'col1': ['str'] * 10, 'col2': ['string'] * 10})
+            df = from_pandas_df(df_raw, chunk_size=3)
+
+            with self.assertRaises(TypeError):
+                df['col1'].transform(lambda x: x + '_suffix')
+
+            r = df.transform(lambda x: x + '_suffix')
+            result = self.executor.execute_dataframe(r, concat=True)[0]
+            expected = df_raw.transform(lambda x: x + '_suffix')
+            pd.testing.assert_frame_equal(result, expected)
+
+            r = df['col2'].transform(lambda x: x + '_suffix', dtype=np.dtype('str'))
+            result = self.executor.execute_dataframe(r, concat=True)[0]
+            expected = df_raw['col2'].transform(lambda x: x + '_suffix')
+            pd.testing.assert_series_equal(result, expected)
         finally:
             options.chunk_store_limit = old_chunk_store_limit