Fix duplicate execution (#3301)

* Fix duplicate execute

* Fix

Co-authored-by: 刘宝 <po.lb@antgroup.com>

mars/core/graph/builder/chunk.py

@@ -124,7 +124,7 @@ def __init__(
         self,
         tileable_graph: TileableGraph,
         tile_context: TileContext,
-        processed_chunks: Set[ChunkType],
+        processed_chunks: Set[str],
         chunk_to_fetch: Dict[ChunkType, ChunkType],
         add_nodes: Callable,
     ):
@@ -301,11 +301,12 @@ def _iter(self):
 
         if chunk_graph is not None:
             # last tiled chunks, add them to processed
-            # so that fetch chunk can be generated
-            processed_chunks = [
-                c.chunk if isinstance(c, FUSE_CHUNK_TYPE) else c
+            # so that fetch chunk can be generated.
+            # Use chunk key as the key to make sure the copied chunk can be build to a fetch.
+            processed_chunks = (
+                c.chunk.key if isinstance(c, FUSE_CHUNK_TYPE) else c.key
                 for c in chunk_graph.result_chunks
-            ]
+            )
             self._processed_chunks.update(processed_chunks)
 
         result_chunks = []
@@ -389,7 +390,7 @@ def __init__(
         self.tile_context = TileContext() if tile_context is None else tile_context
         self.tile_context.set_tileables(set(graph))
 
-        self._processed_chunks: Set[ChunkType] = set()
+        self._processed_chunks: Set[str] = set()
         self._chunk_to_fetch: Dict[ChunkType, ChunkType] = dict()
 
         tiler_cls = Tiler if tiler_cls is None else tiler_cls
@@ -402,7 +403,7 @@ def __init__(
         )
 
     def _process_node(self, entity: EntityType):
-        if entity in self._processed_chunks:
+        if entity.key in self._processed_chunks:
             if entity not in self._chunk_to_fetch:
                 # gen fetch
                 fetch_chunk = build_fetch(entity).data
@@ -413,7 +414,7 @@ def _process_node(self, entity: EntityType):
     def _select_inputs(self, inputs: List[ChunkType]):
         new_inputs = []
         for inp in inputs:
-            if inp in self._processed_chunks:
+            if inp.key in self._processed_chunks:
                 # gen fetch
                 if inp not in self._chunk_to_fetch:
                     fetch_chunk = build_fetch(inp).data
@@ -424,7 +425,7 @@ def _select_inputs(self, inputs: List[ChunkType]):
         return new_inputs
 
     def _if_add_node(self, node: EntityType, visited: Set):
-        return node not in visited and node not in self._processed_chunks
+        return node not in visited and node.key not in self._processed_chunks
 
     def _build(self) -> Iterable[Union[TileableGraph, ChunkGraph]]:
         tile_iterator = iter(self.tiler)

mars/dataframe/base/rechunk.py

@@ -157,7 +157,7 @@ def tile(cls, op: "DataFrameRechunk"):
                 params["dtypes"] = pd.concat([c.dtypes for c in inp_chunks_arr[0]])
             if len(inp_slice_chunks) == 1:
                 c = inp_slice_chunks[0]
-                cc = c.op.copy().reset_key().new_chunk(c.op.inputs, kws=[params])
+                cc = c.op.copy().new_chunk(c.op.inputs, kws=[params])
                 out_chunks.append(cc)
             else:
                 out_chunk = DataFrameConcat(

mars/dataframe/base/tests/test_base_execution.py

@@ -91,6 +91,15 @@ def test_to_cpu_execution(setup_gpu):
 
 
 def test_rechunk_execution(setup):
+    ns = np.random.RandomState(0)
+    df = pd.DataFrame(ns.rand(100, 10), columns=["a" + str(i) for i in range(10)])
+
+    # test rechunk after sort
+    mdf = DataFrame(df, chunk_size=10)
+    result = mdf.sort_values("a0").rechunk(chunk_size=10).execute().fetch()
+    expected = df.sort_values("a0")
+    pd.testing.assert_frame_equal(result, expected)
+
     data = pd.DataFrame(np.random.rand(8, 10))
     df = from_pandas_df(pd.DataFrame(data), chunk_size=3)
     df2 = df.rechunk((3, 4))

mars/dataframe/sort/core.py

@@ -85,7 +85,7 @@ def _tile_head(cls, op: "DataFrameSortOperand"):
                 shape = tuple(shape)
                 concat_params["shape"] = shape
                 if len(to_combine_chunks) == 1:
-                    c = to_combine_chunks[0]
+                    c = to_combine_chunks[0].copy()
                     c._index = chunk_index
                 else:
                     c = DataFrameConcat(

mars/dataframe/utils.py

@@ -1412,7 +1412,7 @@ def _concat_chunks(merge_chunks: List[ChunkType], output_index: int):
             # concat previous chunks
             if len(to_merge_chunks) == 1:
                 # do not generate concat op for 1 input.
-                c = to_merge_chunks[0]
+                c = to_merge_chunks[0].copy()
                 c._index = (
                     (len(n_split),) if df_or_series.ndim == 1 else (len(n_split), 0)
                 )

mars/services/task/supervisor/preprocessor.py

@@ -36,7 +36,7 @@ def __init__(
         self,
         tileable_graph: TileableGraph,
         tile_context: TileContext,
-        processed_chunks: Set[ChunkType],
+        processed_chunks: Set[str],
         chunk_to_fetch: Dict[ChunkType, ChunkType],
         add_nodes: Callable,
         cancelled: asyncio.Event = None,