# Buggy code
def execute(cls, ctx, op):
    from xgboost import DMatrix

    raw_data = data = ctx[op.data.key]
    if isinstance(data, tuple):
        data = ToDMatrix.get_xgb_dmatrix(data)
    else:
        data = data.spmatrix if hasattr(data, "spmatrix") else data
        data = DMatrix(data)
    result = op.model.predict(data)

    if isinstance(op.outputs[0], DATAFRAME_CHUNK_TYPE):
        result = pd.DataFrame(result, index=raw_data.index)
    elif isinstance(op.outputs[0], SERIES_CHUNK_TYPE):
        result = pd.Series(result, index=raw_data.index, name="predictions")

    ctx[op.outputs[0].key] = result


# Traceback (ValueError)
In [1]: import mars.dataframe as md

In [2]: import numpy as np

In [3]: df = md.DataFrame({'a': np.random.randint(100, size=10), 'b':np.random.rand(10), 'label': np.random.randint(2, size=10)}, chunk_size=4)

In [4]: df = df[df.a > 0]

In [5]: df._shape = (10, 3)

In [6]: data = df[['a', 'b']]

In [7]: label = df['label']

In [8]: from mars.learn.contrib import xgboost as xgb

In [9]: xgb.MarsDMatrix(data=data, label=label).execute()
Out[9]:
Traceback (most recent call last):
File "/Users/qinxuye/miniconda3/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3331, in run_code
exec(code_obj, self.user_global_ns, self.user_ns)
File "<ipython-input-13-43292b123748>", line 1, in <module>
xgb.MarsDMatrix(data=data, label=label).execute()
File "/Users/qinxuye/Workspace/mars/mars/learn/contrib/xgboost/dmatrix.py", line 327, in to_dmatrix
outs.execute(session=session, **(run_kwargs or dict()))
File "/Users/qinxuye/Workspace/mars/mars/core.py", line 736, in execute
return super().execute(session=session, **kw)
File "/Users/qinxuye/Workspace/mars/mars/core.py", line 379, in execute
return run()
File "/Users/qinxuye/Workspace/mars/mars/core.py", line 374, in run
session.run(self, **kw)
File "/Users/qinxuye/Workspace/mars/mars/session.py", line 505, in run
result = self._sess.run(*tileables, **kw)
File "/Users/qinxuye/Workspace/mars/mars/session.py", line 111, in run
res = self._executor.execute_tileables(tileables, **kw)
File "/Users/qinxuye/Workspace/mars/mars/utils.py", line 458, in _inner
return func(*args, **kwargs)
File "/Users/qinxuye/Workspace/mars/mars/executor.py", line 866, in execute_tileables
tileables, tileable_graph=tileable_graph)
File "/Users/qinxuye/Workspace/mars/mars/utils.py", line 458, in _inner
return func(*args, **kwargs)
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 348, in build
tileables, tileable_graph=tileable_graph)
File "/Users/qinxuye/Workspace/mars/mars/utils.py", line 458, in _inner
return func(*args, **kwargs)
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 262, in build
self._on_tile_failure(tileable_data.op, exc_info)
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 301, in inner
raise exc_info[1].with_traceback(exc_info[2]) from None
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 242, in build
tiled = self._tile(tileable_data, tileable_graph)
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 337, in _tile
return super()._tile(tileable_data, tileable_graph)
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 201, in _tile
tds[0]._inplace_tile()
File "/Users/qinxuye/Workspace/mars/mars/core.py", line 168, in _inplace_tile
return handler.inplace_tile(self)
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 136, in inplace_tile
dispatched = self.dispatch(to_tile.op)
File "/Users/qinxuye/Workspace/mars/mars/utils.py", line 458, in _inner
return func(*args, **kwargs)
File "/Users/qinxuye/Workspace/mars/mars/tiles.py", line 119, in dispatch
tiled = op_cls.tile(op)
File "/Users/qinxuye/Workspace/mars/mars/learn/contrib/xgboost/dmatrix.py", line 255, in tile
return cls._tile_multi_output(op)
File "/Users/qinxuye/Workspace/mars/mars/learn/contrib/xgboost/dmatrix.py", line 139, in _tile_multi_output
label = label.rechunk({0: nsplit})._inplace_tile()
File "/Users/qinxuye/Workspace/mars/mars/tensor/rechunk/rechunk.py", line 97, in rechunk
chunk_size = get_nsplits(tensor, chunk_size, tensor.dtype.itemsize)
File "/Users/qinxuye/Workspace/mars/mars/tensor/rechunk/core.py", line 38, in get_nsplits
return decide_chunk_sizes(tileable.shape, chunk_size, itemsize)
File "/Users/qinxuye/Workspace/mars/mars/tensor/utils.py", line 551, in decide_chunk_sizes
return normalize_chunk_sizes(shape, tuple(chunk_size[j] for j in range(len(shape))))
File "/Users/qinxuye/Workspace/mars/mars/tensor/utils.py", line 66, in normalize_chunk_sizes
raise ValueError('chunks shape should be of the same length, '
ValueError: chunks shape should be of the same length, got shape: 10, chunks: (nan, nan, nan)

# Fixed code
def execute(cls, ctx, op):
    from xgboost import DMatrix

    raw_data = data = ctx[op.data.key]
    if isinstance(data, tuple):
        data = ToDMatrix.get_xgb_dmatrix(data)
    else:
        data = data.spmatrix if hasattr(data, "spmatrix") else data
        data = DMatrix(data)

    # do not pass arguments that are None
    kwargs = dict((k, v) for k, v in op.kwargs.items() if v is not None)
    result = op.model.predict(data, **kwargs)

    if isinstance(op.outputs[0], DATAFRAME_CHUNK_TYPE):
        result = pd.DataFrame(result, index=raw_data.index)
    elif isinstance(op.outputs[0], SERIES_CHUNK_TYPE):
        result = pd.Series(result, index=raw_data.index, name="predictions")

    ctx[op.outputs[0].key] = result