Repository: mars-project/mars

Issue:
[BUG] DataFrameIndexAlign failed on ray shuffle with AssertionError
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
Mars `DataFrameIndexAlign` operand failed on ray shuffle with AssertionError on check assert len(subtask_chunk_graph.result_chunks) == 1:
```
def _get_subtask_out_info(
    subtask_chunk_graph: ChunkGraph, is_mapper: bool, n_reducers: int = None
):
    # output_keys might be duplicate in chunk graph, use dict to deduplicate.
    # output_keys order should be consistent with remote `execute_subtask`,
    # dict can preserve insert order.
    output_keys = {}
    for chunk in subtask_chunk_graph.result_chunks:
        if isinstance(
            chunk.op, VirtualOperand
        ):  # FIXME(chaokunyang) no need to check this?
            continue
        elif is_mapper:
            assert (
                len(subtask_chunk_graph.result_chunks) == 1
            ), subtask_chunk_graph.result_chunks
            assert n_reducers is not None
            return set(), n_reducers
        else:
            output_keys[chunk.key] = 1
    return output_keys.keys(), len(output_keys)
```
![graphviz (6)](https://user-images.githubusercontent.com/12445254/185032948-52deeb6b-8ccb-4019-8021-53dcc065baa2.svg)


**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version: 3.8
2. The version of Mars you use: master
3. Versions of crucial packages, such as numpy, scipy and pandas
4. Full stack of the error.
```
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
mars/deploy/oscar/session.py:1890: in execute
    return session.execute(
mars/deploy/oscar/session.py:1684: in execute
    execution_info: ExecutionInfo = fut.result(
../../../../anaconda3/envs/mars3.8/lib/python3.8/concurrent/futures/_base.py:444: in result
    return self.__get_result()
../../../../anaconda3/envs/mars3.8/lib/python3.8/concurrent/futures/_base.py:389: in __get_result
    raise self._exception
mars/deploy/oscar/session.py:1870: in _execute
    await execution_info
../../../../anaconda3/envs/mars3.8/lib/python3.8/asyncio/tasks.py:695: in _wrap_awaitable
    return (yield from awaitable.__await__())
mars/deploy/oscar/session.py:105: in wait
    return await self._aio_task
mars/deploy/oscar/session.py:953: in _run_in_background
    raise task_result.error.with_traceback(task_result.traceback)
mars/services/task/supervisor/processor.py:369: in run
    await self._process_stage_chunk_graph(*stage_args)
mars/services/task/supervisor/processor.py:247: in _process_stage_chunk_graph
    chunk_to_result = await self._executor.execute_subtask_graph(
mars/services/task/execution/ray/executor.py:507: in execute_subtask_graph
    output_keys, out_count = _get_subtask_out_info(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

subtask_chunk_graph = <mars.core.graph.entity.ChunkGraph object at 0x7fea91c08e00>, is_mapper = True, n_reducers = 3

    def _get_subtask_out_info(
        subtask_chunk_graph: ChunkGraph, is_mapper: bool, n_reducers: int = None
    ):
        # output_keys might be duplicate in chunk graph, use dict to deduplicate.
        # output_keys order should be consistent with remote `execute_subtask`,
        # dict can preserve insert order.
        output_keys = {}
        for chunk in subtask_chunk_graph.result_chunks:
            if isinstance(
                chunk.op, VirtualOperand
            ):  # FIXME(chaokunyang) no need to check this?
                continue
            elif is_mapper:
>               assert (
                    len(subtask_chunk_graph.result_chunks) == 1
                ), subtask_chunk_graph.result_chunks
E               AssertionError: [DataFrameChunkData <op=DataFrameIndexAlign, stage=map, key=c7a56e28260a78c9b936be2c72feb298_0>, DataFrameChunkData <op=DataFrameIndexAlign, stage=map, key=ceab58328a4b66ef381aceeb58b8cf0f_0>, DataFrameChunkData <op=DataFrameIndexAlign, stage=map, key=c738082cdc61d25a1ed37e5fcabe0b3e_0>]

mars/services/task/execution/ray/executor.py:259: AssertionError

```
5. Minimized code to reproduce the error.
` MARS_CI_BACKEND=ray pytest -v -s mars/dataframe/indexing/tests/test_indexing_execution.py::test_align_execution`
**Expected behavior**
A clear and concise description of what you expected to happen.

**Additional context**
This bug is different from #3221, which the reducer is mapper too. The  root cause here is that mars fusion #2435 fused mutiple mapper into one subtask, which breaks `Fetch_BY_INDEX` shuffle mode. We need to extends mars fusion to make mapper isolated from mapper siblings.



Discussion:

Resolving patch:
diff --git a/mars/services/task/analyzer/analyzer.py b/mars/services/task/analyzer/analyzer.py
--- a/mars/services/task/analyzer/analyzer.py
+++ b/mars/services/task/analyzer/analyzer.py
@@ -439,6 +439,26 @@ def gen_subtask_graph(
         # states
         visited = set()
         logic_key_to_subtasks = defaultdict(list)
+        if self._shuffle_fetch_type == ShuffleFetchType.FETCH_BY_INDEX:
+            for chunk in self._chunk_graph.topological_iter():
+                if not isinstance(chunk.op, ShuffleProxy):
+                    continue
+                # Can't use `OperandStage.map` to find mappers directly, since `stage` of some operand
+                # such as `DataFrameIndexAlign` are `OperandStage.map` but not a shuffle mapper sometimes.
+                mapper_chunks = self._chunk_graph.predecessors(chunk)
+                for mapper_chunk in mapper_chunks:
+                    chunk_color = chunk_to_colors[mapper_chunk]
+                    same_color_chunks = color_to_chunks[chunk_color]
+                    mappers = [
+                        c for c in same_color_chunks if c.op.stage == OperandStage.map
+                    ]
+                    if len(mappers) > 1:
+                        # ensure every subtask contains only at most one mapper
+                        for mapper in mappers:
+                            same_color_chunks.remove(mapper)
+                            mapper_color = coloring.next_color()
+                            chunk_to_colors[mapper] = mapper_color
+                            color_to_chunks[mapper_color] = [mapper]
         for chunk in self._chunk_graph.topological_iter():
             if chunk in visited or isinstance(chunk.op, Fetch):
                 # skip fetch chunk
diff --git a/mars/services/task/analyzer/fusion.py b/mars/services/task/analyzer/fusion.py
--- a/mars/services/task/analyzer/fusion.py
+++ b/mars/services/task/analyzer/fusion.py
@@ -48,7 +48,7 @@ def __init__(
 
         self._coloring_iter = itertools.count()
 
-    def _next_color(self) -> int:
+    def next_color(self) -> int:
         return next(self._coloring_iter)
 
     @classmethod
@@ -82,7 +82,7 @@ def _color_init_nodes(self) -> Dict[OperandType, int]:
         for chunk, band in self.chunk_to_bands.items():
             band_priority = (band, chunk.op.priority)
             if band_priority not in band_priority_to_colors:
-                band_priority_to_colors[band_priority] = self._next_color()
+                band_priority_to_colors[band_priority] = self.next_color()
 
         band_priority_to_color_list = defaultdict(list)
         for (band, priority), color in band_priority_to_colors.items():
@@ -94,7 +94,7 @@ def _color_init_nodes(self) -> Dict[OperandType, int]:
             color = band_priority_to_color_list[band, priority][-1]
             size = color_to_size[color]
             if size >= self.initial_same_color_num:
-                color = self._next_color()
+                color = self.next_color()
                 band_priority_to_color_list[band, priority].append(color)
             color_to_size[color] += 1
             op_to_colors[chunk.op] = color
@@ -125,17 +125,17 @@ def color(self) -> Dict[ChunkType, int]:
             if len(predecessors) == 1 and predecessors[0] in broadcaster_chunk_set:
                 # TODO: handle situation that chunks which specify reassign_workers
                 # predecessor is broadcaster, just allocate a new color
-                color = self._next_color()
+                color = self.next_color()
             elif len(pred_colors) == 1:
                 if self._can_color_same(chunk, predecessors):
                     # predecessors have only 1 color, will color with same one
                     color = next(iter(pred_colors))
                 else:
-                    color = self._next_color()
+                    color = self.next_color()
             else:
                 # has more than 1 color, color a new one
                 assert len(pred_colors) > 1
-                color = self._next_color()
+                color = self.next_color()
 
             op_to_colors[chunk.op] = chunk_to_colors[chunk] = color
 
@@ -155,7 +155,7 @@ def color(self) -> Dict[ChunkType, int]:
                 stack = []
                 for succ in self.chunk_graph.iter_successors(chunk):
                     if chunk_to_colors[succ] == chunk_color:
-                        new_color = op_to_colors[succ.op] = self._next_color()
+                        new_color = op_to_colors[succ.op] = self.next_color()
                         for c in succ.op.outputs:
                             if c not in self.chunk_graph:  # pragma: no cover
                                 continue
@@ -178,7 +178,7 @@ def color(self) -> Dict[ChunkType, int]:
                         if node_input_same_color:
                             node_new_color = node_pred_colors[0]
                         else:
-                            node_new_color = self._next_color()
+                            node_new_color = self.next_color()
                         op_to_colors[node.op] = node_new_color
                         for c in node.op.outputs:
                             if c not in self.chunk_graph:  # pragma: no cover


Test patch:
diff --git a/mars/dataframe/indexing/tests/test_indexing_execution.py b/mars/dataframe/indexing/tests/test_indexing_execution.py
--- a/mars/dataframe/indexing/tests/test_indexing_execution.py
+++ b/mars/dataframe/indexing/tests/test_indexing_execution.py
@@ -1625,8 +1625,7 @@ def test_add_prefix_suffix(setup):
     pd.testing.assert_series_equal(r.execute().fetch(), raw.add_suffix("_item"))
 
 
-# TODO https://github.com/mars-project/mars/issues/3221
-# @pytest.mark.ray_dag
+@pytest.mark.ray_dag
 @pytest.mark.parametrize("join", ["outer", "left"])
 def test_align_execution(setup, join):
     rs = np.random.RandomState(0)