# Buggy code
def execute(cls, ctx, op):
    def _base_concat(chunk, inputs):
        # auto generated concat when executing a DataFrame, Series or Index
        if chunk.op.output_types[0] == OutputType.dataframe:
            return _auto_concat_dataframe_chunks(chunk, inputs)
        elif chunk.op.output_types[0] == OutputType.series:
            return _auto_concat_series_chunks(chunk, inputs)
        elif chunk.op.output_types[0] == OutputType.index:
            return _auto_concat_index_chunks(chunk, inputs)
        elif chunk.op.output_types[0] == OutputType.categorical:
            return _auto_concat_categorical_chunks(chunk, inputs)
        else:  # pragma: no cover
            raise TypeError(
                "Only DataFrameChunk, SeriesChunk, IndexChunk, "
                "and CategoricalChunk can be automatically concatenated"
            )

    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

    def _auto_concat_series_chunks(chunk, inputs):
        # auto generated concat when executing a Series
        if len(inputs) == 1:
            concat = inputs[0]
        else:
            xdf = pd if isinstance(inputs[0], pd.Series) or cudf is None else cudf
            if chunk.op.axis is not None:
                concat = xdf.concat(inputs, axis=chunk.op.axis)
            else:
                concat = xdf.concat(inputs)
        if getattr(chunk.index_value, "should_be_monotonic", False):
            concat.sort_index(inplace=True)
        return concat

    def _auto_concat_index_chunks(chunk, inputs):
        if len(inputs) == 1:
            xdf = pd if isinstance(inputs[0], pd.Index) or cudf is None else cudf
            concat_df = xdf.DataFrame(index=inputs[0])
        else:
            xdf = pd if isinstance(inputs[0], pd.Index) or cudf is None else cudf
            empty_dfs = [xdf.DataFrame(index=inp) for inp in inputs]
            concat_df = xdf.concat(empty_dfs, axis=0)
        if getattr(chunk.index_value, "should_be_monotonic", False):
            concat_df.sort_index(inplace=True)
        return concat_df.index

    def _auto_concat_categorical_chunks(_, inputs):
        if len(inputs) == 1:  # pragma: no cover
            return inputs[0]
        else:
            # convert categorical into array
            arrays = [np.asarray(inp) for inp in inputs]
            array = np.concatenate(arrays)
            return pd.Categorical(
                array, categories=inputs[0].categories, ordered=inputs[0].ordered
            )

    chunk = op.outputs[0]
    inputs = [ctx[input.key] for input in op.inputs]

    if isinstance(inputs[0], tuple):
        ctx[chunk.key] = tuple(
            _base_concat(chunk, [input[i] for input in inputs])
            for i in range(len(inputs[0]))
        )
    else:
        ctx[chunk.key] = _base_concat(chunk, inputs)


# 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 execute(cls, ctx, op):
    def _base_concat(chunk, inputs):
        # auto generated concat when executing a DataFrame, Series or Index
        if chunk.op.output_types[0] == OutputType.dataframe:
            return _auto_concat_dataframe_chunks(chunk, inputs)
        elif chunk.op.output_types[0] == OutputType.series:
            return _auto_concat_series_chunks(chunk, inputs)
        elif chunk.op.output_types[0] == OutputType.index:
            return _auto_concat_index_chunks(chunk, inputs)
        elif chunk.op.output_types[0] == OutputType.categorical:
            return _auto_concat_categorical_chunks(chunk, inputs)
        else:  # pragma: no cover
            raise TypeError(
                "Only DataFrameChunk, SeriesChunk, IndexChunk, "
                "and CategoricalChunk can be automatically concatenated"
            )

    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

    def _auto_concat_series_chunks(chunk, inputs):
        # auto generated concat when executing a Series
        if len(inputs) == 1:
            concat = inputs[0]
        else:
            xdf = pd if isinstance(inputs[0], pd.Series) or cudf is None else cudf
            if chunk.op.axis is not None:
                concat = xdf.concat(inputs, axis=chunk.op.axis)
            else:
                concat = xdf.concat(inputs)
        if getattr(chunk.index_value, "should_be_monotonic", False):
            concat.sort_index(inplace=True)
        return concat

    def _auto_concat_index_chunks(chunk, inputs):
        if len(inputs) == 1:
            xdf = pd if isinstance(inputs[0], pd.Index) or cudf is None else cudf
            concat_df = xdf.DataFrame(index=inputs[0])
        else:
            xdf = pd if isinstance(inputs[0], pd.Index) or cudf is None else cudf
            empty_dfs = [xdf.DataFrame(index=inp) for inp in inputs]
            concat_df = xdf.concat(empty_dfs, axis=0)
        if getattr(chunk.index_value, "should_be_monotonic", False):
            concat_df.sort_index(inplace=True)
        return concat_df.index

    def _auto_concat_categorical_chunks(_, inputs):
        if len(inputs) == 1:  # pragma: no cover
            return inputs[0]
        else:
            # convert categorical into array
            arrays = [np.asarray(inp) for inp in inputs]
            array = np.concatenate(arrays)
            return pd.Categorical(
                array, categories=inputs[0].categories, ordered=inputs[0].ordered
            )

    chunk = op.outputs[0]
    inputs = [ctx[input.key] for input in op.inputs]

    if isinstance(inputs[0], tuple):
        ctx[chunk.key] = tuple(
            _base_concat(chunk, [input[i] for input in inputs])
            for i in range(len(inputs[0]))
        )
    else:
        ctx[chunk.key] = _base_concat(chunk, inputs)