Repository: mars-project/mars

Issue:
[BUG] Optimization of filter-setitem failed
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

Optimization of filter-setitem failed.

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version
2. The version of Mars you use
3. Versions of crucial packages, such as numpy, scipy and pandas
4. Full stack of the error.
5. Minimized code to reproduce the error.

```
In [13]: df = md.DataFrame(mt.random.rand(10, 4, chunk_size=4), columns=list('ab
    ...: cd'))

In [14]: df['e'] = df['a'] * df['b'] - 2

In [15]: df['f'] = df['e'] * 2 - df['c']

In [16]: df.groupby('a').sum().execute()
Unexpected error happens in <function TaskProcessor.optimize at 0x7ff2c31a1280>
Traceback (most recent call last):
  File "/Users/xuyeqin/Workspace/mars/mars/services/task/supervisor/processor.py", line 57, in inner
    return await func(processor, *args, **kwargs)
  File "/Users/xuyeqin/Workspace/mars/mars/services/task/supervisor/processor.py", line 131, in optimize
    return await asyncio.to_thread(self._preprocessor.optimize)
  File "/Users/xuyeqin/Workspace/mars/mars/lib/aio/_threads.py", line 36, in to_thread
    return await loop.run_in_executor(None, func_call)
  File "/Users/xuyeqin/miniconda3/envs/mars3.8/lib/python3.8/concurrent/futures/thread.py", line 57, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/Users/xuyeqin/Workspace/mars/mars/services/task/supervisor/preprocessor.py", line 127, in optimize
    self.tileable_optimization_records = optimize_tileable_graph(
  File "/Users/xuyeqin/Workspace/mars/mars/optimization/logical/tileable/core.py", line 36, in optimize
    return TileableOptimizer.optimize(tileable_graph)
  File "/Users/xuyeqin/Workspace/mars/mars/core/mode.py", line 77, in _inner
    return func(*args, **kwargs)
  File "/Users/xuyeqin/Workspace/mars/mars/optimization/logical/core.py", line 271, in optimize
    rule.apply(op)
  File "/Users/xuyeqin/Workspace/mars/mars/optimization/logical/tileable/arithmetic_query.py", line 260, in apply
    self._graph.add_edge(in_tileable, new_node)
  File "mars/core/graph/core.pyx", line 93, in mars.core.graph.core.DirectedGraph.add_edge
  File "mars/core/graph/core.pyx", line 100, in mars.core.graph.core.DirectedGraph._add_edge
KeyError: 'Node DataFrame(op=DataFrameSetitem) does not exist in the directed graph'
  0%|                                                   | 0/100 [00:00<?, ?it/s]
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-16-aafe9450034d> in <module>
----> 1 df.groupby('a').sum().execute()

~/Workspace/mars/mars/core/entity/tileables.py in execute(self, session, **kw)
    460 
    461     def execute(self, session=None, **kw):
--> 462         result = self.data.execute(session=session, **kw)
    463         if isinstance(result, TILEABLE_TYPE):
    464             return self

~/Workspace/mars/mars/core/entity/executable.py in execute(self, session, **kw)
     96 
     97         session = _get_session(self, session)
---> 98         return execute(self, session=session, **kw)
     99 
    100     def _check_session(self, session: SessionType, action: str):

~/Workspace/mars/mars/deploy/oscar/session.py in execute(tileable, session, wait, new_session_kwargs, show_progress, progress_update_interval, *tileables, **kwargs)
   1771         session = get_default_or_create(**(new_session_kwargs or dict()))
   1772     session = _ensure_sync(session)
-> 1773     return session.execute(
   1774         tileable,
   1775         *tileables,

~/Workspace/mars/mars/deploy/oscar/session.py in execute(self, tileable, show_progress, *tileables, **kwargs)
   1571         fut = asyncio.run_coroutine_threadsafe(coro, self._loop)
   1572         try:
-> 1573             execution_info: ExecutionInfo = fut.result(
   1574                 timeout=self._isolated_session.timeout
   1575             )

~/miniconda3/envs/mars3.8/lib/python3.8/concurrent/futures/_base.py in result(self, timeout)
    437                 raise CancelledError()
    438             elif self._state == FINISHED:
--> 439                 return self.__get_result()
    440             else:
    441                 raise TimeoutError()

~/miniconda3/envs/mars3.8/lib/python3.8/concurrent/futures/_base.py in __get_result(self)
    386     def __get_result(self):
    387         if self._exception:
--> 388             raise self._exception
    389         else:
    390             return self._result

~/Workspace/mars/mars/deploy/oscar/session.py in _execute(session, wait, show_progress, progress_update_interval, cancelled, *tileables, **kwargs)
   1722                 while not cancelled.is_set():
   1723                     try:
-> 1724                         await asyncio.wait_for(
   1725                             asyncio.shield(execution_info), progress_update_interval
   1726                         )

~/miniconda3/envs/mars3.8/lib/python3.8/asyncio/tasks.py in wait_for(fut, timeout, loop)
    481 
    482         if fut.done():
--> 483             return fut.result()
    484         else:
    485             fut.remove_done_callback(cb)

~/miniconda3/envs/mars3.8/lib/python3.8/asyncio/tasks.py in _wrap_awaitable(awaitable)
    682     that will later be wrapped in a Task by ensure_future().
    683     """
--> 684     return (yield from awaitable.__await__())
    685 
    686 _wrap_awaitable._is_coroutine = _is_coroutine

~/Workspace/mars/mars/deploy/oscar/session.py in wait()
    100 
    101             async def wait():
--> 102                 return await self._aio_task
    103 
    104             self._future_local.future = fut = asyncio.run_coroutine_threadsafe(

~/Workspace/mars/mars/deploy/oscar/session.py in _run_in_background(self, tileables, task_id, progress, profiling)
    901                     )
    902                 if task_result.error:
--> 903                     raise task_result.error.with_traceback(task_result.traceback)
    904             if cancelled:
    905                 return

~/Workspace/mars/mars/services/task/supervisor/processor.py in inner(processor, *args, **kwargs)
     55     async def inner(processor: "TaskProcessor", *args, **kwargs):
     56         try:
---> 57             return await func(processor, *args, **kwargs)
     58         except:  # noqa: E722  # nosec  # pylint: disable=bare-except  # pragma: no cover
     59             if log_when_error:

~/Workspace/mars/mars/services/task/supervisor/processor.py in optimize(self)
    129         # optimization, run it in executor,
    130         # since optimization may be a CPU intensive operation
--> 131         return await asyncio.to_thread(self._preprocessor.optimize)
    132 
    133     @_record_error

~/Workspace/mars/mars/lib/aio/_threads.py in to_thread(func, *args, **kwargs)
     34     ctx = contextvars.copy_context()
     35     func_call = functools.partial(ctx.run, func, *args, **kwargs)
---> 36     return await loop.run_in_executor(None, func_call)

~/miniconda3/envs/mars3.8/lib/python3.8/concurrent/futures/thread.py in run(self)
     55 
     56         try:
---> 57             result = self.fn(*self.args, **self.kwargs)
     58         except BaseException as exc:
     59             self.future.set_exception(exc)

~/Workspace/mars/mars/services/task/supervisor/preprocessor.py in optimize(self)
    125         if self._config.optimize_tileable_graph:
    126             # enable optimization
--> 127             self.tileable_optimization_records = optimize_tileable_graph(
    128                 self.tileable_graph
    129             )

~/Workspace/mars/mars/optimization/logical/tileable/core.py in optimize(tileable_graph)
     34 
     35 def optimize(tileable_graph: TileableGraph) -> OptimizationRecords:
---> 36     return TileableOptimizer.optimize(tileable_graph)

~/Workspace/mars/mars/core/mode.py in _inner(*args, **kwargs)
     75             def _inner(*args, **kwargs):
     76                 with enter_mode(**mode_name_to_value):
---> 77                     return func(*args, **kwargs)
     78 
     79         else:

~/Workspace/mars/mars/optimization/logical/core.py in optimize(cls, graph)
    269                 if rule.match(op):
    270                     optimized = True
--> 271                     rule.apply(op)
    272         if optimized:
    273             cls._replace_inputs(graph, records)

~/Workspace/mars/mars/optimization/logical/tileable/arithmetic_query.py in apply(self, op)
    258 
    259         self._replace_node(node, new_node)
--> 260         self._graph.add_edge(in_tileable, new_node)
    261         self._records.append_record(
    262             OptimizationRecord(node, new_node, OptimizationRecordType.replace)

~/Workspace/mars/mars/core/graph/core.pyx in mars.core.graph.core.DirectedGraph.add_edge()

~/Workspace/mars/mars/core/graph/core.pyx in mars.core.graph.core.DirectedGraph._add_edge()

KeyError: 'Node DataFrame(op=DataFrameSetitem) does not exist in the directed graph'
```




Discussion:

Resolving patch:
diff --git a/mars/optimization/logical/core.py b/mars/optimization/logical/core.py
--- a/mars/optimization/logical/core.py
+++ b/mars/optimization/logical/core.py
@@ -17,7 +17,7 @@
 from collections import defaultdict
 from dataclasses import dataclass
 from enum import Enum
-from typing import Dict, List, Tuple, Type
+from typing import Dict, List, Optional, Tuple, Type
 
 from ...core import OperandType, ChunkType, EntityType, enter_mode
 from ...core.graph import EntityGraph
@@ -59,10 +59,12 @@ def append_record(self, record: OptimizationRecord):
         ):
             self._optimized_chunk_to_records[record.new_chunk] = record
 
-    def get_optimization_result(self, original_chunk: ChunkType) -> ChunkType:
+    def get_optimization_result(
+        self, original_chunk: ChunkType, default: Optional[ChunkType] = None
+    ) -> ChunkType:
         chunk = original_chunk
         if chunk not in self._original_chunk_to_records:
-            return
+            return default
         while chunk in self._original_chunk_to_records:
             record = self._original_chunk_to_records[chunk]
             if record.record_type == OptimizationRecordType.replace:
@@ -72,10 +74,12 @@ def get_optimization_result(self, original_chunk: ChunkType) -> ChunkType:
                 return None
         return chunk
 
-    def get_original_chunk(self, optimized_chunk: ChunkType) -> ChunkType:
+    def get_original_chunk(
+        self, optimized_chunk: ChunkType, default: Optional[ChunkType] = None
+    ) -> ChunkType:
         chunk = optimized_chunk
         if chunk not in self._optimized_chunk_to_records:
-            return
+            return default
         while chunk in self._optimized_chunk_to_records:
             record = self._optimized_chunk_to_records[chunk]
             if record.record_type == OptimizationRecordType.replace:
@@ -151,28 +155,25 @@ def _replace_node(self, original_node: EntityType, new_node: EntityType):
         for succ in successors:
             self._graph.add_edge(new_node, succ)
 
-    @classmethod
-    def _add_collapsable_predecessor(cls, node: EntityType, predecessor: EntityType):
-        if predecessor not in cls._preds_to_remove:
-            cls._preds_to_remove[predecessor] = {node}
+    def _add_collapsable_predecessor(self, node: EntityType, predecessor: EntityType):
+        pred_original = self._records.get_original_chunk(predecessor, predecessor)
+        if predecessor not in self._preds_to_remove:
+            self._preds_to_remove[pred_original] = {node}
         else:
-            cls._preds_to_remove[predecessor].add(node)
+            self._preds_to_remove[pred_original].add(node)
 
     def _remove_collapsable_predecessors(self, node: EntityType):
         node = self._records.get_optimization_result(node) or node
         preds_opt_to_remove = []
         for pred in self._graph.predecessors(node):
-            pred_original = self._records.get_original_chunk(pred)
-            pred_original = pred_original if pred_original is not None else pred
-
-            pred_opt = self._records.get_optimization_result(pred)
-            pred_opt = pred_opt if pred_opt is not None else pred
+            pred_original = self._records.get_original_chunk(pred, pred)
+            pred_opt = self._records.get_optimization_result(pred, pred)
 
             if pred_opt in self._graph.results or pred_original in self._graph.results:
                 continue
             affect_succ = self._preds_to_remove.get(pred_original) or []
             affect_succ_opt = [
-                self._records.get_optimization_result(s) or s for s in affect_succ
+                self._records.get_optimization_result(s, s) for s in affect_succ
             ]
             if all(s in affect_succ_opt for s in self._graph.successors(pred)):
                 preds_opt_to_remove.append((pred_original, pred_opt))
diff --git a/mars/optimization/logical/tileable/arithmetic_query.py b/mars/optimization/logical/tileable/arithmetic_query.py
--- a/mars/optimization/logical/tileable/arithmetic_query.py
+++ b/mars/optimization/logical/tileable/arithmetic_query.py
@@ -105,13 +105,12 @@ def _is_select_dataframe_column(tileable) -> bool:
             and index_op.mask is None
         )
 
-    @classmethod
-    def _extract_eval_expression(cls, tileable) -> EvalExtractRecord:
+    def _extract_eval_expression(self, tileable) -> EvalExtractRecord:
         if is_scalar(tileable):
             if isinstance(tileable, (int, bool, str, bytes, np.integer, np.bool_)):
                 return EvalExtractRecord(expr=repr(tileable))
             else:
-                var_name = f"__eval_scalar_var{cls._next_var_id()}"
+                var_name = f"__eval_scalar_var{self._next_var_id()}"
                 var_dict = {var_name: tileable}
                 return EvalExtractRecord(expr=f"@{var_name}", variables=var_dict)
 
@@ -121,15 +120,15 @@ def _extract_eval_expression(cls, tileable) -> EvalExtractRecord:
         if tileable in _extract_result_cache:
             return _extract_result_cache[tileable]
 
-        if cls._is_select_dataframe_column(tileable):
-            result = cls._extract_column_select(tileable)
+        if self._is_select_dataframe_column(tileable):
+            result = self._extract_column_select(tileable)
         elif isinstance(tileable.op, DataFrameUnaryUfunc):
-            result = cls._extract_unary(tileable)
+            result = self._extract_unary(tileable)
         elif isinstance(tileable.op, DataFrameBinopUfunc):
             if tileable.op.fill_value is not None or tileable.op.level is not None:
                 result = EvalExtractRecord()
             else:
-                result = cls._extract_binary(tileable)
+                result = self._extract_binary(tileable)
         else:
             result = EvalExtractRecord()
 
@@ -140,35 +139,33 @@ def _extract_eval_expression(cls, tileable) -> EvalExtractRecord:
     def _extract_column_select(cls, tileable) -> EvalExtractRecord:
         return EvalExtractRecord(tileable.inputs[0], f"`{tileable.op.col_names}`")
 
-    @classmethod
-    def _extract_unary(cls, tileable) -> EvalExtractRecord:
+    def _extract_unary(self, tileable) -> EvalExtractRecord:
         op = tileable.op
         func_name = getattr(op, "_func_name") or getattr(op, "_bin_func_name")
         if func_name not in _func_name_to_builder:  # pragma: no cover
             return EvalExtractRecord()
 
-        in_tileable, expr, variables = cls._extract_eval_expression(op.inputs[0])
+        in_tileable, expr, variables = self._extract_eval_expression(op.inputs[0])
         if in_tileable is None:
             return EvalExtractRecord()
 
-        cls._add_collapsable_predecessor(tileable, op.inputs[0])
+        self._add_collapsable_predecessor(tileable, op.inputs[0])
         return EvalExtractRecord(
             in_tileable, _func_name_to_builder[func_name](expr), variables
         )
 
-    @classmethod
-    def _extract_binary(cls, tileable) -> EvalExtractRecord:
+    def _extract_binary(self, tileable) -> EvalExtractRecord:
         op = tileable.op
         func_name = getattr(op, "_func_name", None) or getattr(op, "_bit_func_name")
         if func_name not in _func_name_to_builder:  # pragma: no cover
             return EvalExtractRecord()
 
-        lhs_tileable, lhs_expr, lhs_vars = cls._extract_eval_expression(op.lhs)
+        lhs_tileable, lhs_expr, lhs_vars = self._extract_eval_expression(op.lhs)
         if lhs_tileable is not None:
-            cls._add_collapsable_predecessor(tileable, op.lhs)
-        rhs_tileable, rhs_expr, rhs_vars = cls._extract_eval_expression(op.rhs)
+            self._add_collapsable_predecessor(tileable, op.lhs)
+        rhs_tileable, rhs_expr, rhs_vars = self._extract_eval_expression(op.rhs)
         if rhs_tileable is not None:
-            cls._add_collapsable_predecessor(tileable, op.rhs)
+            self._add_collapsable_predecessor(tileable, op.rhs)
 
         if lhs_expr is None or rhs_expr is None:
             return EvalExtractRecord()
@@ -190,6 +187,9 @@ def _extract_binary(cls, tileable) -> EvalExtractRecord:
     def apply(self, op: OperandType):
         node = op.outputs[0]
         in_tileable, expr, variables = self._extract_eval_expression(node)
+        opt_in_tileable = self._records.get_optimization_result(
+            in_tileable, in_tileable
+        )
 
         new_op = DataFrameEval(
             _key=node.op.key,
@@ -199,13 +199,13 @@ def apply(self, op: OperandType):
             parser="pandas",
             is_query=False,
         )
-        new_node = new_op.new_tileable([in_tileable], **node.params).data
-        new_node._key = node.key
-        new_node._id = node.id
+        new_node = new_op.new_tileable(
+            [opt_in_tileable], _key=node.key, _id=node.id, **node.params
+        ).data
 
         self._remove_collapsable_predecessors(node)
         self._replace_node(node, new_node)
-        self._graph.add_edge(in_tileable, new_node)
+        self._graph.add_edge(opt_in_tileable, new_node)
 
         self._records.append_record(
             OptimizationRecord(node, new_node, OptimizationRecordType.replace)
@@ -241,34 +241,40 @@ def _get_optimized_eval_op(self, op: OperandType) -> OperandType:
     def _get_input_columnar_node(self, op: OperandType) -> ENTITY_TYPE:
         raise NotImplementedError
 
-    def apply(self, op: DataFrameIndex):
-        node = op.outputs[0]
-        in_tileable = op.inputs[0]
-        in_columnar_node = self._get_input_columnar_node(op)
-
-        new_op = self._build_new_eval_op(op)
-        new_op._key = node.op.key
-
-        new_node = new_op.new_tileable([in_tileable], **node.params).data
-        new_node._key = node.key
-        new_node._id = node.id
-
-        self._add_collapsable_predecessor(node, in_columnar_node)
-        self._remove_collapsable_predecessors(node)
+    def _update_op_node(self, old_node: ENTITY_TYPE, new_node: ENTITY_TYPE):
+        self._replace_node(old_node, new_node)
+        for in_tileable in new_node.inputs:
+            self._graph.add_edge(in_tileable, new_node)
 
-        self._replace_node(node, new_node)
-        self._graph.add_edge(in_tileable, new_node)
+        original_node = self._records.get_original_chunk(old_node, old_node)
         self._records.append_record(
-            OptimizationRecord(node, new_node, OptimizationRecordType.replace)
+            OptimizationRecord(original_node, new_node, OptimizationRecordType.replace)
         )
 
         # check node if it's in result
         try:
-            i = self._graph.results.index(node)
+            i = self._graph.results.index(old_node)
             self._graph.results[i] = new_node
         except ValueError:
             pass
 
+    def apply(self, op: DataFrameIndex):
+        node = op.outputs[0]
+        in_tileable = op.inputs[0]
+        in_columnar_node = self._get_input_columnar_node(op)
+        opt_in_tileable = self._records.get_optimization_result(
+            in_tileable, in_tileable
+        )
+
+        new_op = self._build_new_eval_op(op)
+        new_node = new_op.new_tileable(
+            [opt_in_tileable], _key=node.key, _id=node.id, **node.params
+        ).data
+
+        self._add_collapsable_predecessor(node, in_columnar_node)
+        self._remove_collapsable_predecessors(node)
+        self._update_op_node(node, new_node)
+
 
 @register_tileable_optimization_rule([DataFrameIndex])
 class DataFrameBoolEvalToQuery(_DataFrameEvalRewriteRule):
@@ -287,6 +293,7 @@ def _get_input_columnar_node(self, op: OperandType) -> ENTITY_TYPE:
     def _build_new_eval_op(self, op: OperandType):
         in_eval_op = self._get_optimized_eval_op(op)
         return DataFrameEval(
+            _key=op.key,
             _output_types=get_output_types(op.outputs[0]),
             expr=in_eval_op.expr,
             variables=in_eval_op.variables,
@@ -308,6 +315,7 @@ def _get_input_columnar_node(self, op: DataFrameSetitem) -> ENTITY_TYPE:
     def _build_new_eval_op(self, op: DataFrameSetitem):
         in_eval_op = self._get_optimized_eval_op(op)
         return DataFrameEval(
+            _key=op.key,
             _output_types=get_output_types(op.outputs[0]),
             expr=f"`{op.indexes}` = {in_eval_op.expr}",
             variables=in_eval_op.variables,
@@ -315,3 +323,43 @@ def _build_new_eval_op(self, op: DataFrameSetitem):
             is_query=False,
             self_target=True,
         )
+
+    def apply(self, op: DataFrameIndex):
+        super().apply(op)
+
+        node = op.outputs[0]
+        opt_node = self._records.get_optimization_result(node, node)
+        if not isinstance(opt_node.op, DataFrameEval):  # pragma: no cover
+            return
+
+        # when encountering consecutive SetItems, expressions can be
+        # merged as a multiline expression
+        pred_opt_node = opt_node.inputs[0]
+        if (
+            isinstance(pred_opt_node.op, DataFrameEval)
+            and opt_node.op.parser == pred_opt_node.op.parser == "pandas"
+            and not opt_node.op.is_query
+            and not pred_opt_node.op.is_query
+            and opt_node.op.self_target
+            and pred_opt_node.op.self_target
+        ):
+            new_expr = pred_opt_node.op.expr + "\n" + opt_node.op.expr
+            new_variables = (pred_opt_node.op.variables or dict()).copy()
+            new_variables.update(opt_node.op.variables or dict())
+
+            new_op = DataFrameEval(
+                _key=op.key,
+                _output_types=get_output_types(op.outputs[0]),
+                expr=new_expr,
+                variables=new_variables,
+                parser="pandas",
+                is_query=False,
+                self_target=True,
+            )
+            new_node = new_op.new_tileable(
+                pred_opt_node.inputs, _key=node.key, _id=node.id, **node.params
+            ).data
+
+            self._add_collapsable_predecessor(opt_node, pred_opt_node)
+            self._remove_collapsable_predecessors(opt_node)
+            self._update_op_node(opt_node, new_node)


Test patch:
diff --git a/mars/optimization/logical/tileable/tests/test_arithmetic_query.py b/mars/optimization/logical/tileable/tests/test_arithmetic_query.py
--- a/mars/optimization/logical/tileable/tests/test_arithmetic_query.py
+++ b/mars/optimization/logical/tileable/tests/test_arithmetic_query.py
@@ -153,14 +153,25 @@ def test_eval_setitem_to_eval(setup):
     df2 = md.DataFrame(raw2, chunk_size=10)
     df3 = df1.merge(df2, on="A", suffixes=("", "_"))
     df3["K"] = df3["A"] * (1 - df3["B"])
+    df3["L"] = df3["K"] - df3["A"]
+    df3["M"] = df3["K"] + df3["L"]
+
     graph = TileableGraph([df3.data])
     next(TileableGraphBuilder(graph).build())
     records = optimize(graph)
     opt_df3 = records.get_optimization_result(df3.data)
-    assert opt_df3.op.expr == "`K` = (`A`) * ((1) - (`B`))"
+    assert opt_df3.op.expr == "\n".join(
+        [
+            "`K` = (`A`) * ((1) - (`B`))",
+            "`L` = (`K`) - (`A`)",
+            "`M` = (`K`) + (`L`)",
+        ]
+    )
     assert len(graph) == 4
     assert len([n for n in graph if isinstance(n.op, DataFrameEval)]) == 1
 
     r_df3 = raw.merge(raw2, on="A", suffixes=("", "_"))
     r_df3["K"] = r_df3["A"] * (1 - r_df3["B"])
+    r_df3["L"] = r_df3["K"] - r_df3["A"]
+    r_df3["M"] = r_df3["K"] + r_df3["L"]
     pd.testing.assert_frame_equal(df3.execute().fetch(), r_df3)