Spawn serialization of executable graphs

mars-project · Dec 9, 2020 · aae932e · aae932e
1 parent 0f2c63e
commit aae932e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/mars/dataframe/reduction/aggregation.py b/mars/dataframe/reduction/aggregation.py
@@ -605,7 +605,11 @@ def _execute_map(cls, ctx, op: "DataFrameAggregate"):
 
         # map according to map groups
         ret_map_dfs = dict()
+        in_cols_set = set(in_data.columns) if in_data.ndim == 2 else None
         for input_key, output_key, cols, func in op.pre_funcs:
+            if cols and in_cols_set == set(cols):
+                cols = None
+
             src_df = in_data if cols is None else in_data[cols]
             if input_key == output_key:
                 ret_map_dfs[output_key] = src_df

diff --git a/mars/dataframe/reduction/core.py b/mars/dataframe/reduction/core.py
@@ -870,7 +870,7 @@ def compile(self) -> ReductionSteps:
 
         for key, step in self._output_key_to_post_steps.items():
             cols = self._output_key_to_post_cols[key]
-            if set(cols) == set(referred_cols):
+            if cols and set(cols) == set(referred_cols):
                 post_cols = None
             else:
                 post_cols = cols