# Buggy code
def _gen_shuffle_chunks(cls, op, in_df, chunks):
    # generate map chunks
    map_chunks = []
    chunk_shape = (in_df.chunk_shape[0], 1)
    for chunk in chunks:
        # no longer consider as_index=False for the intermediate phases,
        # will do reset_index at last if so
        map_op = DataFrameGroupByOperand(
            stage=OperandStage.map,
            shuffle_size=chunk_shape[0],
            output_types=[OutputType.dataframe_groupby],
        )
        map_chunks.append(
            map_op.new_chunk(
                [chunk],
                shape=(np.nan, np.nan),
                index=chunk.index,
                index_value=op.outputs[0].index_value,
            )
        )

    proxy_chunk = DataFrameShuffleProxy(output_types=[OutputType.dataframe]).new_chunk(
        map_chunks, shape=()
    )

    # generate reduce chunks
    reduce_chunks = []
    for out_idx in itertools.product(*(range(s) for s in chunk_shape)):
        reduce_op = DataFrameGroupByOperand(
            stage=OperandStage.reduce,
            shuffle_key=",".join(str(idx) for idx in out_idx),
            output_types=[OutputType.dataframe_groupby],
        )
        reduce_chunks.append(
            reduce_op.new_chunk(
                [proxy_chunk],
                shape=(np.nan, np.nan),
                index=out_idx,
                index_value=op.outputs[0].index_value,
            )
        )

    return reduce_chunks


# 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 _gen_shuffle_chunks(cls, op, in_df, chunks):
    # generate map chunks
    map_chunks = []
    chunk_shape = (in_df.chunk_shape[0], 1)
    for chunk in chunks:
        # no longer consider as_index=False for the intermediate phases,
        # will do reset_index at last if so
        map_op = DataFrameGroupByOperand(
            stage=OperandStage.map,
            shuffle_size=chunk_shape[0],
            output_types=[OutputType.dataframe_groupby],
        )
        map_chunks.append(
            map_op.new_chunk(
                [chunk],
                shape=(np.nan, np.nan),
                index=chunk.index,
                index_value=op.outputs[0].index_value,
            )
        )

    proxy_chunk = DataFrameShuffleProxy(output_types=[OutputType.dataframe]).new_chunk(
        map_chunks, shape=()
    )

    # generate reduce chunks
    reduce_chunks = []
    for out_idx in itertools.product(*(range(s) for s in chunk_shape)):
        reduce_op = DataFrameGroupByOperand(
            stage=OperandStage.reduce,
            shuffle_key=",".join(str(idx) for idx in out_idx),
            output_types=[OutputType.dataframe_groupby],
        )
        reduce_chunks.append(
            reduce_op.new_chunk(
                [proxy_chunk], shape=(np.nan, np.nan), index=out_idx, index_value=None
            )
        )

    return reduce_chunks