Repository: mars-project/mars

Issue:
[BUG] Tile of rechunk() outputs inconsistent dtypes and columns_value
**Describe the bug**
DataFrame.rechunk() produces chunks with different number of dtypes and columns_value, which may impair descendant chunks.

Code to reproduce this:

```python                                                                                                                                                                   
In [2]: import pandas as pd                                                                                                                                                                                          

In [4]: from mars.dataframe.datasource.dataframe import from_pandas as from_pandas_df 

In [6]: cols = [chr(ord('A') + i) for i in range(10)] 

In [8]: df_raw = pd.DataFrame(dict((c, [i ** 2 for i in range(20)]) for c in cols))                                                                                                                           

In [9]: df = from_pandas_df(df_raw, chunk_size=5)       

In [13]: rechunked = df.rechunk((20, 1)).tiles()                                                                                                                                                              

In [14]: rechunked.chunks[0].columns_value.to_pandas()                                                                                                                                                        
Out[14]: Index(['A'], dtype='object')

In [15]: rechunked.chunks[0].dtypes                                                                                                                                                                           
Out[15]: 
A    int64
B    int64
C    int64
D    int64
E    int64
dtype: object
```


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
@@ -14,6 +14,8 @@
 
 import itertools
 
+import pandas as pd
+
 from ... import opcodes as OperandDef
 from ...serialize import KeyField, AnyField, Int32Field, Int64Field
 from ...tensor.rechunk.core import get_nsplits, plan_rechunks, compute_rechunk_slices
@@ -92,7 +94,7 @@ def rechunk(a, chunk_size, threshold=None, chunk_size_limit=None):
     return op(a)
 
 
-def _concat_dataframe_index_and_columns(to_concat_chunks):
+def _concat_dataframe_meta(to_concat_chunks):
     if to_concat_chunks[0].index_value.to_pandas().empty:
         index_value = to_concat_chunks[0].index_value
     else:
@@ -101,7 +103,10 @@ def _concat_dataframe_index_and_columns(to_concat_chunks):
 
     idx_to_columns_value = dict((c.index[1], c.columns_value) for c in to_concat_chunks if c.index[0] == 0)
     columns_value = merge_index_value(idx_to_columns_value, store_data=True)
-    return index_value, columns_value
+
+    idx_to_dtypes = dict((c.index[1], c.dtypes) for c in to_concat_chunks if c.index[0] == 0)
+    dtypes = pd.concat([v[1] for v in list(sorted(idx_to_dtypes.items()))])
+    return index_value, columns_value, dtypes
 
 
 def _concat_series_index(to_concat_chunks):
@@ -137,7 +142,8 @@ def compute_rechunk(a, chunk_size):
                                                       object_type=ObjectType.dataframe)
                 merge_chunk = merge_chunk_op.new_chunk([old_chunk], shape=merge_chunk_shape,
                                                        index=merge_idx, index_value=new_index_value,
-                                                       columns_value=new_columns_value, dtypes=old_chunk.dtypes)
+                                                       columns_value=new_columns_value,
+                                                       dtypes=old_chunk.dtypes.iloc[chunk_slice[1]])
             else:
                 merge_chunk_op = SeriesIlocGetItem(chunk_slice, sparse=old_chunk.op.sparse,
                                                    object_type=ObjectType.series)
@@ -160,11 +166,11 @@ def compute_rechunk(a, chunk_size):
         else:
             if is_dataframe:
                 chunk_op = DataFrameConcat(object_type=ObjectType.dataframe)
-                index_value, columns_value = _concat_dataframe_index_and_columns(to_merge)
+                index_value, columns_value, dtypes = _concat_dataframe_meta(to_merge)
                 out_chunk = chunk_op.new_chunk(to_merge, shape=chunk_shape,
                                                index=idx, index_value=index_value,
                                                columns_value=columns_value,
-                                               dtypes=to_merge[0].dtypes)
+                                               dtypes=dtypes)
             else:
                 chunk_op = DataFrameConcat(object_type=ObjectType.series)
                 index_value = _concat_series_index(to_merge)


Test patch:
diff --git a/mars/dataframe/base/tests/test_base.py b/mars/dataframe/base/tests/test_base.py
--- a/mars/dataframe/base/tests/test_base.py
+++ b/mars/dataframe/base/tests/test_base.py
@@ -93,7 +93,8 @@ def testToCPU(self):
         self.assertIs(df2, to_cpu(df2))
 
     def testRechunk(self):
-        df = from_pandas_df(pd.DataFrame(np.random.rand(10, 10)), chunk_size=3)
+        raw = pd.DataFrame(np.random.rand(10, 10))
+        df = from_pandas_df(raw, chunk_size=3)
         df2 = df.rechunk(4).tiles()
 
         self.assertEqual(df2.shape, (10, 10))
@@ -102,19 +103,26 @@ def testRechunk(self):
         self.assertEqual(df2.chunks[0].shape, (4, 4))
         pd.testing.assert_index_equal(df2.chunks[0].index_value.to_pandas(), pd.RangeIndex(4))
         pd.testing.assert_index_equal(df2.chunks[0].columns_value.to_pandas(), pd.RangeIndex(4))
+        pd.testing.assert_series_equal(df2.chunks[0].dtypes, raw.dtypes[:4])
 
         self.assertEqual(df2.chunks[2].shape, (4, 2))
         pd.testing.assert_index_equal(df2.chunks[2].index_value.to_pandas(), pd.RangeIndex(4))
         pd.testing.assert_index_equal(df2.chunks[2].columns_value.to_pandas(), pd.RangeIndex(8, 10))
+        pd.testing.assert_series_equal(df2.chunks[2].dtypes, raw.dtypes[-2:])
 
         self.assertEqual(df2.chunks[-1].shape, (2, 2))
         pd.testing.assert_index_equal(df2.chunks[-1].index_value.to_pandas(), pd.RangeIndex(8, 10))
         pd.testing.assert_index_equal(df2.chunks[-1].columns_value.to_pandas(), pd.RangeIndex(8, 10))
+        pd.testing.assert_series_equal(df2.chunks[-1].dtypes, raw.dtypes[-2:])
+
+        for c in df2.chunks:
+            self.assertEqual(c.shape[1], len(c.dtypes))
+            self.assertEqual(len(c.columns_value.to_pandas()), len(c.dtypes))
 
         columns = [np.random.bytes(10) for _ in range(10)]
         index = np.random.randint(-100, 100, size=(4,))
-        data = pd.DataFrame(np.random.rand(4, 10), index=index, columns=columns)
-        df = from_pandas_df(data, chunk_size=3)
+        raw = pd.DataFrame(np.random.rand(4, 10), index=index, columns=columns)
+        df = from_pandas_df(raw, chunk_size=3)
         df2 = df.rechunk(6).tiles()
 
         self.assertEqual(df2.shape, (4, 10))
@@ -123,10 +131,16 @@ def testRechunk(self):
         self.assertEqual(df2.chunks[0].shape, (4, 6))
         pd.testing.assert_index_equal(df2.chunks[0].index_value.to_pandas(), df.index_value.to_pandas())
         pd.testing.assert_index_equal(df2.chunks[0].columns_value.to_pandas(), pd.Index(columns[:6]))
+        pd.testing.assert_series_equal(df2.chunks[0].dtypes, raw.dtypes[:6])
 
         self.assertEqual(df2.chunks[1].shape, (4, 4))
         pd.testing.assert_index_equal(df2.chunks[1].index_value.to_pandas(), df.index_value.to_pandas())
         pd.testing.assert_index_equal(df2.chunks[1].columns_value.to_pandas(), pd.Index(columns[6:]))
+        pd.testing.assert_series_equal(df2.chunks[1].dtypes, raw.dtypes[-4:])
+
+        for c in df2.chunks:
+            self.assertEqual(c.shape[1], len(c.dtypes))
+            self.assertEqual(len(c.columns_value.to_pandas()), len(c.dtypes))
 
         # test Series rechunk
         series = from_pandas_series(pd.Series(np.random.rand(10,)), chunk_size=3)