Repository: mars-project/mars

Issue:
Optimize df.merge by avoiding shuffles
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Is your feature request related to a problem? Please describe.**
Mars always need to shuffle all dataframes when joining them, however, if one of them has only one chunk, shuffle could be avoided by the strategy like the mapjoin in Hive.


Discussion:

Resolving patch:
diff --git a/mars/dataframe/merge/merge.py b/mars/dataframe/merge/merge.py
--- a/mars/dataframe/merge/merge.py
+++ b/mars/dataframe/merge/merge.py
@@ -208,6 +208,51 @@ def _gen_shuffle_chunks(cls, op, out_shape, shuffle_on, df):
                                     index_value=proxy_chunk.index_value, columns_value=proxy_chunk.columns_value))
         return reduce_chunks
 
+    @classmethod
+    def _tile_one_chunk(cls, op, left, right):
+        df = op.outputs[0]
+        if len(left.chunks) == 1 and len(right.chunks) == 1:
+            merge_op = op.copy().reset_key()
+            out_chunk = merge_op.new_chunk([left.chunks[0], right.chunks[0]],
+                                           shape=df.shape,
+                                           index=left.chunks[0].index,
+                                           index_value=df.index_value,
+                                           columns_value=df.columns_value)
+            out_chunks = [out_chunk]
+            nsplits = ((np.nan,), (df.shape[1],))
+        elif len(left.chunks) == 1:
+            out_chunks = []
+            left_chunk = left.chunks[0]
+            for c in right.chunks:
+                merge_op = op.copy().reset_key()
+                out_chunk = merge_op.new_chunk([left_chunk, c],
+                                               shape=(np.nan, df.shape[1]),
+                                               index=c.index,
+                                               index_value=infer_index_value(left_chunk.index_value,
+                                                                             c.index_value),
+                                               columns_value=df.columns_value)
+                out_chunks.append(out_chunk)
+            nsplits = ((np.nan,) * len(right.chunks), (df.shape[1],))
+        else:
+            out_chunks = []
+            right_chunk = right.chunks[0]
+            for c in left.chunks:
+                merge_op = op.copy().reset_key()
+                out_chunk = merge_op.new_chunk([c, right_chunk],
+                                               shape=(np.nan, df.shape[1]),
+                                               index=c.index,
+                                               index_value=infer_index_value(right_chunk.index_value,
+                                                                             c.index_value),
+                                               columns_value=df.columns_value)
+                out_chunks.append(out_chunk)
+            nsplits = ((np.nan,) * len(left.chunks), (df.shape[1],))
+
+        new_op = op.copy()
+        return new_op.new_dataframes(op.inputs, df.shape,
+                                     nsplits=nsplits,
+                                     chunks=out_chunks, dtypes=df.dtypes,
+                                     index_value=df.index_value, columns_value=df.columns_value)
+
     @classmethod
     def tile(cls, op):
         df = op.outputs[0]
@@ -222,6 +267,9 @@ def tile(cls, op):
             check_chunks_unknown_shape([right], TilesError)
             right = right.rechunk({1: right.shape[1]})._inplace_tile()
 
+        if len(left.chunks) == 1 or len(right.chunks) == 1:
+            return cls._tile_one_chunk(op, left, right)
+
         left_row_chunk_size = left.chunk_shape[0]
         right_row_chunk_size = right.chunk_shape[0]
         out_row_chunk_size = max(left_row_chunk_size, right_row_chunk_size)
@@ -241,7 +289,8 @@ def tile(cls, op):
             merge_op = op.copy().reset_key()
             out_chunk = merge_op.new_chunk([left_chunk, right_chunk], shape=(np.nan, df.shape[1]),
                                            index=left_chunk.index,
-                                           index_value=infer_index_value(left_chunk.index_value, right_chunk.index_value),
+                                           index_value=infer_index_value(left_chunk.index_value,
+                                                                         right_chunk.index_value),
                                            columns_value=df.columns_value)
             out_chunks.append(out_chunk)
 


Test patch:
diff --git a/mars/dataframe/merge/tests/test_merge.py b/mars/dataframe/merge/tests/test_merge.py
--- a/mars/dataframe/merge/tests/test_merge.py
+++ b/mars/dataframe/merge/tests/test_merge.py
@@ -17,6 +17,7 @@
 
 from mars.operands import OperandStage
 from mars.executor import Executor
+from mars.tiles import get_tiled
 from mars.tests.core import TestBase
 from mars.dataframe.datasource.dataframe import from_pandas
 from mars.dataframe.merge import DataFrameMergeAlign, DataFrameShuffleMerge
@@ -156,3 +157,44 @@ def testJoinOn(self):
                     self.assertEqual(rchunk.op.index_shuffle_size, 3)
                     self.assertEqual(rchunk.op.shuffle_on, None)
                 pd.testing.assert_index_equal(chunk.columns_value.to_pandas(), df.columns_value.to_pandas())
+
+    def testMergeOneChunk(self):
+        df1 = pd.DataFrame({'lkey': ['foo', 'bar', 'baz', 'foo'],
+                            'value': [1, 2, 3, 5]})
+        df2 = pd.DataFrame({'rkey': ['foo', 'bar', 'baz', 'foo'],
+                            'value': [5, 6, 7, 8]})
+
+        # all have one chunk
+        mdf1 = from_pandas(df1)
+        mdf2 = from_pandas(df2)
+        df = mdf1.merge(mdf2, left_on='lkey', right_on='rkey')
+        tiled = df.tiles()
+
+        self.assertEqual(tiled.chunk_shape, (1, 1))
+        self.assertEqual(tiled.chunks[0].inputs[0].key, get_tiled(mdf1).chunks[0].key)
+        self.assertEqual(tiled.chunks[0].inputs[1].key, get_tiled(mdf2).chunks[0].key)
+
+        # left has one chunk
+        mdf1 = from_pandas(df1)
+        mdf2 = from_pandas(df2, chunk_size=2)
+        df = mdf1.merge(mdf2, left_on='lkey', right_on='rkey')
+        tiled = df.tiles()
+
+        self.assertEqual(tiled.chunk_shape, (2, 1))
+        self.assertEqual(tiled.chunks[0].inputs[0].key, get_tiled(mdf1).chunks[0].key)
+        self.assertEqual(tiled.chunks[0].inputs[1].key, get_tiled(mdf2).chunks[0].key)
+        self.assertEqual(tiled.chunks[1].inputs[0].key, get_tiled(mdf1).chunks[0].key)
+        self.assertEqual(tiled.chunks[1].inputs[1].key, get_tiled(mdf2).chunks[1].key)
+
+        # right has one chunk
+        mdf1 = from_pandas(df1, chunk_size=2)
+        mdf2 = from_pandas(df2)
+        df = mdf1.merge(mdf2, left_on='lkey', right_on='rkey')
+        tiled = df.tiles()
+
+        self.assertEqual(tiled.chunk_shape, (2, 1))
+        self.assertEqual(tiled.chunks[0].inputs[0].key, get_tiled(mdf1).chunks[0].key)
+        self.assertEqual(tiled.chunks[0].inputs[1].key, get_tiled(mdf2).chunks[0].key)
+        self.assertEqual(tiled.chunks[1].inputs[0].key, get_tiled(mdf1).chunks[1].key)
+        self.assertEqual(tiled.chunks[1].inputs[1].key, get_tiled(mdf2).chunks[0].key)
+
diff --git a/mars/dataframe/merge/tests/test_merge_execution.py b/mars/dataframe/merge/tests/test_merge_execution.py
--- a/mars/dataframe/merge/tests/test_merge_execution.py
+++ b/mars/dataframe/merge/tests/test_merge_execution.py
@@ -233,3 +233,42 @@ def testJoinOn(self):
         expected4.set_index('a2', inplace=True)
         result4.set_index('a2', inplace=True)
         pd.testing.assert_frame_equal(sort_dataframe_inplace(expected4, 0), sort_dataframe_inplace(result4, 0))
+
+    def testMergeOneChunk(self):
+        df1 = pd.DataFrame({'lkey': ['foo', 'bar', 'baz', 'foo'],
+                            'value': [1, 2, 3, 5]}, index=['a1', 'a2', 'a3', 'a4'])
+        df2 = pd.DataFrame({'rkey': ['foo', 'bar', 'baz', 'foo'],
+                            'value': [5, 6, 7, 8]}, index=['a1', 'a2', 'a3', 'a4'])
+
+        # all have one chunk
+        mdf1 = from_pandas(df1)
+        mdf2 = from_pandas(df2)
+
+        expected = df1.merge(df2, left_on='lkey', right_on='rkey')
+        jdf = mdf1.merge(mdf2, left_on='lkey', right_on='rkey')
+        result = self.executor.execute_dataframe(jdf, concat=True)[0]
+
+        pd.testing.assert_frame_equal(expected.sort_values(by=expected.columns[1]).reset_index(drop=True),
+                                      result.sort_values(by=result.columns[1]).reset_index(drop=True))
+
+        # left have one chunk
+        mdf1 = from_pandas(df1)
+        mdf2 = from_pandas(df2, chunk_size=2)
+
+        expected = df1.merge(df2, left_on='lkey', right_on='rkey')
+        jdf = mdf1.merge(mdf2, left_on='lkey', right_on='rkey')
+        result = self.executor.execute_dataframe(jdf, concat=True)[0]
+
+        pd.testing.assert_frame_equal(expected.sort_values(by=expected.columns[1]).reset_index(drop=True),
+                                      result.sort_values(by=result.columns[1]).reset_index(drop=True))
+
+        # right have one chunk
+        mdf1 = from_pandas(df1, chunk_size=3)
+        mdf2 = from_pandas(df2)
+
+        expected = df1.merge(df2, left_on='lkey', right_on='rkey')
+        jdf = mdf1.merge(mdf2, left_on='lkey', right_on='rkey')
+        result = self.executor.execute_dataframe(jdf, concat=True)[0]
+
+        pd.testing.assert_frame_equal(expected.sort_values(by=expected.columns[1]).reset_index(drop=True),
+                                      result.sort_values(by=result.columns[1]).reset_index(drop=True))