Repository: mars-project/mars

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.



Discussion:

Resolving patch:
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 @@ class ChunkStandardizeRangeIndex(DataFrameOperand, DataFrameOperandMixin):
     _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 @@ def apply(self, node):
 
         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 @@ def _replace_copied_tilebale(self, graph):
             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 @@ def _replace_copied_tilebale(self, graph):
             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


Test patch:
diff --git a/mars/optimizes/tileable_graph/tests/test_column_pruning.py b/mars/optimizes/tileable_graph/tests/test_column_pruning.py
--- a/mars/optimizes/tileable_graph/tests/test_column_pruning.py
+++ b/mars/optimizes/tileable_graph/tests/test_column_pruning.py
@@ -23,11 +23,14 @@
 import mars.dataframe as md
 from mars.core import ExecutableTuple
 from mars.config import option_context
-from mars.tests.core import TestBase
+from mars.tests.core import TestBase, ExecutorForTest
 from mars.optimizes.tileable_graph.core import tileable_optimized
 
 
 class Test(TestBase):
+    def setUp(self):
+        self.executor = ExecutorForTest()
+
     def testGroupByPruneReadCSV(self):
         tempdir = tempfile.mkdtemp()
         file_path = os.path.join(tempdir, 'test.csv')
@@ -38,6 +41,12 @@ def testGroupByPruneReadCSV(self):
                                'd': list('abaaaddce')})
             df.to_csv(file_path, index=False)
 
+            # Use test executor
+            mdf = md.read_csv(file_path).groupby('c').agg({'a': 'sum'})
+            result = self.executor.execute_dataframe(mdf)[0]
+            expected = df.groupby('c').agg({'a': 'sum'})
+            pd.testing.assert_frame_equal(result, expected)
+
             mdf = md.read_csv(file_path).groupby('c').agg({'a': 'sum'})
             expected = df.groupby('c').agg({'a': 'sum'})
             pd.testing.assert_frame_equal(mdf.to_pandas(), expected)
@@ -97,3 +106,37 @@ def testGroupByPruneReadCSV(self):
 
         finally:
             shutil.rmtree(tempdir)
+
+    def testExecutedPruning(self):
+        tempdir = tempfile.mkdtemp()
+        file_path = os.path.join(tempdir, 'test.csv')
+        try:
+            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(file_path, index=False)
+
+            in_df = md.read_csv(file_path)
+            mdf = in_df.groupby('c').agg({'a': 'sum'})
+
+            expected = pd_df.groupby('c').agg({'a': 'sum'})
+            pd.testing.assert_frame_equal(mdf.to_pandas(), expected)
+            optimized_df = tileable_optimized[mdf.data]
+            self.assertEqual(optimized_df.inputs[0].op.usecols, ['a', 'c'])
+
+            # make sure in_df has correct columns
+            pd.testing.assert_frame_equal(in_df.to_pandas(), pd_df)
+
+            # skip pruning
+            in_df = md.read_csv(file_path)
+            df1 = in_df.groupby('d').agg({'b': 'min'})
+            df2 = in_df[in_df.d.isin(df1.index)]
+
+            expected1 = pd_df.groupby('d').agg({'b': 'min'})
+            expected2 = pd_df[pd_df.d.isin(expected1.index)]
+
+            pd.testing.assert_frame_equal(df2.to_pandas(), expected2)
+
+        finally:
+            shutil.rmtree(tempdir)