We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] Column pruning leads to a wrong result
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
A clear and concise description of what the bug is.

**To Reproduce**
``` Python
pd_df = pd.DataFrame({'a': [3, 4, 5, 3, 5, 4, 1, 2, 3],
                      'b': [1, 3, 4, 5, 6, 5, 4, 4, 4],
                      'c': list('aabaaddce'),
                      'd': list('abaaaddce')})
pd_df.to_csv('test.csv', index=False)
in_df = md.read_csv('test.csv')
df1 = in_df.groupby('d').agg({'b': 'min'})
df2 = in_df[in_df.d.isin(df1.index)]
df2.execute()
Out[5]: 
   b  d
0  1  a
1  3  b
2  4  a
3  5  a
4  6  a
5  5  d
6  4  d
7  4  c
8  4  e
```

the output only has two columns which should be 4.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mars/optimizes/tileable_graph/column_pruning.py`
Content:
```
1 #!/usr/bin/env python
2 # -*- coding: utf-8 -*-
3 # Copyright 1999-2020 Alibaba Group Holding Ltd.
4 #
5 # Licensed under the Apache License, Version 2.0 (the "License");
6 # you may not use this file except in compliance with the License.
7 # You may obtain a copy of the License at
8 #
9 #      http://www.apache.org/licenses/LICENSE-2.0
10 #
11 # Unless required by applicable law or agreed to in writing, software
12 # distributed under the License is distributed on an "AS IS" BASIS,
13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 # See the License for the specific language governing permissions and
15 # limitations under the License.
16 
17 from ...utils import copy_tileables
18 from ...dataframe.utils import parse_index
19 from ...dataframe.groupby.aggregation import DataFrameGroupByAgg
20 from ...dataframe.datasource.read_csv import DataFrameReadCSV
21 from .core import TileableOptimizeRule, register
22 
23 
24 class GroupbyPruneReadCSV(TileableOptimizeRule):
25     """
26     An experimental implementation for tileable optimization rule.
27     This rule works only when groupby aggregation operation follows the read CSV files,
28     we can prune the columns that not used by the following operations when read the files.
29     """
30     def match(self, node):
31         if isinstance(node.inputs[0].op, DataFrameReadCSV) and \
32                 node.inputs[0] not in self._optimizer_context.result_tileables:
33             by_data = node.op.groupby_params.get('by')
34             by_columns = by_data if isinstance(by_data, (list, tuple)) else [by_data]
35             if isinstance(node.op.func, (str, list)):
36                 # Passing func name(s) means perform on all columns.
37                 return False
38             elif len(set(by_columns + list(node.op.func))) == len(node.inputs[0].op.usecols or node.inputs[0].dtypes):
39                 # If performs on all columns, no need to prune.
40                 return False
41             return True
42         return False
43 
44     def apply(self, node):
45         by_data = node.op.groupby_params.get('by')
46         by_columns = by_data if isinstance(by_data, (list, tuple)) else [by_data]
47         agg_columns = list(node.op.func)
48         input_node = node.inputs[0]
49         selected_columns = [c for c in list(input_node.dtypes.index) if c in by_columns + agg_columns]
50         if input_node in self._optimizer_context:
51             new_input = self._optimizer_context[input_node]
52             selected_columns = [
53                 c for c in list(input_node.dtypes.index) if c in selected_columns + new_input.op._usecols]
54         else:
55             new_input = copy_tileables([input_node])[0].data
56 
57         new_input._shape = (input_node.shape[0], len(selected_columns))
58         new_input._dtypes = input_node.dtypes[selected_columns]
59         new_input._columns_value = parse_index(new_input._dtypes.index, store_data=True)
60         new_input.op._usecols = selected_columns
61         new_node = copy_tileables([node], inputs=[new_input])[0].data
62 
63         self._optimizer_context[node] = new_node
64         self._optimizer_context[input_node] = new_input
65         return new_node
66 
67 
68 register(DataFrameGroupByAgg, GroupbyPruneReadCSV)
69
```

Path: `mars/dataframe/base/standardize_range_index.py`
Content:
```
1 #!/usr/bin/env python
2 # -*- coding: utf-8 -*-
3 # Copyright 1999-2020 Alibaba Group Holding Ltd.
4 #
5 # Licensed under the Apache License, Version 2.0 (the "License");
6 # you may not use this file except in compliance with the License.
7 # You may obtain a copy of the License at
8 #
9 #      http://www.apache.org/licenses/LICENSE-2.0
10 #
11 # Unless required by applicable law or agreed to in writing, software
12 # distributed under the License is distributed on an "AS IS" BASIS,
13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 # See the License for the specific language governing permissions and
15 # limitations under the License.
16 
17 import pandas as pd
18 
19 from ... import opcodes as OperandDef
20 from ...utils import lazy_import
21 from ...serialize import Int32Field
22 from ..operands import DataFrameOperandMixin, DataFrameOperand
23 
24 
25 cudf = lazy_import('cudf', globals=globals())
26 
27 
28 class ChunkStandardizeRangeIndex(DataFrameOperand, DataFrameOperandMixin):
29     _op_type_ = OperandDef.STANDARDIZE_RANGE_INDEX
30 
31     _axis = Int32Field('axis')
32 
33     def __init__(self, prepare_inputs=None, axis=None, object_type=None, **kwargs):
34         super().__init__(__prepare_inputs=prepare_inputs, _axis=axis, _object_type=object_type, **kwargs)
35 
36     @property
37     def axis(self):
38         return self._axis
39 
40     @classmethod
41     def execute(cls, ctx, op):
42         xdf = cudf if op.gpu else pd
43         in_data = ctx[op.inputs[-1].key].copy()
44         input_keys = [c.key for c in op.inputs[:-1]]
45         metas = ctx.get_chunk_metas(input_keys)
46         index_start = sum([m.chunk_shape[op.axis] for m in metas])
47         if op.axis == 0:
48             in_data.index = xdf.RangeIndex(index_start, index_start + len(in_data))
49         else:
50             in_data.columns = xdf.RangeIndex(index_start, index_start + in_data.shape[1])
51         ctx[op.outputs[0].key] = in_data
52
```

Path: `mars/optimizes/tileable_graph/core.py`
Content:
```
1 #!/usr/bin/env python
2 # -*- coding: utf-8 -*-
3 # Copyright 1999-2020 Alibaba Group Holding Ltd.
4 #
5 # Licensed under the Apache License, Version 2.0 (the "License");
6 # you may not use this file except in compliance with the License.
7 # You may obtain a copy of the License at
8 #
9 #      http://www.apache.org/licenses/LICENSE-2.0
10 #
11 # Unless required by applicable law or agreed to in writing, software
12 # distributed under the License is distributed on an "AS IS" BASIS,
13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 # See the License for the specific language governing permissions and
15 # limitations under the License.
16 
17 import weakref
18 from collections import defaultdict
19 
20 from ...graph import DAG
21 from ...graph_builder import TileableGraphBuilder
22 from ...utils import copy_tileables, kernel_mode, enter_build_mode
23 
24 _rules = defaultdict(list)
25 
26 tileable_optimized = weakref.WeakKeyDictionary()
27 
28 
29 class TileableOptimizeRule(object):
30     def __init__(self, optimized_context):
31         self._optimizer_context = optimized_context
32 
33     def match(self, node):
34         raise NotImplementedError
35 
36     def apply(self, node):
37         raise NotImplementedError
38 
39 
40 class OptimizeContext(weakref.WeakKeyDictionary):
41     def __init__(self, dict=None):
42         weakref.WeakKeyDictionary.__init__(self, dict=dict)
43         self._result_tileables = []
44 
45     @property
46     def result_tileables(self):
47         return self._result_tileables
48 
49     def append_result_tileables(self, tileables):
50         self._result_tileables.extend(tileables)
51 
52 
53 class OptimizeIntegratedTileableGraphBuilder(TileableGraphBuilder):
54     def __init__(self, **kw):
55         self._optimizer_context = OptimizeContext()
56         super().__init__(**kw)
57         self._node_processor = self._apply_rules(self._node_processor, self._optimizer_context)
58 
59     @staticmethod
60     def _apply_rules(node_processor, optimizer_context):
61         def inner(node):
62             node = node_processor(node) if node_processor is not None else node
63             if type(node.op) in _rules:
64                 for rule in _rules[type(node.op)]:
65                     ruler = rule(optimizer_context)
66                     if ruler.match(node):
67                         node = rule(optimizer_context).apply(node)
68             return node
69 
70         return inner
71 
72     def _mapping_tileables(self, tileables):
73         for t in tileables:
74             if t in self._optimizer_context:
75                 tileable_optimized[t] = self._optimizer_context[t]
76 
77     def _replace_copied_tilebale(self, graph):
78         if len(self._optimizer_context) == 0:
79             return graph
80 
81         new_graph = DAG()
82         reversed_mapping = weakref.WeakKeyDictionary((v, k) for k, v in self._optimizer_context.items())
83         for n in graph.topological_iter():
84             if n in reversed_mapping:
85                 new_node = n
86             elif any(inp in self._optimizer_context for inp in n.inputs):
87                 new_inputs = [self._optimizer_context.get(i, i) for i in n.inputs]
88                 new_tileables = copy_tileables(n.op.outputs, inputs=new_inputs)
89                 for t, new_t in zip(n.op.outputs, new_tileables):
90                     self._optimizer_context[t] = new_t.data
91                     if t is n:
92                         new_node = new_t.data
93             else:
94                 new_node = n
95             new_graph.add_node(new_node)
96             for inp in new_node.inputs:
97                 new_graph.add_node(inp)
98                 new_graph.add_edge(inp, new_node)
99         return new_graph
100 
101     @kernel_mode
102     @enter_build_mode
103     def build(self, tileables, tileable_graph=None):
104         self._optimizer_context.append_result_tileables(tileables)
105         graph = super().build(tileables, tileable_graph=tileable_graph)
106         graph = self._replace_copied_tilebale(graph)
107         self._mapping_tileables(tileables)
108         return graph
109 
110 
111 def register(op_type, rule):
112     _rules[op_type].append(rule)
113
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/mars/dataframe/base/standardize_range_index.py b/mars/dataframe/base/standardize_range_index.py
--- a/mars/dataframe/base/standardize_range_index.py
+++ b/mars/dataframe/base/standardize_range_index.py
@@ -31,7 +31,7 @@
     _axis = Int32Field('axis')
 
     def __init__(self, prepare_inputs=None, axis=None, object_type=None, **kwargs):
-        super().__init__(__prepare_inputs=prepare_inputs, _axis=axis, _object_type=object_type, **kwargs)
+        super().__init__(_prepare_inputs=prepare_inputs, _axis=axis, _object_type=object_type, **kwargs)
 
     @property
     def axis(self):
diff --git a/mars/optimizes/tileable_graph/column_pruning.py b/mars/optimizes/tileable_graph/column_pruning.py
--- a/mars/optimizes/tileable_graph/column_pruning.py
+++ b/mars/optimizes/tileable_graph/column_pruning.py
@@ -62,7 +62,7 @@
 
         self._optimizer_context[node] = new_node
         self._optimizer_context[input_node] = new_input
-        return new_node
+        return node
 
 
 register(DataFrameGroupByAgg, GroupbyPruneReadCSV)
diff --git a/mars/optimizes/tileable_graph/core.py b/mars/optimizes/tileable_graph/core.py
--- a/mars/optimizes/tileable_graph/core.py
+++ b/mars/optimizes/tileable_graph/core.py
@@ -79,15 +79,19 @@
             return graph
 
         new_graph = DAG()
-        reversed_mapping = weakref.WeakKeyDictionary((v, k) for k, v in self._optimizer_context.items())
+        replaced_tileables = weakref.WeakKeyDictionary()
         for n in graph.topological_iter():
-            if n in reversed_mapping:
-                new_node = n
-            elif any(inp in self._optimizer_context for inp in n.inputs):
-                new_inputs = [self._optimizer_context.get(i, i) for i in n.inputs]
+            if graph.count_predecessors(n) == 0:
+                if n in self._optimizer_context and \
+                        all(suc in self._optimizer_context for suc in graph.successors(n)):
+                    replaced_tileables[n] = new_node = self._optimizer_context[n]
+                else:
+                    new_node = n
+            elif any(inp in replaced_tileables for inp in n.inputs):
+                new_inputs = [replaced_tileables.get(i, i) for i in n.inputs]
                 new_tileables = copy_tileables(n.op.outputs, inputs=new_inputs)
                 for t, new_t in zip(n.op.outputs, new_tileables):
-                    self._optimizer_context[t] = new_t.data
+                    replaced_tileables[t] = new_t.data
                     if t is n:
                         new_node = new_t.data
             else:
@@ -96,6 +100,7 @@
             for inp in new_node.inputs:
                 new_graph.add_node(inp)
                 new_graph.add_edge(inp, new_node)
+        self._optimizer_context.update(replaced_tileables)
         return new_graph
 
     @kernel_mode

{"golden_diff": "diff --git a/mars/dataframe/base/standardize_range_index.py b/mars/dataframe/base/standardize_range_index.py\n--- a/mars/dataframe/base/standardize_range_index.py\n+++ b/mars/dataframe/base/standardize_range_index.py\n@@ -31,7 +31,7 @@\n     _axis = Int32Field('axis')\n \n     def __init__(self, prepare_inputs=None, axis=None, object_type=None, **kwargs):\n-        super().__init__(__prepare_inputs=prepare_inputs, _axis=axis, _object_type=object_type, **kwargs)\n+        super().__init__(_prepare_inputs=prepare_inputs, _axis=axis, _object_type=object_type, **kwargs)\n \n     @property\n     def axis(self):\ndiff --git a/mars/optimizes/tileable_graph/column_pruning.py b/mars/optimizes/tileable_graph/column_pruning.py\n--- a/mars/optimizes/tileable_graph/column_pruning.py\n+++ b/mars/optimizes/tileable_graph/column_pruning.py\n@@ -62,7 +62,7 @@\n \n         self._optimizer_context[node] = new_node\n         self._optimizer_context[input_node] = new_input\n-        return new_node\n+        return node\n \n \n register(DataFrameGroupByAgg, GroupbyPruneReadCSV)\ndiff --git a/mars/optimizes/tileable_graph/core.py b/mars/optimizes/tileable_graph/core.py\n--- a/mars/optimizes/tileable_graph/core.py\n+++ b/mars/optimizes/tileable_graph/core.py\n@@ -79,15 +79,19 @@\n             return graph\n \n         new_graph = DAG()\n-        reversed_mapping = weakref.WeakKeyDictionary((v, k) for k, v in self._optimizer_context.items())\n+        replaced_tileables = weakref.WeakKeyDictionary()\n         for n in graph.topological_iter():\n-            if n in reversed_mapping:\n-                new_node = n\n-            elif any(inp in self._optimizer_context for inp in n.inputs):\n-                new_inputs = [self._optimizer_context.get(i, i) for i in n.inputs]\n+            if graph.count_predecessors(n) == 0:\n+                if n in self._optimizer_context and \\\n+                        all(suc in self._optimizer_context for suc in graph.successors(n)):\n+                    replaced_tileables[n] = new_node = self._optimizer_context[n]\n+                else:\n+                    new_node = n\n+            elif any(inp in replaced_tileables for inp in n.inputs):\n+                new_inputs = [replaced_tileables.get(i, i) for i in n.inputs]\n                 new_tileables = copy_tileables(n.op.outputs, inputs=new_inputs)\n                 for t, new_t in zip(n.op.outputs, new_tileables):\n-                    self._optimizer_context[t] = new_t.data\n+                    replaced_tileables[t] = new_t.data\n                     if t is n:\n                         new_node = new_t.data\n             else:\n@@ -96,6 +100,7 @@\n             for inp in new_node.inputs:\n                 new_graph.add_node(inp)\n                 new_graph.add_edge(inp, new_node)\n+        self._optimizer_context.update(replaced_tileables)\n         return new_graph\n \n     @kernel_mode\n", "issue": "[BUG] Column pruning leads to a wrong result\n<!--\r\nThank you for your contribution!\r\n\r\nPlease review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.\r\n-->\r\n\r\n**Describe the bug**\r\nA clear and concise description of what the bug is.\r\n\r\n**To Reproduce**\r\n``` Python\r\npd_df = pd.DataFrame({'a': [3, 4, 5, 3, 5, 4, 1, 2, 3],\r\n                      'b': [1, 3, 4, 5, 6, 5, 4, 4, 4],\r\n                      'c': list('aabaaddce'),\r\n                      'd': list('abaaaddce')})\r\npd_df.to_csv('test.csv', index=False)\r\nin_df = md.read_csv('test.csv')\r\ndf1 = in_df.groupby('d').agg({'b': 'min'})\r\ndf2 = in_df[in_df.d.isin(df1.index)]\r\ndf2.execute()\r\nOut[5]: \r\n   b  d\r\n0  1  a\r\n1  3  b\r\n2  4  a\r\n3  5  a\r\n4  6  a\r\n5  5  d\r\n6  4  d\r\n7  4  c\r\n8  4  e\r\n```\r\n\r\nthe output only has two columns which should be 4.\r\n\n", "before_files": [{"content": "#!/usr/bin/env python\n# -*- coding: utf-8 -*-\n# Copyright 1999-2020 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom ...utils import copy_tileables\nfrom ...dataframe.utils import parse_index\nfrom ...dataframe.groupby.aggregation import DataFrameGroupByAgg\nfrom ...dataframe.datasource.read_csv import DataFrameReadCSV\nfrom .core import TileableOptimizeRule, register\n\n\nclass GroupbyPruneReadCSV(TileableOptimizeRule):\n    \"\"\"\n    An experimental implementation for tileable optimization rule.\n    This rule works only when groupby aggregation operation follows the read CSV files,\n    we can prune the columns that not used by the following operations when read the files.\n    \"\"\"\n    def match(self, node):\n        if isinstance(node.inputs[0].op, DataFrameReadCSV) and \\\n                node.inputs[0] not in self._optimizer_context.result_tileables:\n            by_data = node.op.groupby_params.get('by')\n            by_columns = by_data if isinstance(by_data, (list, tuple)) else [by_data]\n            if isinstance(node.op.func, (str, list)):\n                # Passing func name(s) means perform on all columns.\n                return False\n            elif len(set(by_columns + list(node.op.func))) == len(node.inputs[0].op.usecols or node.inputs[0].dtypes):\n                # If performs on all columns, no need to prune.\n                return False\n            return True\n        return False\n\n    def apply(self, node):\n        by_data = node.op.groupby_params.get('by')\n        by_columns = by_data if isinstance(by_data, (list, tuple)) else [by_data]\n        agg_columns = list(node.op.func)\n        input_node = node.inputs[0]\n        selected_columns = [c for c in list(input_node.dtypes.index) if c in by_columns + agg_columns]\n        if input_node in self._optimizer_context:\n            new_input = self._optimizer_context[input_node]\n            selected_columns = [\n                c for c in list(input_node.dtypes.index) if c in selected_columns + new_input.op._usecols]\n        else:\n            new_input = copy_tileables([input_node])[0].data\n\n        new_input._shape = (input_node.shape[0], len(selected_columns))\n        new_input._dtypes = input_node.dtypes[selected_columns]\n        new_input._columns_value = parse_index(new_input._dtypes.index, store_data=True)\n        new_input.op._usecols = selected_columns\n        new_node = copy_tileables([node], inputs=[new_input])[0].data\n\n        self._optimizer_context[node] = new_node\n        self._optimizer_context[input_node] = new_input\n        return new_node\n\n\nregister(DataFrameGroupByAgg, GroupbyPruneReadCSV)\n", "path": "mars/optimizes/tileable_graph/column_pruning.py"}, {"content": "#!/usr/bin/env python\n# -*- coding: utf-8 -*-\n# Copyright 1999-2020 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport pandas as pd\n\nfrom ... import opcodes as OperandDef\nfrom ...utils import lazy_import\nfrom ...serialize import Int32Field\nfrom ..operands import DataFrameOperandMixin, DataFrameOperand\n\n\ncudf = lazy_import('cudf', globals=globals())\n\n\nclass ChunkStandardizeRangeIndex(DataFrameOperand, DataFrameOperandMixin):\n    _op_type_ = OperandDef.STANDARDIZE_RANGE_INDEX\n\n    _axis = Int32Field('axis')\n\n    def __init__(self, prepare_inputs=None, axis=None, object_type=None, **kwargs):\n        super().__init__(__prepare_inputs=prepare_inputs, _axis=axis, _object_type=object_type, **kwargs)\n\n    @property\n    def axis(self):\n        return self._axis\n\n    @classmethod\n    def execute(cls, ctx, op):\n        xdf = cudf if op.gpu else pd\n        in_data = ctx[op.inputs[-1].key].copy()\n        input_keys = [c.key for c in op.inputs[:-1]]\n        metas = ctx.get_chunk_metas(input_keys)\n        index_start = sum([m.chunk_shape[op.axis] for m in metas])\n        if op.axis == 0:\n            in_data.index = xdf.RangeIndex(index_start, index_start + len(in_data))\n        else:\n            in_data.columns = xdf.RangeIndex(index_start, index_start + in_data.shape[1])\n        ctx[op.outputs[0].key] = in_data\n", "path": "mars/dataframe/base/standardize_range_index.py"}, {"content": "#!/usr/bin/env python\n# -*- coding: utf-8 -*-\n# Copyright 1999-2020 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport weakref\nfrom collections import defaultdict\n\nfrom ...graph import DAG\nfrom ...graph_builder import TileableGraphBuilder\nfrom ...utils import copy_tileables, kernel_mode, enter_build_mode\n\n_rules = defaultdict(list)\n\ntileable_optimized = weakref.WeakKeyDictionary()\n\n\nclass TileableOptimizeRule(object):\n    def __init__(self, optimized_context):\n        self._optimizer_context = optimized_context\n\n    def match(self, node):\n        raise NotImplementedError\n\n    def apply(self, node):\n        raise NotImplementedError\n\n\nclass OptimizeContext(weakref.WeakKeyDictionary):\n    def __init__(self, dict=None):\n        weakref.WeakKeyDictionary.__init__(self, dict=dict)\n        self._result_tileables = []\n\n    @property\n    def result_tileables(self):\n        return self._result_tileables\n\n    def append_result_tileables(self, tileables):\n        self._result_tileables.extend(tileables)\n\n\nclass OptimizeIntegratedTileableGraphBuilder(TileableGraphBuilder):\n    def __init__(self, **kw):\n        self._optimizer_context = OptimizeContext()\n        super().__init__(**kw)\n        self._node_processor = self._apply_rules(self._node_processor, self._optimizer_context)\n\n    @staticmethod\n    def _apply_rules(node_processor, optimizer_context):\n        def inner(node):\n            node = node_processor(node) if node_processor is not None else node\n            if type(node.op) in _rules:\n                for rule in _rules[type(node.op)]:\n                    ruler = rule(optimizer_context)\n                    if ruler.match(node):\n                        node = rule(optimizer_context).apply(node)\n            return node\n\n        return inner\n\n    def _mapping_tileables(self, tileables):\n        for t in tileables:\n            if t in self._optimizer_context:\n                tileable_optimized[t] = self._optimizer_context[t]\n\n    def _replace_copied_tilebale(self, graph):\n        if len(self._optimizer_context) == 0:\n            return graph\n\n        new_graph = DAG()\n        reversed_mapping = weakref.WeakKeyDictionary((v, k) for k, v in self._optimizer_context.items())\n        for n in graph.topological_iter():\n            if n in reversed_mapping:\n                new_node = n\n            elif any(inp in self._optimizer_context for inp in n.inputs):\n                new_inputs = [self._optimizer_context.get(i, i) for i in n.inputs]\n                new_tileables = copy_tileables(n.op.outputs, inputs=new_inputs)\n                for t, new_t in zip(n.op.outputs, new_tileables):\n                    self._optimizer_context[t] = new_t.data\n                    if t is n:\n                        new_node = new_t.data\n            else:\n                new_node = n\n            new_graph.add_node(new_node)\n            for inp in new_node.inputs:\n                new_graph.add_node(inp)\n                new_graph.add_edge(inp, new_node)\n        return new_graph\n\n    @kernel_mode\n    @enter_build_mode\n    def build(self, tileables, tileable_graph=None):\n        self._optimizer_context.append_result_tileables(tileables)\n        graph = super().build(tileables, tileable_graph=tileable_graph)\n        graph = self._replace_copied_tilebale(graph)\n        self._mapping_tileables(tileables)\n        return graph\n\n\ndef register(op_type, rule):\n    _rules[op_type].append(rule)\n", "path": "mars/optimizes/tileable_graph/core.py"}], "after_files": [{"content": "#!/usr/bin/env python\n# -*- coding: utf-8 -*-\n# Copyright 1999-2020 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom ...utils import copy_tileables\nfrom ...dataframe.utils import parse_index\nfrom ...dataframe.groupby.aggregation import DataFrameGroupByAgg\nfrom ...dataframe.datasource.read_csv import DataFrameReadCSV\nfrom .core import TileableOptimizeRule, register\n\n\nclass GroupbyPruneReadCSV(TileableOptimizeRule):\n    \"\"\"\n    An experimental implementation for tileable optimization rule.\n    This rule works only when groupby aggregation operation follows the read CSV files,\n    we can prune the columns that not used by the following operations when read the files.\n    \"\"\"\n    def match(self, node):\n        if isinstance(node.inputs[0].op, DataFrameReadCSV) and \\\n                node.inputs[0] not in self._optimizer_context.result_tileables:\n            by_data = node.op.groupby_params.get('by')\n            by_columns = by_data if isinstance(by_data, (list, tuple)) else [by_data]\n            if isinstance(node.op.func, (str, list)):\n                # Passing func name(s) means perform on all columns.\n                return False\n            elif len(set(by_columns + list(node.op.func))) == len(node.inputs[0].op.usecols or node.inputs[0].dtypes):\n                # If performs on all columns, no need to prune.\n                return False\n            return True\n        return False\n\n    def apply(self, node):\n        by_data = node.op.groupby_params.get('by')\n        by_columns = by_data if isinstance(by_data, (list, tuple)) else [by_data]\n        agg_columns = list(node.op.func)\n        input_node = node.inputs[0]\n        selected_columns = [c for c in list(input_node.dtypes.index) if c in by_columns + agg_columns]\n        if input_node in self._optimizer_context:\n            new_input = self._optimizer_context[input_node]\n            selected_columns = [\n                c for c in list(input_node.dtypes.index) if c in selected_columns + new_input.op._usecols]\n        else:\n            new_input = copy_tileables([input_node])[0].data\n\n        new_input._shape = (input_node.shape[0], len(selected_columns))\n        new_input._dtypes = input_node.dtypes[selected_columns]\n        new_input._columns_value = parse_index(new_input._dtypes.index, store_data=True)\n        new_input.op._usecols = selected_columns\n        new_node = copy_tileables([node], inputs=[new_input])[0].data\n\n        self._optimizer_context[node] = new_node\n        self._optimizer_context[input_node] = new_input\n        return node\n\n\nregister(DataFrameGroupByAgg, GroupbyPruneReadCSV)\n", "path": "mars/optimizes/tileable_graph/column_pruning.py"}, {"content": "#!/usr/bin/env python\n# -*- coding: utf-8 -*-\n# Copyright 1999-2020 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport pandas as pd\n\nfrom ... import opcodes as OperandDef\nfrom ...utils import lazy_import\nfrom ...serialize import Int32Field\nfrom ..operands import DataFrameOperandMixin, DataFrameOperand\n\n\ncudf = lazy_import('cudf', globals=globals())\n\n\nclass ChunkStandardizeRangeIndex(DataFrameOperand, DataFrameOperandMixin):\n    _op_type_ = OperandDef.STANDARDIZE_RANGE_INDEX\n\n    _axis = Int32Field('axis')\n\n    def __init__(self, prepare_inputs=None, axis=None, object_type=None, **kwargs):\n        super().__init__(_prepare_inputs=prepare_inputs, _axis=axis, _object_type=object_type, **kwargs)\n\n    @property\n    def axis(self):\n        return self._axis\n\n    @classmethod\n    def execute(cls, ctx, op):\n        xdf = cudf if op.gpu else pd\n        in_data = ctx[op.inputs[-1].key].copy()\n        input_keys = [c.key for c in op.inputs[:-1]]\n        metas = ctx.get_chunk_metas(input_keys)\n        index_start = sum([m.chunk_shape[op.axis] for m in metas])\n        if op.axis == 0:\n            in_data.index = xdf.RangeIndex(index_start, index_start + len(in_data))\n        else:\n            in_data.columns = xdf.RangeIndex(index_start, index_start + in_data.shape[1])\n        ctx[op.outputs[0].key] = in_data\n", "path": "mars/dataframe/base/standardize_range_index.py"}, {"content": "#!/usr/bin/env python\n# -*- coding: utf-8 -*-\n# Copyright 1999-2020 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport weakref\nfrom collections import defaultdict\n\nfrom ...graph import DAG\nfrom ...graph_builder import TileableGraphBuilder\nfrom ...utils import copy_tileables, kernel_mode, enter_build_mode\n\n_rules = defaultdict(list)\n\ntileable_optimized = weakref.WeakKeyDictionary()\n\n\nclass TileableOptimizeRule(object):\n    def __init__(self, optimized_context):\n        self._optimizer_context = optimized_context\n\n    def match(self, node):\n        raise NotImplementedError\n\n    def apply(self, node):\n        raise NotImplementedError\n\n\nclass OptimizeContext(weakref.WeakKeyDictionary):\n    def __init__(self, dict=None):\n        weakref.WeakKeyDictionary.__init__(self, dict=dict)\n        self._result_tileables = []\n\n    @property\n    def result_tileables(self):\n        return self._result_tileables\n\n    def append_result_tileables(self, tileables):\n        self._result_tileables.extend(tileables)\n\n\nclass OptimizeIntegratedTileableGraphBuilder(TileableGraphBuilder):\n    def __init__(self, **kw):\n        self._optimizer_context = OptimizeContext()\n        super().__init__(**kw)\n        self._node_processor = self._apply_rules(self._node_processor, self._optimizer_context)\n\n    @staticmethod\n    def _apply_rules(node_processor, optimizer_context):\n        def inner(node):\n            node = node_processor(node) if node_processor is not None else node\n            if type(node.op) in _rules:\n                for rule in _rules[type(node.op)]:\n                    ruler = rule(optimizer_context)\n                    if ruler.match(node):\n                        node = rule(optimizer_context).apply(node)\n            return node\n\n        return inner\n\n    def _mapping_tileables(self, tileables):\n        for t in tileables:\n            if t in self._optimizer_context:\n                tileable_optimized[t] = self._optimizer_context[t]\n\n    def _replace_copied_tilebale(self, graph):\n        if len(self._optimizer_context) == 0:\n            return graph\n\n        new_graph = DAG()\n        replaced_tileables = weakref.WeakKeyDictionary()\n        for n in graph.topological_iter():\n            if graph.count_predecessors(n) == 0:\n                if n in self._optimizer_context and \\\n                        all(suc in self._optimizer_context for suc in graph.successors(n)):\n                    replaced_tileables[n] = new_node = self._optimizer_context[n]\n                else:\n                    new_node = n\n            elif any(inp in replaced_tileables for inp in n.inputs):\n                new_inputs = [replaced_tileables.get(i, i) for i in n.inputs]\n                new_tileables = copy_tileables(n.op.outputs, inputs=new_inputs)\n                for t, new_t in zip(n.op.outputs, new_tileables):\n                    replaced_tileables[t] = new_t.data\n                    if t is n:\n                        new_node = new_t.data\n            else:\n                new_node = n\n            new_graph.add_node(new_node)\n            for inp in new_node.inputs:\n                new_graph.add_node(inp)\n                new_graph.add_edge(inp, new_node)\n        self._optimizer_context.update(replaced_tileables)\n        return new_graph\n\n    @kernel_mode\n    @enter_build_mode\n    def build(self, tileables, tileable_graph=None):\n        self._optimizer_context.append_result_tileables(tileables)\n        graph = super().build(tileables, tileable_graph=tileable_graph)\n        graph = self._replace_copied_tilebale(graph)\n        self._mapping_tileables(tileables)\n        return graph\n\n\ndef register(op_type, rule):\n    _rules[op_type].append(rule)\n", "path": "mars/optimizes/tileable_graph/core.py"}]}