Repository: mars-project/mars

Issue:
[BUG] Failed when `on` field is in `index` for `md.merge` 
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
Failed when `on` field is in `index` for `md.merge`.

**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 protobuf
4. Full stack of the error.
5. Minimized code to reproduce the error.

```
In [4]: df = pd.DataFrame({'a': np.arange(10), 'b': np.random.rand(10)})        

In [5]: df2 = df.copy()                                                         

In [6]: df2.set_index('a', inplace=True)                                        

In [7]: df2                                                                     
Out[7]: 
          b
a          
0  0.984265
1  0.544014
2  0.592392
3  0.269762
4  0.236130
5  0.846061
6  0.308780
7  0.604834
8  0.973824
9  0.867099

In [8]: df.merge(df2, on='a')  # can work for pandas                                                 
Out[8]: 
   a       b_x       b_y
0  0  0.984265  0.984265
1  1  0.544014  0.544014
2  2  0.592392  0.592392
3  3  0.269762  0.269762
4  4  0.236130  0.236130
5  5  0.846061  0.846061
6  6  0.308780  0.308780
7  7  0.604834  0.604834
8  8  0.973824  0.973824
9  9  0.867099  0.867099

In [9]: import mars.dataframe as md                                             

In [10]: mdf = md.DataFrame(df)                                                 

In [11]: mdf2 = md.DataFrame(df2)                                               

In [12]: mdf.merge(mdf2, on='a')  # cannot work for mars dataframe                                              
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-12-bd6a81883d3a> in <module>
----> 1 mdf.merge(mdf2, on='a')

~/Workspace/mars/mars/dataframe/merge/merge.py in merge(df, right, how, on, left_on, right_on, left_index, right_index, sort, suffixes, copy, indicator, strategy, validate)
    350         left_index=left_index, right_index=right_index, sort=sort, suffixes=suffixes,
    351         copy=copy, indicator=indicator, validate=validate, object_type=ObjectType.dataframe)
--> 352     return op(df, right)
    353 
    354 

~/Workspace/mars/mars/dataframe/merge/merge.py in __call__(self, left, right)
    174                                   left_index=self.left_index, right_index=self.right_index,
    175                                   sort=self.sort, suffixes=self.suffixes,
--> 176                                   copy=self.copy_, indicator=self.indicator, validate=self.validate)
    177 
    178         # the `index_value` doesn't matter.

~/miniconda3/lib/python3.7/site-packages/pandas/core/frame.py in merge(self, right, how, on, left_on, right_on, left_index, right_index, sort, suffixes, copy, indicator, validate)
   7292             copy=copy,
   7293             indicator=indicator,
-> 7294             validate=validate,
   7295         )
   7296 

~/miniconda3/lib/python3.7/site-packages/pandas/core/reshape/merge.py in merge(left, right, how, on, left_on, right_on, left_index, right_index, sort, suffixes, copy, indicator, validate)
     84         copy=copy,
     85         indicator=indicator,
---> 86         validate=validate,
     87     )
     88     return op.get_result()

~/miniconda3/lib/python3.7/site-packages/pandas/core/reshape/merge.py in __init__(self, left, right, how, on, left_on, right_on, axis, left_index, right_index, sort, suffixes, copy, indicator, validate)
    625             self.right_join_keys,
    626             self.join_names,
--> 627         ) = self._get_merge_keys()
    628 
    629         # validate the merge keys dtypes. We may need to coerce

~/miniconda3/lib/python3.7/site-packages/pandas/core/reshape/merge.py in _get_merge_keys(self)
    981                     if not is_rkey(rk):
    982                         if rk is not None:
--> 983                             right_keys.append(right._get_label_or_level_values(rk))
    984                         else:
    985                             # work-around for merge_asof(right_index=True)

~/miniconda3/lib/python3.7/site-packages/pandas/core/generic.py in _get_label_or_level_values(self, key, axis)
   1689             values = self.axes[axis].get_level_values(key)._values
   1690         else:
-> 1691             raise KeyError(key)
   1692 
   1693         # Check for duplicates

KeyError: 'a'
```

**Expected behavior**
A clear and concise description of what you expected to happen.

**Additional context**
Add any other context about the problem here.



Discussion:

Resolving patch:
diff --git a/mars/dataframe/merge/merge.py b/mars/dataframe/merge/merge.py
--- a/mars/dataframe/merge/merge.py
+++ b/mars/dataframe/merge/merge.py
@@ -60,8 +60,22 @@ def _set_inputs(self, inputs):
     def execute_map(cls, ctx, op):
         chunk = op.outputs[0]
         df = ctx[op.inputs[0].key]
+        shuffle_on = op.shuffle_on
+
+        if shuffle_on is not None:
+            # shuffle on field may be resident in index
+            to_reset_index_names = []
+            if not isinstance(shuffle_on, (list, tuple)):
+                if shuffle_on not in df.dtypes:
+                    to_reset_index_names.append(shuffle_on)
+            else:
+                for son in shuffle_on:
+                    if son not in df.dtypes:
+                        to_reset_index_names.append(shuffle_on)
+            if len(to_reset_index_names) > 0:
+                df = df.reset_index(to_reset_index_names)
 
-        filters = hash_dataframe_on(df, op.shuffle_on, op.index_shuffle_size)
+        filters = hash_dataframe_on(df, shuffle_on, op.index_shuffle_size)
 
         # shuffle on index
         for index_idx, index_filter in enumerate(filters):
@@ -159,15 +173,35 @@ def indicator(self):
     def validate(self):
         return self._validate
 
+    @staticmethod
+    def _make_value(dtype):
+        # special handle for datetime64 and timedelta64
+        dispatch = {
+            np.datetime64: pd.Timestamp,
+            np.timedelta64: pd.Timedelta,
+        }
+        # otherwise, just use dtype.type itself to convert
+        convert = dispatch.get(dtype.type, dtype.type)
+        return convert(1)
+
+    @staticmethod
+    def _make_data(obj):
+        empty_df = build_empty_df(obj.dtypes, index=obj.index_value.to_pandas()[:0])
+        dtypes = empty_df.dtypes
+        record = [_DataFrameMergeBase._make_value(dtype) for dtype in empty_df.dtypes]
+        if isinstance(empty_df.index, pd.MultiIndex):
+            index = tuple(_DataFrameMergeBase._make_value(level.dtype)
+                          for level in empty_df.index.levels)
+            empty_df.loc[index,] = record
+        else:
+            index = _DataFrameMergeBase._make_value(empty_df.index.dtype)
+            empty_df.loc[index] = record
+        # make sure dtypes correct for MultiIndex
+        empty_df = empty_df.astype(dtypes, copy=False)
+        return empty_df
+
     def __call__(self, left, right):
-        empty_left, empty_right = build_empty_df(left.dtypes), build_empty_df(right.dtypes)
-        # left should have values to keep columns order.
-        gen_left_data = [np.random.rand(1).astype(dt)[0] for dt in left.dtypes]
-        empty_left = empty_left.append(pd.DataFrame([gen_left_data], columns=list(empty_left.columns))
-                                       .astype(left.dtypes))
-        gen_right_data = [np.random.rand(1).astype(dt)[0] for dt in right.dtypes]
-        empty_right = empty_right.append(pd.DataFrame([gen_right_data], columns=list(empty_right.columns))
-                                         .astype(right.dtypes))
+        empty_left, empty_right = self._make_data(left), self._make_data(right)
         # this `merge` will check whether the combination of those arguments is valid
         merged = empty_left.merge(empty_right, how=self.how, on=self.on,
                                   left_on=self.left_on, right_on=self.right_on,
diff --git a/mars/dataframe/utils.py b/mars/dataframe/utils.py
--- a/mars/dataframe/utils.py
+++ b/mars/dataframe/utils.py
@@ -392,7 +392,7 @@ def build_split_idx_to_origin_idx(splits, increase=True):
 
 def build_empty_df(dtypes, index=None):
     columns = dtypes.index
-    df = pd.DataFrame(columns=columns)
+    df = pd.DataFrame(columns=columns, index=index)
     for c, d in zip(columns, dtypes):
         df[c] = pd.Series(dtype=d, index=index)
     return df


Test patch:
diff --git a/mars/dataframe/merge/tests/test_merge_execution.py b/mars/dataframe/merge/tests/test_merge_execution.py
--- a/mars/dataframe/merge/tests/test_merge_execution.py
+++ b/mars/dataframe/merge/tests/test_merge_execution.py
@@ -31,9 +31,15 @@ def setUp(self):
     def testMerge(self):
         df1 = pd.DataFrame(np.arange(20).reshape((4, 5)) + 1, columns=['a', 'b', 'c', 'd', 'e'])
         df2 = pd.DataFrame(np.arange(20).reshape((5, 4)) + 1, columns=['a', 'b', 'x', 'y'])
+        df3 = df1.copy()
+        df3.index = pd.RangeIndex(2, 6, name='index')
+        df4 = df1.copy()
+        df4.index = pd.MultiIndex.from_tuples([(i, i + 1) for i in range(4)], names=['i1', 'i2'])
 
         mdf1 = from_pandas(df1, chunk_size=2)
         mdf2 = from_pandas(df2, chunk_size=2)
+        mdf3 = from_pandas(df3, chunk_size=3)
+        mdf4 = from_pandas(df4, chunk_size=2)
 
         # Note [Index of Merge]
         #
@@ -88,6 +94,24 @@ def testMerge(self):
         result5 = self.executor.execute_dataframe(jdf5, concat=True)[0]
         pd.testing.assert_frame_equal(sort_dataframe_inplace(expected5, 0), sort_dataframe_inplace(result5, 0))
 
+        # merge when some on is index
+        expected6 = df3.merge(df2, how='inner', left_on='index', right_on='a')
+        jdf6 = mdf3.merge(mdf2, how='inner', left_on='index', right_on='a')
+        result6 = self.executor.execute_dataframe(jdf6, concat=True)[0]
+        pd.testing.assert_frame_equal(sort_dataframe_inplace(expected6, 0), sort_dataframe_inplace(result6, 0))
+
+        # merge when on is in MultiIndex
+        expected7 = df4.merge(df2, how='inner', left_on='i1', right_on='a')
+        jdf7 = mdf4.merge(mdf2, how='inner', left_on='i1', right_on='a')
+        result7 = self.executor.execute_dataframe(jdf7, concat=True)[0]
+        pd.testing.assert_frame_equal(sort_dataframe_inplace(expected7, 0), sort_dataframe_inplace(result7, 0))
+
+        # merge when on is in MultiIndex, and on not in index
+        expected8 = df4.merge(df2, how='inner', on=['a', 'b'])
+        jdf8 = mdf4.merge(mdf2, how='inner', on=['a', 'b'])
+        result8 = self.executor.execute_dataframe(jdf8, concat=True)[0]
+        pd.testing.assert_frame_equal(sort_dataframe_inplace(expected8, 0), sort_dataframe_inplace(result8, 0))
+
     def testJoin(self):
         df1 = pd.DataFrame([[1, 3, 3], [4, 2, 6], [7, 8, 9]], index=['a1', 'a2', 'a3'])
         df2 = pd.DataFrame([[1, 2, 3], [1, 5, 6], [7, 8, 9]], index=['a1', 'b2', 'b3']) + 1