[BACKPORT] Fix the error raised when inferring dtype in DataFrame.transform (#1424) (#1427)

Author: hekaisheng

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
 

mars/dataframe/base/transform.py

@@ -171,21 +171,35 @@ def tile(cls, op: "TransformOperand"):
     def _infer_df_func_returns(self, in_dtypes, dtypes):
         if self.object_type == ObjectType.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.object_type == ObjectType.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._object_type = ObjectType.dataframe
         else: