Repository: mars-project/mars

Issue:
[BUG] df.loc[:, ['col']] led to an incorrect result
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->


**To Reproduce**
``` Python
In [49]: mdf = md.DataFrame({'A' : ['foo', 'bar', 'foo', 'bar',
    ...:                           'foo', 'bar'],
    ...:                    'B' : ['one', 'one', 'two', 'three',
    ...:                           'two', 'two'],
    ...:                    'C' : [1, 5, 5, 2, None, 5],
    ...:                    'D' : [2.0, 5., 8., 1., 2., None]}, chunk_size=5, index=[0, 1, 2, 0, 1,
    ...: 2]
    ...:                    )

In [50]: selected1 = mdf.loc[:, ['A', 'C']]

In [51]: selected2 = mdf[['A', 'C']]

In [52]: selected1[mdf['D'] > 2].execute()
100%|██████████████████████████████████████████████████████████| 100.0/100 [00:00<00:00, 362.46it/s]
Out[52]: 
     A    C
1  bar  5.0
2  foo  5.0
1  bar  5.0

In [53]: selected2[mdf['D'] > 2].execute()
100%|██████████████████████████████████████████████████████████| 100.0/100 [00:00<00:00, 589.57it/s]
Out[53]: 
     A    C
1  bar  5.0
2  foo  5.0
```
`selected1` is equivalent to `selected2` while they produced different results, `selected2` should be the correct answer.


Discussion:

Resolving patch:
diff --git a/mars/dataframe/indexing/index_lib.py b/mars/dataframe/indexing/index_lib.py
--- a/mars/dataframe/indexing/index_lib.py
+++ b/mars/dataframe/indexing/index_lib.py
@@ -44,7 +44,7 @@
     to_numpy,
     normalize_chunk_sizes,
 )
-from ...utils import classproperty, has_unknown_shape
+from ...utils import classproperty, has_unknown_shape, is_full_slice
 from ..core import SERIES_CHUNK_TYPE, SERIES_TYPE, IndexValue
 from ..utils import parse_index
 from .utils import convert_labels_into_positions
@@ -190,10 +190,13 @@ def set_chunk_index_info(
             "dtypes": None,
         }
         if index_info.input_axis == 0:
-            index = chunk_input.index_value.to_pandas()
-            kw["index_value"] = parse_index(
-                index[slc], chunk_input, slc, store_data=False
-            )
+            if is_full_slice(slc):
+                kw["index_value"] = chunk_input.index_value
+            else:
+                index = chunk_input.index_value.to_pandas()
+                kw["index_value"] = parse_index(
+                    index[slc], chunk_input, slc, store_data=False
+                )
         else:
             assert index_info.input_axis == 1
             index = chunk_input.columns_value.to_pandas()
@@ -272,12 +275,17 @@ def set_chunk_index_info(
             "dtypes": None,
         }
         if index_info.input_axis == 0:
-            index = chunk_input.index_value.to_pandas()
-            start, stop = index.slice_locs(slc.start, slc.stop, slc.step, kind="loc")
-            pos_slc = slice(start, stop, slc.step)
-            kw["index_value"] = parse_index(
-                index[pos_slc], chunk_input, slc, store_data=False
-            )
+            if is_full_slice(index):
+                kw["index_value"] = chunk_input.index_value
+            else:
+                index = chunk_input.index_value.to_pandas()
+                start, stop = index.slice_locs(
+                    slc.start, slc.stop, slc.step, kind="loc"
+                )
+                pos_slc = slice(start, stop, slc.step)
+                kw["index_value"] = parse_index(
+                    index[pos_slc], chunk_input, slc, store_data=False
+                )
         else:
             assert index_info.input_axis == 1
             dtypes = chunk_input.dtypes


Test patch:
diff --git a/mars/dataframe/indexing/tests/test_indexing.py b/mars/dataframe/indexing/tests/test_indexing.py
--- a/mars/dataframe/indexing/tests/test_indexing.py
+++ b/mars/dataframe/indexing/tests/test_indexing.py
@@ -430,6 +430,18 @@ def test_iloc_setitem():
     ]
     assert series.chunks[1].op.value == 2
 
+    raw = pd.DataFrame(
+        np.random.rand(9, 2),
+        index=["a1", "a2", "a3"] * 3,
+        columns=["x", "y"],
+    )
+    df = md.DataFrame(raw, chunk_size=4)
+    iloc_df = df.iloc[:, 1:]
+    tiled_df, tiled_iloc_df = tile(df, iloc_df)
+    # for full slice, index_value should be same as input chunk
+    for loc_chunk, chunk in zip(tiled_iloc_df.chunks, tiled_df.chunks):
+        assert loc_chunk.index_value.key == chunk.index_value.key
+
     # fancy index
     series = md.Series(pd.Series(np.arange(10)), chunk_size=3)
     series.iloc[[2, 4, 9]] = 3
@@ -620,6 +632,19 @@ def test_dataframe_loc():
     with pytest.raises(KeyError):
         _ = df.loc[:, ["non_exist"]]
 
+    # test loc chunk's index_value
+    raw = pd.DataFrame(
+        np.random.rand(9, 2),
+        index=["a1", "a2", "a3"] * 3,
+        columns=["x", "y"],
+    )
+    df = md.DataFrame(raw, chunk_size=4)
+    loc_df = df.loc[:, ["x"]]
+    tiled_df, tiled_loc_df = tile(df, loc_df)
+    # for full slice, index_value should be same as input chunk
+    for loc_chunk, chunk in zip(tiled_loc_df.chunks, tiled_df.chunks):
+        assert loc_chunk.index_value.key == chunk.index_value.key
+
 
 def test_loc_use_iloc():
     raw = pd.DataFrame([[1, 3, 3], [4, 2, 6], [7, 8, 9]], columns=["x", "y", "z"])