Repository: mars-project/mars

Issue:
[BUG] XGBoost training failed
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

XGBoost training failed.

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version
2. The version of Mars you use
3. Versions of crucial packages, such as numpy, scipy and protobuf
4. Full stack of the error.
5. Minimized code to reproduce the error.

```
Traceback (most recent call last):
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/scheduler/graph.py", line 410, in _execute_graph
    self.prepare_graph(compose=compose)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/utils.py", line 377, in _wrapped
    return func(*args, **kwargs)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/utils.py", line 451, in _inner
    return func(*args, **kwargs)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/scheduler/graph.py", line 648, in prepare_graph
    self._target_tileable_datas + fetch_tileables, tileable_graph)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/utils.py", line 451, in _inner
    return func(*args, **kwargs)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/tiles.py", line 348, in build
    tileables, tileable_graph=tileable_graph)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/utils.py", line 451, in _inner
    return func(*args, **kwargs)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/tiles.py", line 262, in build
    self._on_tile_failure(tileable_data.op, exc_info)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/tiles.py", line 301, in inner
    raise exc_info[1].with_traceback(exc_info[2]) from None
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/tiles.py", line 242, in build
    tiled = self._tile(tileable_data, tileable_graph)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/tiles.py", line 337, in _tile
    return super()._tile(tileable_data, tileable_graph)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/tiles.py", line 203, in _tile
    tds = on_tile(tileable_data.op.outputs, tds)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/scheduler/graph.py", line 630, in on_tile
    return self.context.wraps(handler.dispatch)(first.op)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/context.py", line 72, in h
    return func(*args, **kwargs)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/utils.py", line 451, in _inner
    return func(*args, **kwargs)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/tiles.py", line 119, in dispatch
    tiled = op_cls.tile(op)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/learn/contrib/xgboost/dmatrix.py", line 255, in tile
    return cls._tile_multi_output(op)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/learn/contrib/xgboost/dmatrix.py", line 139, in _tile_multi_output
    label = label.rechunk({0: nsplit})._inplace_tile()
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/tensor/rechunk/rechunk.py", line 81, in rechunk
    chunk_size = get_nsplits(tensor, chunk_size, tensor.dtype.itemsize)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/tensor/rechunk/core.py", line 38, in get_nsplits
    return decide_chunk_sizes(tileable.shape, chunk_size, itemsize)
  File "/home/admin/work/_public-mars-0.6.2.zip/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 "/home/admin/work/_public-mars-0.6.2.zip/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: 1000000, chunks: (nan,)
```



Discussion:

Resolving patch:
diff --git a/mars/dataframe/base/rechunk.py b/mars/dataframe/base/rechunk.py
--- a/mars/dataframe/base/rechunk.py
+++ b/mars/dataframe/base/rechunk.py
@@ -23,6 +23,7 @@
 from ...tensor.utils import calc_sliced_size
 from ...tiles import TilesError
 from ...utils import check_chunks_unknown_shape
+from ..initializer import DataFrame as asdataframe, Series as asseries
 from ..operands import DataFrameOperand, DataFrameOperandMixin, DATAFRAME_TYPE
 from ..utils import indexing_index_value, merge_index_value
 
@@ -41,6 +42,10 @@ def __init__(self, chunk_size=None, threshold=None, chunk_size_limit=None, outpu
                          _chunk_size_limit=chunk_size_limit, _output_types=output_types,
                          _reassign_worker=reassign_worker, **kw)
 
+    @property
+    def input(self):
+        return self._input
+
     @property
     def chunk_size(self):
         return self._chunk_size
@@ -69,8 +74,15 @@ def __call__(self, x):
     @classmethod
     def tile(cls, op):
         check_chunks_unknown_shape(op.inputs, TilesError)
+
+        a = op.input
+        a = asdataframe(a) if a.ndim == 2 else asseries(a)
+        chunk_size = _get_chunk_size(a, op.chunk_size)
+        if chunk_size == a.nsplits:
+            return [a]
+
         out = op.outputs[0]
-        new_chunk_size = op.chunk_size
+        new_chunk_size = chunk_size
         if isinstance(out, DATAFRAME_TYPE):
             itemsize = max(getattr(dt, 'itemsize', 8) for dt in out.dtypes)
         else:
@@ -88,16 +100,25 @@ def tile(cls, op):
         return [out]
 
 
-def rechunk(a, chunk_size, threshold=None, chunk_size_limit=None, reassign_worker=False):
+def _get_chunk_size(a, chunk_size):
     if isinstance(a, DATAFRAME_TYPE):
         itemsize = max(getattr(dt, 'itemsize', 8) for dt in a.dtypes)
     else:
         itemsize = a.dtype.itemsize
-    chunk_size = get_nsplits(a, chunk_size, itemsize)
-    if chunk_size == a.nsplits:
-        return a
+    return get_nsplits(a, chunk_size, itemsize)
+
 
-    op = DataFrameRechunk(chunk_size, threshold, chunk_size_limit, reassign_worker=reassign_worker)
+def rechunk(a, chunk_size, threshold=None, chunk_size_limit=None, reassign_worker=False):
+    if not any(pd.isna(s) for s in a.shape):
+        # do client check only when no unknown shape,
+        # real nsplits will be recalculated inside `tile`
+        chunk_size = _get_chunk_size(a, chunk_size)
+        if chunk_size == a.nsplits:
+            return a
+
+    op = DataFrameRechunk(chunk_size=chunk_size, threshold=threshold,
+                          chunk_size_limit=chunk_size_limit,
+                          reassign_worker=reassign_worker)
     return op(a)
 
 
diff --git a/mars/dataframe/datastore/to_csv.py b/mars/dataframe/datastore/to_csv.py
--- a/mars/dataframe/datastore/to_csv.py
+++ b/mars/dataframe/datastore/to_csv.py
@@ -176,7 +176,7 @@ def tile(cls, op: 'DataFrameToCSV'):
         in_df = op.input
         out_df = op.outputs[0]
 
-        if in_df.ndim == 2:
+        if in_df.ndim == 2 and in_df.chunk_shape[1] > 1:
             # make sure only 1 chunk on the column axis
             in_df = in_df.rechunk({1: in_df.shape[1]})._inplace_tile()
 
diff --git a/mars/tensor/rechunk/rechunk.py b/mars/tensor/rechunk/rechunk.py
--- a/mars/tensor/rechunk/rechunk.py
+++ b/mars/tensor/rechunk/rechunk.py
@@ -14,12 +14,15 @@
 
 import itertools
 
+import numpy as np
+
 from ... import opcodes as OperandDef
 from ...serialize import KeyField, AnyField, BoolField, Int32Field, Int64Field
 from ...utils import check_chunks_unknown_shape
 from ...tiles import TilesError
 from ..utils import calc_sliced_size
 from ..operands import TensorHasInput, TensorOperandMixin
+from ..datasource import tensor as astensor
 from .core import plan_rechunks, get_nsplits, compute_rechunk_slices
 
 
@@ -39,6 +42,10 @@ def __init__(self, chunk_size=None, threshold=None, chunk_size_limit=None,
                          _reassign_worker=reassign_worker,
                          _dtype=dtype, _sparse=sparse, **kw)
 
+    @property
+    def input(self):
+        return self._input
+
     @property
     def chunk_size(self):
         return self._chunk_size
@@ -61,6 +68,12 @@ def __call__(self, tensor):
     @classmethod
     def tile(cls, op):
         check_chunks_unknown_shape(op.inputs, TilesError)
+
+        tensor = astensor(op.input)
+        chunk_size = get_nsplits(tensor, op.chunk_size, tensor.dtype.itemsize)
+        if chunk_size == tensor.nsplits:
+            return [tensor]
+
         new_chunk_size = op.chunk_size
         steps = plan_rechunks(op.inputs[0], new_chunk_size, op.inputs[0].dtype.itemsize,
                               threshold=op.threshold,
@@ -78,9 +91,12 @@ def tile(cls, op):
 
 def rechunk(tensor, chunk_size, threshold=None, chunk_size_limit=None,
             reassign_worker=False):
-    chunk_size = get_nsplits(tensor, chunk_size, tensor.dtype.itemsize)
-    if chunk_size == tensor.nsplits:
-        return tensor
+    if not any(np.isnan(s) for s in tensor.shape):
+        # do client check only when tensor has no unknown shape,
+        # otherwise, recalculate chunk_size in `tile`
+        chunk_size = get_nsplits(tensor, chunk_size, tensor.dtype.itemsize)
+        if chunk_size == tensor.nsplits:
+            return tensor
 
     op = TensorRechunk(chunk_size, threshold, chunk_size_limit, reassign_worker=reassign_worker,
                        dtype=tensor.dtype, sparse=tensor.issparse())


Test patch:
diff --git a/mars/dataframe/datastore/tests/test_datastore_execute.py b/mars/dataframe/datastore/tests/test_datastore_execute.py
--- a/mars/dataframe/datastore/tests/test_datastore_execute.py
+++ b/mars/dataframe/datastore/tests/test_datastore_execute.py
@@ -24,7 +24,7 @@
 from mars.dataframe import DataFrame
 from mars.deploy.local.core import new_cluster
 from mars.session import new_session
-from mars.tests.core import mock, TestBase, ExecutorForTest
+from mars.tests.core import mock, TestBase
 
 try:
     import vineyard
@@ -47,7 +47,7 @@
 class Test(TestBase):
     def setUp(self):
         super().setUp()
-        self.executor = ExecutorForTest()
+        self.ctx, self.executor = self._create_test_context()
 
     def testToCSVExecution(self):
         index = pd.RangeIndex(100, 0, -1, name='index')
@@ -83,6 +83,18 @@ def testToCSVExecution(self):
             pd.testing.assert_frame_equal(result, raw)
             pd.testing.assert_frame_equal(dfs[1].set_index('index'), raw.iloc[33: 66])
 
+            with self.ctx:
+                # test df with unknown shape
+                df2 = DataFrame(raw, chunk_size=(50, 2))
+                df2 = df2[df2['col1'] < 1]
+                path2 = os.path.join(base_path, 'out2.csv')
+                r = df2.to_csv(path2)
+                self.executor.execute_dataframes([r])
+
+                result = pd.read_csv(path2, dtype=raw.dtypes.to_dict())
+                result.set_index('index', inplace=True)
+                pd.testing.assert_frame_equal(result, raw)
+
             # SERIES TESTS
             series = md.Series(raw.col1, chunk_size=33)
 
diff --git a/mars/learn/contrib/xgboost/tests/test_classifier.py b/mars/learn/contrib/xgboost/tests/test_classifier.py
--- a/mars/learn/contrib/xgboost/tests/test_classifier.py
+++ b/mars/learn/contrib/xgboost/tests/test_classifier.py
@@ -117,6 +117,17 @@ def testLocalClassifier(self):
 
         self.assertEqual(prediction.ndim, 1)
 
+        # test train with unknown shape
+        cond = X[:, 0] > 0
+        X3 = X[cond]
+        y3 = y[cond]
+        classifier = XGBClassifier(verbosity=1, n_estimators=2)
+        classifier.fit(X3, y3)
+        prediction = classifier.predict(X)
+
+        self.assertEqual(prediction.ndim, 1)
+        self.assertEqual(prediction.shape[0], len(self.X))
+
         classifier = XGBClassifier(verbosity=1, n_estimators=2)
         with self.assertRaises(TypeError):
             classifier.fit(X, y, wrong_param=1)