# Buggy code
def _auto_concat_dataframe_chunks(chunk, inputs):
    xdf = (
        pd if isinstance(inputs[0], (pd.DataFrame, pd.Series)) or cudf is None else cudf
    )

    if chunk.op.axis is not None:
        return xdf.concat(inputs, axis=op.axis)

    # auto generated concat when executing a DataFrame
    if len(inputs) == 1:
        ret = inputs[0]
    else:
        max_rows = max(inp.index[0] for inp in chunk.inputs)
        min_rows = min(inp.index[0] for inp in chunk.inputs)
        n_rows = max_rows - min_rows + 1
        n_cols = int(len(inputs) // n_rows)
        assert n_rows * n_cols == len(inputs)

        concats = []
        for i in range(n_rows):
            if n_cols == 1:
                concats.append(inputs[i])
            else:
                concat = xdf.concat(
                    [inputs[i * n_cols + j] for j in range(n_cols)], axis=1
                )
                concats.append(concat)

        if xdf is pd:
            # The `sort=False` is to suppress a `FutureWarning` of pandas,
            # when the index or column of chunks to concatenate is not aligned,
            # which may happens for certain ops.
            #
            # See also Note [Columns of Left Join] in test_merge_execution.py.
            ret = xdf.concat(concats, sort=False)
        else:
            ret = xdf.concat(concats)
            # cuDF will lost index name when concat two seriess.
            ret.index.name = concats[0].index.name

    if getattr(chunk.index_value, "should_be_monotonic", False):
        ret.sort_index(inplace=True)
    if getattr(chunk.columns_value, "should_be_monotonic", False):
        ret.sort_index(axis=1, inplace=True)
    return ret


# 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 _auto_concat_dataframe_chunks(chunk, inputs):
    xdf = (
        pd if isinstance(inputs[0], (pd.DataFrame, pd.Series)) or cudf is None else cudf
    )

    if chunk.op.axis is not None:
        return xdf.concat(inputs, axis=op.axis)

    # auto generated concat when executing a DataFrame
    if len(inputs) == 1:
        ret = inputs[0]
    else:
        n_rows = len(set(inp.index[0] for inp in chunk.inputs))
        n_cols = int(len(inputs) // n_rows)
        assert n_rows * n_cols == len(inputs)

        concats = []
        for i in range(n_rows):
            if n_cols == 1:
                concats.append(inputs[i])
            else:
                concat = xdf.concat(
                    [inputs[i * n_cols + j] for j in range(n_cols)], axis=1
                )
                concats.append(concat)

        if xdf is pd:
            # The `sort=False` is to suppress a `FutureWarning` of pandas,
            # when the index or column of chunks to concatenate is not aligned,
            # which may happens for certain ops.
            #
            # See also Note [Columns of Left Join] in test_merge_execution.py.
            ret = xdf.concat(concats, sort=False)
        else:
            ret = xdf.concat(concats)
            # cuDF will lost index name when concat two seriess.
            ret.index.name = concats[0].index.name

    if getattr(chunk.index_value, "should_be_monotonic", False):
        ret.sort_index(inplace=True)
    if getattr(chunk.columns_value, "should_be_monotonic", False):
        ret.sort_index(axis=1, inplace=True)
    return ret