Stop calling user funcs when dtypes is specified (#2587)

Author: wjsi

mars/dataframe/base/map_chunk.py

@@ -16,7 +16,7 @@
 import pandas as pd
 
 from ... import opcodes
-from ...core import recursive_tile
+from ...core import recursive_tile, get_output_types
 from ...core.custom_log import redirect_custom_log
 from ...serialization.serializables import (
     KeyField,
@@ -90,67 +90,98 @@ def _set_inputs(self, inputs):
         super()._set_inputs(inputs)
         self._input = self._inputs[0]
 
-    def __call__(self, df_or_series, index=None, dtypes=None):
+    def _infer_attrs_by_call(self, df_or_series):
         test_obj = (
             build_df(df_or_series, size=2)
             if df_or_series.ndim == 2
             else build_series(df_or_series, size=2, name=df_or_series.name)
         )
-        output_type = self._output_types[0] if self.output_types else None
-
-        # try run to infer meta
-        try:
-            kwargs = self.kwargs or dict()
-            if self.with_chunk_index:
-                kwargs["chunk_index"] = (0,) * df_or_series.ndim
-            with np.errstate(all="ignore"), quiet_stdio():
-                obj = self._func(test_obj, *self._args, **kwargs)
-        except:  # noqa: E722  # nosec
-            if df_or_series.ndim == 1 or output_type == OutputType.series:
-                obj = pd.Series([], dtype=np.dtype(object))
-            elif output_type == OutputType.dataframe and dtypes is not None:
-                obj = build_empty_df(dtypes)
+        kwargs = self.kwargs or dict()
+        if self.with_chunk_index:
+            kwargs["chunk_index"] = (0,) * df_or_series.ndim
+        with np.errstate(all="ignore"), quiet_stdio():
+            obj = self._func(test_obj, *self._args, **kwargs)
+
+        if obj.ndim == 2:
+            output_type = OutputType.dataframe
+            dtypes = obj.dtypes
+            if obj.shape == test_obj.shape:
+                shape = (df_or_series.shape[0], len(dtypes))
+            else:  # pragma: no cover
+                shape = (np.nan, len(dtypes))
+        else:
+            output_type = OutputType.series
+            dtypes = pd.Series([obj.dtype], name=obj.name)
+            if obj.shape == test_obj.shape:
+                shape = df_or_series.shape
             else:
-                raise TypeError(
-                    "Cannot determine `output_type`, "
-                    "you have to specify it as `dataframe` or `series`, "
-                    "for dataframe, `dtypes` is required as well "
-                    "if output_type='dataframe'"
-                )
+                shape = (np.nan,)
 
-        if getattr(obj, "ndim", 0) == 1 or output_type == OutputType.series:
-            shape = self._kwargs.pop("shape", None)
-            if shape is None:
-                # series
-                if obj.shape == test_obj.shape:
-                    shape = df_or_series.shape
-                else:
-                    shape = (np.nan,)
-            if index is None:
-                index = obj.index
+        index_value = parse_index(
+            obj.index, df_or_series, self._func, self._args, self._kwargs
+        )
+        return {
+            "output_type": output_type,
+            "index_value": index_value,
+            "shape": shape,
+            "dtypes": dtypes,
+        }
+
+    def __call__(self, df_or_series, index=None, dtypes=None):
+        output_type = (
+            self.output_types[0]
+            if self.output_types
+            else get_output_types(df_or_series)[0]
+        )
+        shape = self._kwargs.pop("shape", None)
+
+        if dtypes is not None:
+            index = index if index is not None else pd.RangeIndex(-1)
             index_value = parse_index(
                 index, df_or_series, self._func, self._args, self._kwargs
             )
+            if shape is None:  # pragma: no branch
+                shape = (
+                    (np.nan,)
+                    if output_type == OutputType.series
+                    else (np.nan, len(dtypes))
+                )
+        else:
+            # try run to infer meta
+            try:
+                attrs = self._infer_attrs_by_call(df_or_series)
+                output_type = attrs["output_type"]
+                index_value = attrs["index_value"]
+                shape = attrs["shape"]
+                dtypes = attrs["dtypes"]
+            except:  # noqa: E722  # nosec
+                if df_or_series.ndim == 1 or output_type == OutputType.series:
+                    output_type = OutputType.series
+                    index = index if index is not None else pd.RangeIndex(-1)
+                    index_value = parse_index(
+                        index, df_or_series, self._func, self._args, self._kwargs
+                    )
+                    dtypes = pd.Series([np.dtype(object)])
+                    shape = (np.nan,)
+                else:
+                    raise TypeError(
+                        "Cannot determine `output_type`, "
+                        "you have to specify it as `dataframe` or `series`, "
+                        "for dataframe, `dtypes` is required as well "
+                        "if output_type='dataframe'"
+                    )
+
+        if output_type == OutputType.series:
             return self.new_series(
                 [df_or_series],
-                dtype=obj.dtype,
+                dtype=dtypes.iloc[0],
                 shape=shape,
                 index_value=index_value,
-                name=obj.name,
+                name=dtypes.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], len(dtypes))
-            else:
-                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,

mars/dataframe/base/tests/test_base_execution.py

@@ -1791,7 +1791,7 @@ def f4(pdf):
     r = df2.map_chunk(
         lambda x: x["a"].apply(pd.Series), output_type="dataframe", dtypes=dtypes
     )
-    assert r.shape == (2, 3)
+    assert r.shape == (np.nan, 3)
     pd.testing.assert_series_equal(r.dtypes, dtypes)
     result = r.execute().fetch()
     expected = raw2.apply(lambda x: x["a"], axis=1, result_type="expand")