Repository: mars-project/mars

Issue:
[BUG] Serialization failed when calling `mars.learn.utils.shuffle`
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

Serialization failed when calling `mars.learn.utils.shuffle`

**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.

```python
In [14]: from mars.learn.utils import shuffle                                                                                                                                                               

In [15]: X, y = shuffle(X, y)                                                                                                                                                                               
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~/Workspace/mars/mars/serialize/jsonserializer.pyx in mars.serialize.jsonserializer.JsonSerializeProvider.serialize_field()

~/Workspace/mars/mars/serialize/jsonserializer.pyx in mars.serialize.jsonserializer.JsonSerializeProvider._serialize_value()

~/Workspace/mars/mars/serialize/jsonserializer.pyx in mars.serialize.jsonserializer.JsonSerializeProvider._serialize_typed_value()

TypeError: Expected tuple, got numpy.ndarray

The above exception was the direct cause of the following exception:

TypeError                                 Traceback (most recent call last)
<ipython-input-15-0c15ee335b21> in <module>
----> 1 X, y = shuffle(X, y)

~/Workspace/mars/mars/learn/utils/shuffle.py in shuffle(*arrays, **options)
    413     op = LearnShuffle(axes=axes, seeds=seeds,
    414                       output_types=get_output_types(*arrays))
--> 415     shuffled_arrays = op(arrays)
    416     if len(arrays) == 1:
    417         return shuffled_arrays[0]

~/Workspace/mars/mars/learn/utils/shuffle.py in __call__(self, arrays)
     92     def __call__(self, arrays):
     93         params = self._calc_params([ar.params for ar in arrays])
---> 94         return self.new_tileables(arrays, kws=params)
     95 
     96     def _shuffle_index_value(self, index_value):

~/Workspace/mars/mars/operands.py in new_tileables(self, inputs, kws, **kw)
    351         tileables = self._new_tileables(inputs, kws=kws, **kw)
    352         if is_eager_mode():
--> 353             ExecutableTuple(tileables).execute(fetch=False)
    354         return tileables
    355 

~/Workspace/mars/mars/core.py in execute(self, session, **kw)
    626         if session is None:
    627             session = Session.default_or_local()
--> 628         return session.run(self, **kw)
    629 
    630     def fetch(self, session=None, **kw):

~/Workspace/mars/mars/session.py in run(self, *tileables, **kw)
    184         tileables = tuple(mt.tensor(t) if not isinstance(t, (Entity, Base)) else t
    185                           for t in tileables)
--> 186         result = self._sess.run(*tileables, **kw)
    187 
    188         for t in tileables:

~/Workspace/mars/mars/deploy/local/session.py in run(self, *tileables, **kw)
    124 
    125         # submit graph to local cluster
--> 126         self._api.submit_graph(self._session_id, json.dumps(graph.to_json(), separators=(',', ':')),
    127                                graph_key, targets, compose=compose)
    128 

~/Workspace/mars/mars/graph.pyx in mars.graph.DirectedGraph.to_json()

~/Workspace/mars/mars/serialize/core.pyx in mars.serialize.core.Serializable.to_json()

~/Workspace/mars/mars/serialize/core.pyx in mars.serialize.core.Serializable.serialize()

~/Workspace/mars/mars/serialize/core.pyx in mars.serialize.core.Provider.serialize_model()

~/Workspace/mars/mars/serialize/core.pyx in mars.serialize.core.Field.serialize()

~/Workspace/mars/mars/serialize/core.pyx in mars.serialize.core.Field.serialize()

~/Workspace/mars/mars/serialize/jsonserializer.pyx in mars.serialize.jsonserializer.JsonSerializeProvider.serialize_field()

~/Workspace/mars/mars/serialize/jsonserializer.pyx in mars.serialize.jsonserializer.JsonSerializeProvider._serialize_reference()

~/Workspace/mars/mars/serialize/core.pyx in mars.serialize.core.AttributeAsDict.serialize()

~/Workspace/mars/mars/serialize/core.pyx in mars.serialize.core.Provider.serialize_attribute_as_dict()

~/Workspace/mars/mars/serialize/core.pyx in mars.serialize.core.Provider.serialize_model()

~/Workspace/mars/mars/serialize/core.pyx in mars.serialize.core.Field.serialize()

~/Workspace/mars/mars/serialize/core.pyx in mars.serialize.core.Field.serialize()

~/Workspace/mars/mars/serialize/jsonserializer.pyx in mars.serialize.jsonserializer.JsonSerializeProvider.serialize_field()

~/Workspace/mars/mars/serialize/jsonserializer.pyx in mars.serialize.jsonserializer.JsonSerializeProvider.serialize_field()

~/Workspace/mars/mars/serialize/jsonserializer.pyx in mars.serialize.jsonserializer.JsonSerializeProvider._serialize_value()

~/Workspace/mars/mars/serialize/jsonserializer.pyx in mars.serialize.jsonserializer.JsonSerializeProvider._serialize_typed_value()

TypeError: Fail to serialize field `seeds` for LearnShuffle <key=4901c18065398f9e19eec455538bc65a>, reason: Expected tuple, got numpy.ndarray
```




Discussion:

Resolving patch:
diff --git a/mars/learn/utils/shuffle.py b/mars/learn/utils/shuffle.py
--- a/mars/learn/utils/shuffle.py
+++ b/mars/learn/utils/shuffle.py
@@ -131,26 +131,32 @@ def _safe_rechunk(tileable, ax_nsplit):
             return tileable
 
     @classmethod
-    def _calc_chunk_params(cls, in_chunk, axes, output, output_type, chunk_op):
+    def _calc_chunk_params(cls, in_chunk, axes, output, output_type,
+                           chunk_op, no_shuffle: bool):
         params = {'index': in_chunk.index}
         if output_type == OutputType.tensor:
             chunk_shape = list(in_chunk.shape)
             for ax in axes:
-                chunk_shape[ax] = np.nan
+                if not no_shuffle:
+                    chunk_shape[ax] = np.nan
             params['shape'] = tuple(chunk_shape)
             params['dtype'] = in_chunk.dtype
             params['order'] = output.order
         elif output_type == OutputType.dataframe:
             chunk_shape = list(in_chunk.shape)
             if 0 in axes:
-                chunk_shape[0] = np.nan
+                if not no_shuffle:
+                    chunk_shape[0] = np.nan
             params['shape'] = tuple(chunk_shape)
             params['dtypes'] = output.dtypes
             params['columns_value'] = output.columns_value
             params['index_value'] = _shuffle_index_value(chunk_op, in_chunk.index_value)
         else:
             assert output_type == OutputType.series
-            params['shape'] = (np.nan,)
+            if no_shuffle:
+                params['shape'] = in_chunk.shape
+            else:
+                params['shape'] = (np.nan,)
             params['name'] = in_chunk.name
             params['index_value'] = _shuffle_index_value(chunk_op, in_chunk.index_value)
             params['dtype'] = in_chunk.dtype
@@ -208,7 +214,7 @@ def tile(cls, op):
                 for c in inp.chunks:
                     chunk_op = LearnShuffle(axes=inp_axes, seeds=op.seeds[:len(inp_axes)],
                                             output_types=output_types)
-                    params = cls._calc_chunk_params(c, inp_axes, oup, output_type, chunk_op)
+                    params = cls._calc_chunk_params(c, inp_axes, oup, output_type, chunk_op, True)
                     out_chunk = chunk_op.new_chunk([c], kws=[params])
                     chunks.append(out_chunk)
                 out_chunks.append(chunks)
@@ -233,8 +239,8 @@ def tile(cls, op):
                     map_chunk_op = LearnShuffle(
                         stage=OperandStage.map,
                         output_types=output_types, axes=inp_axes,
-                        seeds=[mapper_seeds[j][in_chunk.index[ax]]
-                               for j, ax in enumerate(inp_axes)],
+                        seeds=tuple(mapper_seeds[j][in_chunk.index[ax]]
+                                    for j, ax in enumerate(inp_axes)),
                         reduce_sizes=reduce_sizes
                     )
                     map_chunk = map_chunk_op.new_chunk([in_chunk], **params)
@@ -249,11 +255,11 @@ def tile(cls, op):
                     chunk_op = LearnShuffle(
                         stage=OperandStage.reduce,
                         output_types=output_types, axes=reduce_axes,
-                        seeds=[reducer_seeds[j][c.index[ax]] for j, ax in enumerate(inp_axes)
-                               if reduce_sizes[j] > 1],
+                        seeds=tuple(reducer_seeds[j][c.index[ax]] for j, ax in enumerate(inp_axes)
+                                    if reduce_sizes[j] > 1),
                         reduce_sizes=reduce_sizes_,
                         shuffle_key=shuffle_key)
-                    params = cls._calc_chunk_params(c, inp_axes, oup, output_type, chunk_op)
+                    params = cls._calc_chunk_params(c, inp_axes, oup, output_type, chunk_op, False)
                     reduce_chunk = chunk_op.new_chunk([proxy_chunk], kws=[params])
                     reduce_chunks.append(reduce_chunk)
 
@@ -266,9 +272,10 @@ def tile(cls, op):
                     out_chunks.insert(i, inp.chunks)
                     out_nsplits.insert(i, inp.nsplits)
             assert len(out_chunks) == len(op.outputs)
-        for param, chunks, ns in zip(params, out_chunks, out_nsplits):
+        for i, param, chunks, ns in zip(itertools.count(), params, out_chunks, out_nsplits):
             param['chunks'] = chunks
             param['nsplits'] = ns
+            param['_position_'] = i
         return new_op.new_tileables(op.inputs, kws=params)
 
     @classmethod
@@ -410,7 +417,7 @@ def shuffle(*arrays, **options):
         if len(shapes) > 1:
             raise ValueError('arrays do not have same shape on axis {0}'.format(ax))
 
-    op = LearnShuffle(axes=axes, seeds=seeds,
+    op = LearnShuffle(axes=axes, seeds=tuple(seeds),
                       output_types=get_output_types(*arrays))
     shuffled_arrays = op(arrays)
     if len(arrays) == 1:


Test patch:
diff --git a/mars/learn/utils/tests/test_shuffle.py b/mars/learn/utils/tests/test_shuffle.py
--- a/mars/learn/utils/tests/test_shuffle.py
+++ b/mars/learn/utils/tests/test_shuffle.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import unittest
-
 import numpy as np
 import pandas as pd
 
@@ -22,9 +20,13 @@
 from mars.learn.utils import shuffle
 from mars.learn.utils.shuffle import LearnShuffle
 from mars.tiles import get_tiled
+from mars.tests.core import TestBase, ExecutorForTest
+
 
+class Test(TestBase):
+    def setUp(self):
+        self.executor = ExecutorForTest('numpy')
 
-class Test(unittest.TestCase):
     def testShuffleExpr(self):
         a = mt.random.rand(10, 3, chunk_size=2)
         b = md.DataFrame(mt.random.rand(10, 5), chunk_size=2)
@@ -100,7 +102,7 @@ def testShuffleExecution(self):
         ts2 = mt.array(s2, chunk_size=2)
 
         ret = shuffle(ts1, ts2, axes=[0, 1], random_state=0)
-        res1, res2 = ret.execute()
+        res1, res2 = self.executor.execute_tileables(ret)
 
         # calc row index
         s1_col_0 = s1[:, 0].tolist()
@@ -120,7 +122,7 @@ def testShuffleExecution(self):
 
         for axes in [(0,), (0, 1), (0, 2), (1, 2), (0, 1, 2)]:
             ret = shuffle(t1, t2, axes=axes, random_state=0)
-            res1, res2 = ret.execute()
+            res1, res2 = self.executor.execute_tileables(ret)
 
             self.assertEqual(res1.shape, raw1.shape)
             self.assertEqual(res2.shape, raw2.shape)
@@ -135,7 +137,7 @@ def testShuffleExecution(self):
 
         for axes in [(1,), (0, 1), (1, 2)]:
             ret = shuffle(t3, t4, axes=axes, random_state=0)
-            res3, res4 = ret.execute()
+            res3, res4 = self.executor.execute_tileables(ret)
 
             self.assertEqual(res3.shape, raw3.shape)
             self.assertEqual(res4.shape, raw4.shape)
@@ -152,7 +154,8 @@ def testShuffleExecution(self):
 
         for axes in [(0,), (1,), (0, 1), (1, 2), [0, 1, 2]]:
             ret = shuffle(t5, df, series, axes=axes, random_state=0)
-            res5, res_df, res_series = ret.execute()
+            # skip check nsplits because it's updated
+            res5, res_df, res_series = self.executor.execute_tileables(ret, check_nsplits=False)
 
             self.assertEqual(res5.shape, raw5.shape)
             self.assertEqual(res_df.shape, df.shape)