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

    deliver_by = False  # output by for the upcoming process
    if isinstance(by, list):
        new_by = []
        for v in by:
            if isinstance(v, Base):
                deliver_by = True
                new_by.append(ctx[v.key])
            else:
                new_by.append(v)
        by = new_by

    if isinstance(by, list) or callable(by):
        on = by
    else:
        on = None

    if isinstance(df, tuple):
        filters = hash_dataframe_on(df[0], on, op.shuffle_size, level=op.level)
    else:
        filters = hash_dataframe_on(df, on, op.shuffle_size, level=op.level)

    for index_idx, index_filter in enumerate(filters):
        if is_dataframe_obj:
            group_key = ",".join([str(index_idx), str(chunk.index[1])])
        else:
            group_key = str(index_idx)

        if deliver_by:
            filtered_by = []
            for v in by:
                if isinstance(v, pd.Series):
                    filtered_by.append(v.loc[index_filter])
                else:
                    filtered_by.append(v)
            if isinstance(df, tuple):
                ctx[(chunk.key, group_key)] = tuple(x.loc[index_filter] for x in df) + (
                    filtered_by,
                    deliver_by,
                )
            else:
                ctx[(chunk.key, group_key)] = (
                    df.loc[index_filter],
                    filtered_by,
                    deliver_by,
                )
        else:
            if isinstance(df, tuple):
                ctx[(chunk.key, group_key)] = tuple(x.loc[index_filter] for x in df) + (
                    deliver_by,
                )
            else:
                ctx[(chunk.key, group_key)] = df.loc[index_filter]


# Traceback (KeyError)
2020-12-02 11:19:40,309 mars.scheduler.operands.common 87 ERROR    Attempt 1: Unexpected error KeyError occurred in executing operand 5c7a3b06d448300987640036d2f5a34e in 11.238.145.234:49708
Traceback (most recent call last):
File "/home/admin/work/_public-mars-0.5.5.zip/mars/promise.py", line 378, in _wrapped
return func(*args, **kwargs)
File "/home/admin/work/_public-mars-0.5.5.zip/mars/utils.py", line 377, in _wrapped
return func(*args, **kwargs)
File "/home/admin/work/_public-mars-0.5.5.zip/mars/worker/execution.py", line 564, in execute_graph
quota_request = self._prepare_quota_request(session_id, graph_key)
File "/home/admin/work/_public-mars-0.5.5.zip/mars/worker/execution.py", line 249, in _prepare_quota_request
memory_estimations = self._estimate_calc_memory(session_id, graph_key)
File "/home/admin/work/_public-mars-0.5.5.zip/mars/worker/execution.py", line 213, in _estimate_calc_memory
res = executor.execute_graph(graph_record.graph, graph_record.chunk_targets, mock=True)
File "/home/admin/work/_public-mars-0.5.5.zip/mars/executor.py", line 690, in execute_graph
res = graph_execution.execute(retval)
File "/home/admin/work/_public-mars-0.5.5.zip/mars/executor.py", line 574, in execute
return [self._chunk_results[key] for key in self._keys]
File "/home/admin/work/_public-mars-0.5.5.zip/mars/executor.py", line 574, in <listcomp>
return [self._chunk_results[key] for key in self._keys]
KeyError: '3990ec90331559138b6ecbc6d76fbd0d'

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

    deliver_by = False  # output by for the upcoming process
    if isinstance(by, list):
        new_by = []
        for v in by:
            if isinstance(v, Base):
                deliver_by = True
                new_by.append(ctx[v.key])
            else:
                new_by.append(v)
        by = new_by

    if isinstance(by, list) or callable(by):
        on = by
    else:
        on = None

    if isinstance(df, tuple):
        filters = hash_dataframe_on(df[0], on, op.shuffle_size, level=op.level)
    else:
        filters = hash_dataframe_on(df, on, op.shuffle_size, level=op.level)

    def _take_index(src, f):
        result = src.loc[f]
        if src.index.names:
            result.index.names = src.index.names
        return result

    for index_idx, index_filter in enumerate(filters):
        if is_dataframe_obj:
            group_key = ",".join([str(index_idx), str(chunk.index[1])])
        else:
            group_key = str(index_idx)

        if deliver_by:
            filtered_by = []
            for v in by:
                if isinstance(v, pd.Series):
                    filtered_by.append(_take_index(v, index_filter))
                else:
                    filtered_by.append(v)
            if isinstance(df, tuple):
                ctx[(chunk.key, group_key)] = tuple(
                    _take_index(x, index_filter) for x in df
                ) + (filtered_by, deliver_by)
            else:
                ctx[(chunk.key, group_key)] = (
                    _take_index(df, index_filter),
                    filtered_by,
                    deliver_by,
                )
        else:
            if isinstance(df, tuple):
                ctx[(chunk.key, group_key)] = tuple(
                    _take_index(x, index_filter) for x in df
                ) + (deliver_by,)
            else:
                ctx[(chunk.key, group_key)] = _take_index(df, index_filter)