Repository: mars-project/mars

Issue:
[BUG] DataFrame.add is not working
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
A clear and concise description of what the bug is.

`DataFrame.add` is not working when `chunk_size` less than data size.

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version: 3.7.3
2. The version of Mars you use: master(commit ec9398768)
3. Versions of crucial packages, such as numpy, scipy and protobuf
4. Full stack of the error.

```
Traceback (most recent call last):
  File "/Users/hetao/mars/mars/tiles.py", line 111, in _dispatch
    handler = self._handlers[op_cls]
KeyError: <class 'mars.dataframe.expressions.arithmetic.add.DataFrameAdd'>

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/hetao/mars/mars/core.py", line 554, in execute
    return session.run(self, **kw)
  File "/Users/hetao/mars/mars/session.py", line 115, in run
    result = self._sess.run(*tileables, **kw)
  File "/Users/hetao/mars/mars/session.py", line 56, in run
    res = self._executor.execute_tileables(tileables, **kw)
  File "/Users/hetao/mars/mars/utils.py", line 364, in _wrapped
    return func(*args, **kwargs)
  File "/Users/hetao/mars/mars/executor.py", line 499, in execute_tileables
    tileable.tiles()
  File "/Users/hetao/mars/mars/core.py", line 474, in tiles
    return handler.tiles(self)
  File "/Users/hetao/mars/mars/tiles.py", line 174, in tiles
    tiled = self._dispatch(node.op)
  File "/Users/hetao/mars/mars/utils.py", line 364, in _wrapped
    return func(*args, **kwargs)
  File "/Users/hetao/mars/mars/tiles.py", line 116, in _dispatch
    return op_cls.tile(op)
  File "/Users/hetao/mars/mars/dataframe/expressions/arithmetic/core.py", line 527, in tile
    return cls._tile_both_dataframes(op)
  File "/Users/hetao/mars/mars/dataframe/expressions/arithmetic/core.py", line 516, in _tile_both_dataframes
    out_chunks = cls._gen_out_chunks_with_all_shuffle(op, out_shape, left, right)
  File "/Users/hetao/mars/mars/dataframe/expressions/arithmetic/core.py", line 461, in _gen_out_chunks_with_all_shuffle
    reduce_op.new_chunk([proxy_chunk], shape=(np.nan, np.nan), index=out_idx))
  File "/Users/hetao/mars/mars/core.py", line 650, in new_chunk
    return self.new_chunks(inputs, kws=kws, **kw)[0]
  File "/Users/hetao/mars/mars/core.py", line 644, in new_chunks
    return self._new_chunks(inputs, kws=kws, **kwargs)
  File "/Users/hetao/mars/mars/core.py", line 618, in _new_chunks
    chunk = self._create_chunk(j, index, **create_chunk_kw)
  File "/Users/hetao/mars/mars/dataframe/expressions/arithmetic/core.py", line 188, in _create_chunk
    all_dtypes = [c.op.column_shuffle_segments[index[1]] for c in index_align_map_chunks
  File "/Users/hetao/mars/mars/dataframe/expressions/arithmetic/core.py", line 189, in <listcomp>
    if c.index[0] == index_align_map_chunks[0].index[0]]
IndexError: list index out of range
>>>
```

5. Minimized code to reproduce the error.

```python
>>> import pandas as pd
>>> import numpy as np
>>>
>>> data = pd.DataFrame(np.random.rand(10, 10), index=np.random.randint(-100, 100, size=(10,)),
...                             columns=[np.random.bytes(10) for _ in range(10)])
>>>
...
>>> from mars.dataframe.expressions.datasource.dataframe import from_pandas

>>>
>>> from mars.dataframe.expressions.arithmetic import add
>>> df = from_pandas(data, chunk_size=3)
>>> add(df, df).execute()
```

**Expected behavior**
A clear and concise description of what you expected to happen.

Return the result of `DataFrame.add`.

**Additional context**
Add any other context about the problem here.



Discussion:

Resolving patch:
diff --git a/mars/dataframe/expressions/arithmetic/core.py b/mars/dataframe/expressions/arithmetic/core.py
--- a/mars/dataframe/expressions/arithmetic/core.py
+++ b/mars/dataframe/expressions/arithmetic/core.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import itertools
+import copy
 
 import numpy as np
 try:
@@ -310,9 +311,6 @@ def _check_overlap(cls, chunk_index_min_max):
     @classmethod
     def _get_chunk_index_min_max(cls, df, index_type, axis):
         index = getattr(df, index_type)
-        if not index.is_monotonic_increasing_or_decreasing and \
-                df.chunk_shape[axis] > 1:
-            return
 
         chunk_index_min_max = []
         for i in range(df.chunk_shape[axis]):
@@ -468,6 +466,33 @@ def _gen_out_chunks_with_all_shuffle(cls, op, out_shape, left, right):
 
         return out_chunks
 
+    @classmethod
+    def _is_index_identical(cls, left, right, index_type, axis):
+        if left.chunk_shape[axis] != right.chunk_shape[axis]:
+            return False
+
+        for i in range(left.chunk_shape[axis]):
+            idx = [0, 0]
+            idx[axis] = i
+            if getattr(left.cix[tuple(idx)], index_type).key != \
+                    getattr(right.cix[tuple(idx)], index_type).key:
+                return False
+
+        return True
+
+    @classmethod
+    def _need_shuffle_on_axis(cls, left, right, index_type, axis):
+        if cls._is_index_identical(left, right, index_type, axis):
+            return False
+
+        for df in (left, right):
+            index = getattr(df, index_type)
+            if not index.is_monotonic_increasing_or_decreasing and \
+                    df.chunk_shape[axis] > 1:
+                return True
+
+        return False
+
     @classmethod
     def _tile_both_dataframes(cls, op):
         # if both of the inputs are DataFrames, axis is just ignored
@@ -479,10 +504,9 @@ def _tile_both_dataframes(cls, op):
         # first, we decide the chunk size on each axis
         # we perform the same logic for both index and columns
         for axis, index_type in enumerate(['index_value', 'columns']):
-            # if both of the indexes are monotonic increasing or decreasing
-            left_chunk_index_min_max = cls._get_chunk_index_min_max(left, index_type, axis)
-            right_chunk_index_min_max = cls._get_chunk_index_min_max(right, index_type, axis)
-            if left_chunk_index_min_max is not None and right_chunk_index_min_max is not None:
+            if not cls._need_shuffle_on_axis(left, right, index_type, axis):
+                left_chunk_index_min_max = cls._get_chunk_index_min_max(left, index_type, axis)
+                right_chunk_index_min_max = cls._get_chunk_index_min_max(right, index_type, axis)
                 # no need to do shuffle on this axis
                 if len(left_chunk_index_min_max[0]) == 1 and len(right_chunk_index_min_max[0]) == 1:
                     # both left and right has only 1 chunk
@@ -536,12 +560,23 @@ def _operator(self):
     def _calc_properties(cls, x1, x2):
         dtypes = columns = index = None
         index_shape = column_shape = np.nan
-        if x1.dtypes is not None and x2.dtypes is not None:
+
+        if x1.columns.key == x2.columns.key:
+            dtypes = x1.dtypes
+            column_shape = len(dtypes)
+            columns = copy.copy(x1.columns)
+            columns.value.should_be_monotonic = True
+        elif x1.dtypes is not None and x2.dtypes is not None:
             dtypes = infer_dtypes(x1.dtypes, x2.dtypes, cls._operator)
             column_shape = len(dtypes)
             columns = parse_index(dtypes.index, store_data=True)
             columns.value.should_be_monotonic = True
-        if x1.index_value is not None and x2.index_value is not None:
+
+        if x1.index_value.key == x2.index_value.key:
+            index = copy.copy(x1.index_value)
+            index.value.should_be_monotonic = True
+            index_shape = x1.shape[0]
+        elif x1.index_value is not None and x2.index_value is not None:
             index = infer_index_value(x1.index_value, x2.index_value, cls._operator)
             index.value.should_be_monotonic = True
             if index.key == x1.index_value.key == x2.index_value.key and \
@@ -577,6 +612,7 @@ def __call__(self, x1, x2):
     def rcall(self, x1, x2):
         return self._call(x1, x2)
 
+
 class DataFrameUnaryOpMixin(DataFrameOperandMixin):
     __slots__ = ()
 
diff --git a/mars/dataframe/utils.py b/mars/dataframe/utils.py
--- a/mars/dataframe/utils.py
+++ b/mars/dataframe/utils.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import functools
-import operator
 
 from mars.lib.mmh3 import hash as mmh_hash
 
@@ -23,9 +22,8 @@ def func(x, size):
         return mmh_hash(bytes(x)) % size
 
     f = functools.partial(func, size=size)
-    grouped = sorted(index.groupby(index.map(f)).items(),
-                     key=operator.itemgetter(0))
-    return [g[1] for g in grouped]
+    idx_to_grouped = dict(index.groupby(index.map(f)).items())
+    return [idx_to_grouped.get(i, list()) for i in range(size)]
 
 
 def hash_dtypes(dtypes, size):


Test patch:
diff --git a/mars/dataframe/expressions/tests/test_arithmetic.py b/mars/dataframe/expressions/tests/test_arithmetic.py
--- a/mars/dataframe/expressions/tests/test_arithmetic.py
+++ b/mars/dataframe/expressions/tests/test_arithmetic.py
@@ -352,6 +352,78 @@ def testAddWithAllShuffle(self):
 
         self.assertEqual(len(proxy_keys), 2)
 
+        data4 = pd.DataFrame(np.random.rand(10, 10), index=np.random.randint(-100, 100, size=(10,)),
+                            columns=[np.random.bytes(10) for _ in range(10)])
+        df4 = from_pandas(data4, chunk_size=3)
+
+        data5 = pd.DataFrame(np.random.rand(10, 10), index=np.random.randint(-100, 100, size=(10,)),
+                             columns=[np.random.bytes(10) for _ in range(10)])
+        df5 = from_pandas(data5, chunk_size=3)
+
+        df6 = add(df4, df5)
+
+        # test df6's index and columns
+        pd.testing.assert_index_equal(df6.columns.to_pandas(), (data4 + data5).columns)
+        self.assertTrue(df6.columns.should_be_monotonic)
+        self.assertIsInstance(df6.index_value.value, IndexValue.Int64Index)
+        self.assertTrue(df6.index_value.should_be_monotonic)
+        pd.testing.assert_index_equal(df6.index_value.to_pandas(), pd.Int64Index([]))
+        self.assertNotEqual(df6.index_value.key, df4.index_value.key)
+        self.assertNotEqual(df6.index_value.key, df5.index_value.key)
+        self.assertEqual(df6.shape[1], 20)  # columns is recorded, so we can get it
+
+        df6.tiles()
+
+        self.assertEqual(df6.chunk_shape, (4, 4))
+        proxy_keys = set()
+        for c in df6.chunks:
+            self.assertIsInstance(c.op, DataFrameAdd)
+            self.assertEqual(len(c.inputs), 2)
+            # test left side
+            self.assertIsInstance(c.inputs[0].op, DataFrameIndexAlignReduce)
+            expect_dtypes = pd.concat([hash_dtypes(ic.inputs[0].op.data.dtypes, 4)[c.index[1]]
+                                       for ic in c.inputs[0].inputs[0].inputs if ic.index[0] == 0])
+            pd.testing.assert_series_equal(c.inputs[0].dtypes, expect_dtypes)
+            pd.testing.assert_index_equal(c.inputs[0].columns.to_pandas(), c.inputs[0].dtypes.index)
+            self.assertIsInstance(c.inputs[0].index_value.to_pandas(), type(data1.index))
+            self.assertIsInstance(c.inputs[0].inputs[0].op, DataFrameShuffleProxy)
+            proxy_keys.add(c.inputs[0].inputs[0].op.key)
+            for ic, ci in zip(c.inputs[0].inputs[0].inputs, df4.chunks):
+                self.assertIsInstance(ic.op, DataFrameIndexAlignMap)
+                self.assertEqual(ic.op.index_shuffle_size, 4)
+                self.assertIsInstance(ic.index_value.to_pandas(), type(data1.index))
+                self.assertEqual(ic.op.column_shuffle_size, 4)
+                self.assertIsNotNone(ic.columns)
+                shuffle_segments = ic.op.column_shuffle_segments
+                expected_shuffle_segments = hash_dtypes(ci.data.dtypes, 4)
+                self.assertEqual(len(shuffle_segments), len(expected_shuffle_segments))
+                for ss, ess in zip(shuffle_segments, expected_shuffle_segments):
+                    pd.testing.assert_series_equal(ss, ess)
+                self.assertIs(ic.inputs[0], ci.data)
+            # test right side
+            self.assertIsInstance(c.inputs[1].op, DataFrameIndexAlignReduce)
+            expect_dtypes = pd.concat([hash_dtypes(ic.inputs[0].op.data.dtypes, 4)[c.index[1]]
+                                       for ic in c.inputs[1].inputs[0].inputs if ic.index[0] == 0])
+            pd.testing.assert_series_equal(c.inputs[1].dtypes, expect_dtypes)
+            pd.testing.assert_index_equal(c.inputs[1].columns.to_pandas(), c.inputs[1].dtypes.index)
+            self.assertIsInstance(c.inputs[0].index_value.to_pandas(), type(data1.index))
+            self.assertIsInstance(c.inputs[1].inputs[0].op, DataFrameShuffleProxy)
+            proxy_keys.add(c.inputs[1].inputs[0].op.key)
+            for ic, ci in zip(c.inputs[1].inputs[0].inputs, df5.chunks):
+                self.assertIsInstance(ic.op, DataFrameIndexAlignMap)
+                self.assertEqual(ic.op.index_shuffle_size, 4)
+                self.assertIsInstance(ic.index_value.to_pandas(), type(data1.index))
+                self.assertEqual(ic.op.column_shuffle_size, 4)
+                self.assertIsNotNone(ic.columns)
+                shuffle_segments = ic.op.column_shuffle_segments
+                expected_shuffle_segments = hash_dtypes(ci.data.dtypes, 4)
+                self.assertEqual(len(shuffle_segments), len(expected_shuffle_segments))
+                for ss, ess in zip(shuffle_segments, expected_shuffle_segments):
+                    pd.testing.assert_series_equal(ss, ess)
+                self.assertIs(ic.inputs[0], ci.data)
+
+        self.assertEqual(len(proxy_keys), 2)
+
     def testWithoutShuffleAndWithOneChunk(self):
         # only 1 axis is monotonic
         # data1 with index split into [0...4], [5...9],
@@ -533,6 +605,33 @@ def testWithShuffleAndOneChunk(self):
 
         self.assertEqual(len(proxy_keys), 2)
 
+    def testAddSelf(self):
+        data = pd.DataFrame(np.random.rand(10, 10), index=np.random.randint(-100, 100, size=(10,)),
+                            columns=[np.random.bytes(10) for _ in range(10)])
+        df = from_pandas(data, chunk_size=3)
+        df2 = add(df, df)
+
+        # test df2's index and columns
+        pd.testing.assert_index_equal(df2.columns.to_pandas(), (data + data).columns)
+        self.assertTrue(df2.columns.should_be_monotonic)
+        self.assertIsInstance(df2.index_value.value, IndexValue.Int64Index)
+        self.assertTrue(df2.index_value.should_be_monotonic)
+        pd.testing.assert_index_equal(df2.index_value.to_pandas(), pd.Int64Index([]))
+        self.assertEqual(df2.index_value.key, df.index_value.key)
+        self.assertEqual(df2.columns.key, df.columns.key)
+        self.assertEqual(df2.shape[1], 10)
+
+        df2.tiles()
+
+        self.assertEqual(df2.chunk_shape, df.chunk_shape)
+        for c in df2.chunks:
+            self.assertIsInstance(c.op, DataFrameAdd)
+            self.assertEqual(len(c.inputs), 2)
+            # test the left side
+            self.assertIs(c.inputs[0], df.cix[c.index].data)
+            # test the right side
+            self.assertIs(c.inputs[1], df.cix[c.index].data)
+
     def testAbs(self):
         data1 = pd.DataFrame(np.random.rand(10, 10), index=[0, 10, 2, 3, 4, 5, 6, 7, 8, 9],
                              columns=[4, 1, 3, 2, 10, 5, 9, 8, 6, 7])