# Buggy code
def _tile_head_tail(cls, op):
    from ..merge import DataFrameConcat

    inp = op.input
    out = op.outputs[0]
    combine_size = options.combine_size

    chunks = inp.chunks

    new_chunks = []
    for c in chunks:
        chunk_op = op.copy().reset_key()
        params = out.params
        params["index"] = c.index
        new_chunks.append(chunk_op.new_chunk([c], kws=[params]))
    chunks = new_chunks

    while len(chunks) > 1:
        new_size = ceildiv(len(chunks), combine_size)
        new_chunks = []
        for i in range(new_size):
            in_chunks = chunks[combine_size * i : combine_size * (i + 1)]
            chunk_index = (i, 0) if in_chunks[0].ndim == 2 else (i,)
            if len(inp.shape) == 1:
                shape = (sum(c.shape[0] for c in in_chunks),)
            else:
                shape = (sum(c.shape[0] for c in in_chunks), in_chunks[0].shape[1])
            concat_chunk = DataFrameConcat(
                axis=0, output_types=in_chunks[0].op.output_types
            ).new_chunk(in_chunks, index=chunk_index, shape=shape)
            chunk_op = op.copy().reset_key()
            params = out.params
            params["index"] = chunk_index
            new_chunks.append(chunk_op.new_chunk([concat_chunk], kws=[params]))
        chunks = new_chunks

    new_op = op.copy()
    params = out.params
    params["nsplits"] = tuple((s,) for s in out.shape)
    params["chunks"] = chunks
    return new_op.new_tileables(op.inputs, kws=[params])


# 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 _tile_head_tail(cls, op):
    from ..merge import DataFrameConcat

    inp = op.input
    out = op.outputs[0]
    combine_size = options.combine_size

    chunks = inp.chunks

    new_chunks = []
    for c in chunks:
        chunk_op = op.copy().reset_key()
        params = out.params
        params["index"] = c.index
        params["shape"] = c.shape if np.isnan(c.shape[0]) else out.shape
        new_chunks.append(chunk_op.new_chunk([c], kws=[params]))
    chunks = new_chunks

    while len(chunks) > 1:
        new_size = ceildiv(len(chunks), combine_size)
        new_chunks = []
        for i in range(new_size):
            in_chunks = chunks[combine_size * i : combine_size * (i + 1)]
            chunk_index = (i, 0) if in_chunks[0].ndim == 2 else (i,)
            if len(inp.shape) == 1:
                shape = (sum(c.shape[0] for c in in_chunks),)
            else:
                shape = (sum(c.shape[0] for c in in_chunks), in_chunks[0].shape[1])
            concat_chunk = DataFrameConcat(
                axis=0, output_types=in_chunks[0].op.output_types
            ).new_chunk(in_chunks, index=chunk_index, shape=shape)
            chunk_op = op.copy().reset_key()
            params = out.params
            params["index"] = chunk_index
            params["shape"] = (
                in_chunks[0].shape if np.isnan(in_chunks[0].shape[0]) else out.shape
            )
            new_chunks.append(chunk_op.new_chunk([concat_chunk], kws=[params]))
        chunks = new_chunks

    new_op = op.copy()
    params = out.params
    params["nsplits"] = tuple((s,) for s in out.shape)
    params["chunks"] = chunks
    return new_op.new_tileables(op.inputs, kws=[params])