# Buggy code
def _calc_dataframe_params(cls, chunk_index_to_chunks, chunk_shape):
    dtypes = pd.concat(
        [chunk_index_to_chunks[0, i].dtypes for i in range(chunk_shape[1])]
    )
    columns_value = parse_index(dtypes.index, store_data=True)
    pd_indxes = [
        chunk_index_to_chunks[i, 0].index_value.to_pandas()
        for i in range(chunk_shape[0])
    ]
    pd_index = reduce(lambda x, y: x.append(y), pd_indxes)
    index_value = parse_index(pd_index)
    return {
        "dtypes": dtypes,
        "columns_value": columns_value,
        "index_value": index_value,
    }


# Traceback (KeyError)
Error
Traceback (most recent call last):
File "/Users/wenjun.swj/Code/mars/mars/scheduler/graph.py", line 410, in _execute_graph
self.prepare_graph(compose=compose)
File "/Users/wenjun.swj/Code/mars/mars/utils.py", line 377, in _wrapped
return func(*args, **kwargs)
File "/Users/wenjun.swj/Code/mars/mars/utils.py", line 451, in _inner
return func(*args, **kwargs)
File "/Users/wenjun.swj/Code/mars/mars/scheduler/graph.py", line 646, in prepare_graph
cur_chunk_graph = chunk_graph_builder.build(
File "/Users/wenjun.swj/Code/mars/mars/utils.py", line 451, in _inner
return func(*args, **kwargs)
File "/Users/wenjun.swj/Code/mars/mars/tiles.py", line 347, in build
chunk_graph = super().build(
File "/Users/wenjun.swj/Code/mars/mars/utils.py", line 451, in _inner
return func(*args, **kwargs)
File "/Users/wenjun.swj/Code/mars/mars/tiles.py", line 262, in build
self._on_tile_failure(tileable_data.op, exc_info)
File "/Users/wenjun.swj/Code/mars/mars/tiles.py", line 301, in inner
raise exc_info[1].with_traceback(exc_info[2]) from None
File "/Users/wenjun.swj/Code/mars/mars/tiles.py", line 242, in build
tiled = self._tile(tileable_data, tileable_graph)
File "/Users/wenjun.swj/Code/mars/mars/tiles.py", line 337, in _tile
return super()._tile(tileable_data, tileable_graph)
File "/Users/wenjun.swj/Code/mars/mars/tiles.py", line 203, in _tile
tds = on_tile(tileable_data.op.outputs, tds)
File "/Users/wenjun.swj/Code/mars/mars/scheduler/graph.py", line 630, in on_tile
return self.context.wraps(handler.dispatch)(first.op)
File "/Users/wenjun.swj/Code/mars/mars/context.py", line 72, in h
return func(*args, **kwargs)
File "/Users/wenjun.swj/Code/mars/mars/utils.py", line 451, in _inner
return func(*args, **kwargs)
File "/Users/wenjun.swj/Code/mars/mars/tiles.py", line 119, in dispatch
tiled = op_cls.tile(op)
File "/Users/wenjun.swj/Code/mars/mars/learn/contrib/xgboost/dmatrix.py", line 257, in tile
return cls._tile_single_output(op)
File "/Users/wenjun.swj/Code/mars/mars/learn/contrib/xgboost/dmatrix.py", line 227, in _tile_single_output
data_chunk = concat_chunks(chunks[0])
File "/Users/wenjun.swj/Code/mars/mars/learn/utils/core.py", line 33, in concat_chunks
tileable = chunks[0].op.create_tileable_from_chunks(chunks)
File "/Users/wenjun.swj/Code/mars/mars/dataframe/operands.py", line 212, in create_tileable_from_chunks
params = cls._calc_dataframe_params(chunk_index_to_chunk, chunk_shape)
File "/Users/wenjun.swj/Code/mars/mars/dataframe/operands.py", line 228, in _calc_dataframe_params
pd_indxes = [chunk_index_to_chunks[i, 0].index_value.to_pandas()
File "/Users/wenjun.swj/Code/mars/mars/dataframe/operands.py", line 228, in <listcomp>
pd_indxes = [chunk_index_to_chunks[i, 0].index_value.to_pandas()
KeyError: (2, 0)

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/Users/wenjun.swj/miniconda3/lib/python3.8/unittest/case.py", line 60, in testPartExecutor
yield
File "/Users/wenjun.swj/miniconda3/lib/python3.8/unittest/case.py", line 676, in run
self._callTestMethod(testMethod)
File "/Users/wenjun.swj/miniconda3/lib/python3.8/unittest/case.py", line 633, in _callTestMethod
method()
File "/Users/wenjun.swj/Code/mars/mars/learn/contrib/xgboost/tests/integrated/test_distributed_xgboost.py", line 71, in testDistributedXGBClassifier
classifier.fit(X, y, eval_set=[(X, y)], session=sess, run_kwargs=run_kwargs)
File "/Users/wenjun.swj/Code/mars/mars/learn/contrib/xgboost/classifier.py", line 55, in fit
result = train(params, dtrain, num_boost_round=self.get_num_boosting_rounds(),
File "/Users/wenjun.swj/Code/mars/mars/learn/contrib/xgboost/train.py", line 200, in train
ret = t.execute(session=session, **run_kwargs).fetch(session=session)
File "/Users/wenjun.swj/Code/mars/mars/core.py", line 379, in execute
return run()
File "/Users/wenjun.swj/Code/mars/mars/core.py", line 374, in run
session.run(self, **kw)
File "/Users/wenjun.swj/Code/mars/mars/session.py", line 499, in run
result = self._sess.run(*tileables, **kw)
File "/Users/wenjun.swj/Code/mars/mars/web/session.py", line 214, in run
if self._check_response_finished(graph_url, timeout_val):
File "/Users/wenjun.swj/Code/mars/mars/web/session.py", line 174, in _check_response_finished
raise ExecutionFailed('Graph execution failed.') from exc
mars.errors.ExecutionFailed: 'Graph execution failed.'

# Fixed code
def _calc_dataframe_params(cls, chunk_index_to_chunks, chunk_shape):
    dtypes = pd.concat(
        [
            chunk_index_to_chunks[0, i].dtypes
            for i in range(chunk_shape[1])
            if (0, i) in chunk_index_to_chunks
        ]
    )
    columns_value = parse_index(dtypes.index, store_data=True)
    pd_indexes = [
        chunk_index_to_chunks[i, 0].index_value.to_pandas()
        for i in range(chunk_shape[0])
        if (i, 0) in chunk_index_to_chunks
    ]
    pd_index = reduce(lambda x, y: x.append(y), pd_indexes)
    index_value = parse_index(pd_index)
    return {
        "dtypes": dtypes,
        "columns_value": columns_value,
        "index_value": index_value,
    }