Repository: mars-project/mars

Issue:
[BUG] DataFrame getitem raised `TypeError`
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
Reproduce code:
``` Python
In [9]: df = md.DataFrame({'col': [1,2,3]})                                                         

In [10]: df[['col', 'col']]['col'].execute()                                                        
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~/Documents/mars_dev/mars/mars/serialization/serializables/field.py in __set__(self, instance, value)
    118                     to_check_value = self._on_serialize(to_check_value)
--> 119                 field_type.validate(to_check_value)
    120             except (TypeError, ValueError) as e:

~/Documents/mars_dev/mars/mars/serialization/serializables/field_type.py in validate(self, value)
    107             raise TypeError(
--> 108                 f"value needs to be instance "
    109                 f"of {self.valid_types}, got {type(value)}"

TypeError: value needs to be instance of (<class 'numpy.dtype'>, <class 'pandas.core.dtypes.base.ExtensionDtype'>), got <class 'pandas.core.series.Series'>

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
<ipython-input-10-16733874ef9c> in <module>
----> 1 df[['col', 'col']]['col'].execute()

~/Documents/mars_dev/mars/mars/dataframe/indexing/getitem.py in dataframe_getitem(df, item)
    564             raise KeyError(f"{item} not in columns")
    565         op = DataFrameIndex(col_names=item)
--> 566     return op(df)
    567 
    568 

~/Documents/mars_dev/mars/mars/core/mode.py in _inner(*args, **kwargs)
     75             def _inner(*args, **kwargs):
     76                 with self:
---> 77                     return func(*args, **kwargs)
     78 
     79         else:

~/Documents/mars_dev/mars/mars/dataframe/indexing/getitem.py in __call__(self, df)
    305                     dtype=dtype,
    306                     index_value=df.index_value,
--> 307                     name=self._col_names,
    308                 )
    309         else:

~/Documents/mars_dev/mars/mars/dataframe/operands.py in new_series(self, inputs, shape, dtype, index_value, name, **kw)
    129 
    130         return self.new_seriess(
--> 131             inputs, shape=shape, dtype=dtype, index_value=index_value, name=name, **kw
    132         )[0]
    133 

~/Documents/mars_dev/mars/mars/dataframe/operands.py in new_seriess(self, inputs, shape, dtype, index_value, name, chunks, nsplits, output_limit, kws, **kw)
    119             output_limit=output_limit,
    120             kws=kws,
--> 121             **kw
    122         )
    123 

~/Documents/mars_dev/mars/mars/core/operand/core.py in new_tileables(self, inputs, kws, **kw)
    263             Override the method `_new_tileables` if needed.
    264         """
--> 265         tileables = self._new_tileables(inputs, kws=kws, **kw)
    266         if is_eager_mode():
    267             ExecutableTuple(tileables).execute()

~/Documents/mars_dev/mars/mars/core/operand/core.py in _new_tileables(self, inputs, kws, **kw)
    222             if kws:
    223                 create_tensor_kw.update(kws[j])
--> 224             tileable = self._create_tileable(j, **create_tensor_kw)
    225             tileables.append(tileable)
    226 

~/Documents/mars_dev/mars/mars/core/operand/core.py in _create_tileable(self, output_idx, **kw)
    200 
    201         kw = self._fill_nan_shape(kw)
--> 202         data = tileable_data_type(**kw)
    203         return tileable_type(data)
    204 

~/Documents/mars_dev/mars/mars/dataframe/core.py in __init__(self, op, shape, nsplits, dtype, name, index_value, chunks, **kw)
   1080             _index_value=index_value,
   1081             _chunks=chunks,
-> 1082             **kw,
   1083         )
   1084         self._accessors = dict()

~/Documents/mars_dev/mars/mars/core/entity/tileables.py in __init__(self, *args, **kwargs)
    279             kwargs["_nsplits"] = tuple(tuple(s) for s in kwargs["_nsplits"])
    280 
--> 281         super().__init__(*args, **kwargs)
    282 
    283         if hasattr(self, "_chunks") and self._chunks:

~/Documents/mars_dev/mars/mars/core/entity/core.py in __init__(self, *args, **kwargs)
     37         )
     38         kwargs["_extra_params"] = kwargs.pop("_extra_params", extras)
---> 39         super().__init__(*args, **kwargs)
     40 
     41     @property

~/Documents/mars_dev/mars/mars/core/base.py in __init__(self, *args, **kwargs)
     29 
     30     def __init__(self, *args, **kwargs):
---> 31         super().__init__(*args, **kwargs)
     32 
     33         if self._init_update_key_ and (not hasattr(self, "_key") or not self._key):

~/Documents/mars_dev/mars/mars/serialization/serializables/core.py in __init__(self, *args, **kwargs)
     62 
     63         for key, val in kwargs.items():
---> 64             object.__setattr__(self, key, val)
     65 
     66     def __repr__(self):

~/Documents/mars_dev/mars/mars/serialization/serializables/field.py in __set__(self, instance, value)
    119                 field_type.validate(to_check_value)
    120             except (TypeError, ValueError) as e:
--> 121                 raise type(e)(f"Failed to set `{self._attr_name}`: {str(e)}")
    122         instance._FIELD_VALUES[self._tag] = value
    123 

TypeError: Failed to set `_dtype`: value needs to be instance of (<class 'numpy.dtype'>, <class 'pandas.core.dtypes.base.ExtensionDtype'>), got <class 'pandas.core.series.Series'>
```


Discussion:

Resolving patch:
diff --git a/mars/dataframe/indexing/getitem.py b/mars/dataframe/indexing/getitem.py
--- a/mars/dataframe/indexing/getitem.py
+++ b/mars/dataframe/indexing/getitem.py
@@ -287,18 +287,17 @@ def _set_inputs(self, inputs):
     def __call__(self, df):
         if self.col_names is not None:
             # if col_names is a list, return a DataFrame, else return a Series
-            if isinstance(self._col_names, list):
-                dtypes = df.dtypes[self._col_names]
-                columns = parse_index(pd.Index(self._col_names), store_data=True)
+            dtype = df.dtypes[self._col_names]
+            if isinstance(dtype, pd.Series):
+                columns = parse_index(dtype.index, store_data=True)
                 return self.new_dataframe(
                     [df],
-                    shape=(df.shape[0], len(self._col_names)),
-                    dtypes=dtypes,
+                    shape=(df.shape[0], len(dtype)),
+                    dtypes=dtype,
                     index_value=df.index_value,
                     columns_value=columns,
                 )
             else:
-                dtype = df.dtypes[self._col_names]
                 return self.new_series(
                     [df],
                     shape=(df.shape[0],),
@@ -439,8 +438,8 @@ def tile_with_columns(cls, op):
         in_df = op.inputs[0]
         out_df = op.outputs[0]
         col_names = op.col_names
-        if not isinstance(col_names, list):
-            column_index = calc_columns_index(col_names, in_df)
+        if not isinstance(out_df, DATAFRAME_TYPE):
+            column_index = calc_columns_index(col_names, in_df)[0]
             out_chunks = []
             dtype = in_df.dtypes[col_names]
             for i in range(in_df.chunk_shape[0]):
@@ -471,7 +470,10 @@ def tile_with_columns(cls, op):
             # When chunk columns are ['c1', 'c2', 'c3'], ['c4', 'c5'],
             # selected columns are ['c2', 'c3', 'c4', 'c2'], `column_splits` will be
             # [(['c2', 'c3'], 0), ('c4', 1), ('c2', 0)].
+            if not isinstance(col_names, _list_like_types):
+                col_names = [col_names]
             selected_index = [calc_columns_index(col, in_df) for col in col_names]
+            selected_index = list(itertools.chain.from_iterable(selected_index))
             condition = np.where(np.diff(selected_index))[0] + 1
             column_splits = np.split(col_names, condition)
             column_indexes = np.split(selected_index, condition)
@@ -482,7 +484,7 @@ def tile_with_columns(cls, op):
                 zip(column_splits, column_indexes)
             ):
                 dtypes = in_df.dtypes[columns]
-                column_nsplits.append(len(columns))
+                column_nsplits.append(len(dtypes))
                 for j in range(in_df.chunk_shape[0]):
                     c = in_df.cix[(j, column_idx[0])]
                     index_op = DataFrameIndex(
@@ -490,11 +492,13 @@ def tile_with_columns(cls, op):
                     )
                     out_chunk = index_op.new_chunk(
                         [c],
-                        shape=(c.shape[0], len(columns)),
+                        shape=(c.shape[0], len(dtypes)),
                         index=(j, i),
                         dtypes=dtypes,
                         index_value=c.index_value,
-                        columns_value=parse_index(pd.Index(columns), store_data=True),
+                        columns_value=parse_index(
+                            pd.Index(dtypes.index), store_data=True
+                        ),
                     )
                     out_chunks[j].append(out_chunk)
             out_chunks = [item for cl in out_chunks for item in cl]
@@ -525,8 +529,6 @@ def execute(cls, ctx, op):
                 mask = op.mask
             if hasattr(mask, "reindex_like"):
                 mask = mask.reindex_like(df).fillna(False)
-            if mask.ndim == 2:
-                mask = mask[df.columns.tolist()]
             ctx[op.outputs[0].key] = df[mask]
 
     @classmethod
diff --git a/mars/dataframe/indexing/utils.py b/mars/dataframe/indexing/utils.py
--- a/mars/dataframe/indexing/utils.py
+++ b/mars/dataframe/indexing/utils.py
@@ -24,8 +24,14 @@ def calc_columns_index(column_name, df):
     :return: chunk index on the columns axis
     """
     column_nsplits = df.nsplits[1]
-    column_loc = df.columns_value.to_pandas().get_loc(column_name)
-    return np.searchsorted(np.cumsum(column_nsplits), column_loc + 1)
+    # if has duplicate columns, will return multiple values
+    columns = df.columns_value.to_pandas().to_numpy()
+    column_locs = (columns == column_name).nonzero()[0]
+
+    return [
+        np.searchsorted(np.cumsum(column_nsplits), column_loc + 1)
+        for column_loc in column_locs
+    ]
 
 
 def convert_labels_into_positions(pandas_index, labels):


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
@@ -414,6 +414,9 @@ def test_dataframe_getitem(setup):
     df7 = df[1:7:2]
     pd.testing.assert_frame_equal(df7.execute().fetch(), data[1:7:2])
 
+    df8 = df[["c1", "c1"]]["c1"]
+    pd.testing.assert_frame_equal(df8.execute().fetch(), data[["c1", "c1"]]["c1"])
+
     series3 = df["c1"][0]
     assert series3.execute().fetch() == data["c1"][0]