# Buggy code
def execute_map(cls, ctx, op):
    chunk = op.outputs[0]
    df = ctx[op.inputs[0].key]

    filters = hash_dataframe_on(df, op.shuffle_on, op.index_shuffle_size)

    # shuffle on index
    for index_idx, index_filter in enumerate(filters):
        group_key = ",".join([str(index_idx), str(chunk.index[1])])
        if index_filter is not None and index_filter is not list():
            ctx[(chunk.key, group_key)] = df.loc[index_filter]
        else:
            ctx[(chunk.key, group_key)] = None


# Traceback (KeyError)
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'

# Fixed code
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, shuffle_on, op.index_shuffle_size)

    # shuffle on index
    for index_idx, index_filter in enumerate(filters):
        group_key = ",".join([str(index_idx), str(chunk.index[1])])
        if index_filter is not None and index_filter is not list():
            ctx[(chunk.key, group_key)] = df.loc[index_filter]
        else:
            ctx[(chunk.key, group_key)] = None