Repository: databricks/koalas

Issue:
DataFrame.merge issues with specifying both joins
The `.merge` function does not seem to work when the index is set through `set_index`. Here are two examples that works with pandas but fail with Koalas:

```py
df1 = pd.DataFrame({"a":[0], "x": [-1]})
df2 = pd.DataFrame({"b":[0], "y": [-1]})
df1.set_index("a").merge(df2.set_index("b"), left_index=True, right_index=True)
df1.merge(df2, left_on=["a"], right_on=["b"])
```



Discussion:
I'll take this.

Resolving patch:
diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
--- a/databricks/koalas/frame.py
+++ b/databricks/koalas/frame.py
@@ -2275,7 +2275,10 @@ def shape(self):
         """
         return len(self), len(self.columns)
 
-    def merge(self, right: 'DataFrame', how: str = 'inner', on: str = None,
+    def merge(left, right: 'DataFrame', how: str = 'inner',
+              on: Optional[Union[str, List[str]]] = None,
+              left_on: Optional[Union[str, List[str]]] = None,
+              right_on: Optional[Union[str, List[str]]] = None,
               left_index: bool = False, right_index: bool = False,
               suffixes: Tuple[str, str] = ('_x', '_y')) -> 'DataFrame':
         """
@@ -2298,6 +2301,12 @@ def merge(self, right: 'DataFrame', how: str = 'inner', on: str = None,
         on: Column or index level names to join on. These must be found in both DataFrames. If on
             is None and not merging on indexes then this defaults to the intersection of the
             columns in both DataFrames.
+        left_on: Column or index level names to join on in the left DataFrame. Can also
+            be an array or list of arrays of the length of the left DataFrame.
+            These arrays are treated as if they are columns.
+        right_on: Column or index level names to join on in the right DataFrame. Can also
+            be an array or list of arrays of the length of the right DataFrame.
+            These arrays are treated as if they are columns.
         left_index: Use the index from the left DataFrame as the join key(s). If it is a
             MultiIndex, the number of keys in the other DataFrame (either the index or a number of
             columns) must match the number of levels.
@@ -2313,6 +2322,38 @@ def merge(self, right: 'DataFrame', how: str = 'inner', on: str = None,
 
         Examples
         --------
+        >>> df1 = ks.DataFrame({'lkey': ['foo', 'bar', 'baz', 'foo'],
+        ...                     'value': [1, 2, 3, 5]},
+        ...                    columns=['lkey', 'value'])
+        >>> df2 = ks.DataFrame({'rkey': ['foo', 'bar', 'baz', 'foo'],
+        ...                     'value': [5, 6, 7, 8]},
+        ...                    columns=['rkey', 'value'])
+        >>> df1
+          lkey  value
+        0  foo      1
+        1  bar      2
+        2  baz      3
+        3  foo      5
+        >>> df2
+          rkey  value
+        0  foo      5
+        1  bar      6
+        2  baz      7
+        3  foo      8
+
+        Merge df1 and df2 on the lkey and rkey columns. The value columns have
+        the default suffixes, _x and _y, appended.
+
+        >>> merged = df1.merge(df2, left_on='lkey', right_on='rkey')
+        >>> merged.sort_values(by=['lkey', 'value_x', 'rkey', 'value_y'])
+          lkey  value_x rkey  value_y
+        0  bar        2  bar        6
+        1  baz        3  baz        7
+        2  foo        1  foo        5
+        3  foo        1  foo        8
+        4  foo        5  foo        5
+        5  foo        5  foo        8
+
         >>> left_kdf = ks.DataFrame({'A': [1, 2]})
         >>> right_kdf = ks.DataFrame({'B': ['x', 'y']}, index=[1, 2])
 
@@ -2341,10 +2382,39 @@ def merge(self, right: 'DataFrame', how: str = 'inner', on: str = None,
         As described in #263, joining string columns currently returns None for missing values
             instead of NaN.
         """
-        if on is None and not left_index and not right_index:
-            raise ValueError("At least 'on' or 'left_index' and 'right_index' have to be set")
-        if on is not None and (left_index or right_index):
-            raise ValueError("Only 'on' or 'left_index' and 'right_index' can be set")
+        _to_list = lambda o: o if o is None or is_list_like(o) else [o]
+
+        if on:
+            if left_on or right_on:
+                raise ValueError('Can only pass argument "on" OR "left_on" and "right_on", '
+                                 'not a combination of both.')
+            left_keys = _to_list(on)
+            right_keys = _to_list(on)
+        else:
+            # TODO: need special handling for multi-index.
+            if left_index:
+                left_keys = left._metadata.index_columns
+            else:
+                left_keys = _to_list(left_on)
+            if right_index:
+                right_keys = right._metadata.index_columns
+            else:
+                right_keys = _to_list(right_on)
+
+            if left_keys and not right_keys:
+                raise ValueError('Must pass right_on or right_index=True')
+            if right_keys and not left_keys:
+                raise ValueError('Must pass left_on or left_index=True')
+            if not left_keys and not right_keys:
+                common = list(left.columns.intersection(right.columns))
+                if len(common) == 0:
+                    raise ValueError(
+                        'No common columns to perform merge on. Merge options: '
+                        'left_on=None, right_on=None, left_index=False, right_index=False')
+                left_keys = common
+                right_keys = common
+            if len(left_keys) != len(right_keys):  # type: ignore
+                raise ValueError('len(left_keys) must equal len(right_keys)')
 
         if how == 'full':
             warnings.warn("Warning: While Koalas will accept 'full', you should use 'outer' " +
@@ -2356,41 +2426,55 @@ def merge(self, right: 'DataFrame', how: str = 'inner', on: str = None,
             raise ValueError("The 'how' parameter has to be amongst the following values: ",
                              "['inner', 'left', 'right', 'outer']")
 
-        if on is None:
-            # FIXME Move index string to constant?
-            on = '__index_level_0__'
-
-        left_table = self._sdf.alias('left_table')
+        left_table = left._sdf.alias('left_table')
         right_table = right._sdf.alias('right_table')
 
+        left_key_columns = [left_table[col] for col in left_keys]  # type: ignore
+        right_key_columns = [right_table[col] for col in right_keys]  # type: ignore
+
+        join_condition = reduce(lambda x, y: x & y,
+                                [lkey == rkey for lkey, rkey
+                                 in zip(left_key_columns, right_key_columns)])
+
+        joined_table = left_table.join(right_table, join_condition, how=how)
+
         # Unpack suffixes tuple for convenience
         left_suffix = suffixes[0]
         right_suffix = suffixes[1]
 
         # Append suffixes to columns with the same name to avoid conflicts later
-        duplicate_columns = list(self.columns & right.columns)
-        if duplicate_columns:
-            for duplicate_column_name in duplicate_columns:
-                left_table = left_table.withColumnRenamed(duplicate_column_name,
-                                                          duplicate_column_name + left_suffix)
-                right_table = right_table.withColumnRenamed(duplicate_column_name,
-                                                            duplicate_column_name + right_suffix)
-
-        join_condition = (left_table[on] == right_table[on] if on not in duplicate_columns
-                          else left_table[on + left_suffix] == right_table[on + right_suffix])
-        joined_table = left_table.join(right_table, join_condition, how=how)
+        duplicate_columns = (set(left._metadata.data_columns)
+                             & set(right._metadata.data_columns))
 
-        if on in duplicate_columns:
-            # Merge duplicate key columns
-            joined_table = joined_table.withColumnRenamed(on + left_suffix, on)
-            joined_table = joined_table.drop(on + right_suffix)
+        left_index_columns = set(left._metadata.index_columns)
+        right_index_columns = set(right._metadata.index_columns)
 
-        # Remove auxiliary index
-        # FIXME Move index string to constant?
-        joined_table = joined_table.drop('__index_level_0__')
+        # TODO: in some case, we can keep indexes.
+        exprs = []
+        for col in left_table.columns:
+            if col in left_index_columns:
+                continue
+            scol = left_table[col]
+            if col in duplicate_columns:
+                if col in left_keys and col in right_keys:
+                    pass
+                else:
+                    col = col + left_suffix
+                    scol = scol.alias(col)
+            exprs.append(scol)
+        for col in right_table.columns:
+            if col in right_index_columns:
+                continue
+            scol = right_table[col]
+            if col in duplicate_columns:
+                if col in left_keys and col in right_keys:
+                    continue
+                else:
+                    col = col + right_suffix
+                    scol = scol.alias(col)
+            exprs.append(scol)
 
-        kdf = DataFrame(joined_table)
-        return kdf
+        return DataFrame(joined_table.select(*exprs))
 
     def sample(self, n: Optional[int] = None, frac: Optional[float] = None, replace: bool = False,
                random_state: Optional[int] = None) -> 'DataFrame':


Test patch:
diff --git a/databricks/koalas/tests/test_dataframe.py b/databricks/koalas/tests/test_dataframe.py
--- a/databricks/koalas/tests/test_dataframe.py
+++ b/databricks/koalas/tests/test_dataframe.py
@@ -18,6 +18,7 @@
 
 import numpy as np
 import pandas as pd
+from pyspark.sql.utils import AnalysisException
 
 from databricks import koalas as ks
 from databricks.koalas.testing.utils import ReusedSQLTestCase, SQLTestUtils
@@ -401,58 +402,105 @@ def test_isin(self):
             kdf.isin(1)
 
     def test_merge(self):
-        left_kdf = ks.DataFrame({'A': [1, 2]})
-        right_kdf = ks.DataFrame({'B': ['x', 'y']}, index=[1, 2])
-
-        # Assert only 'on' or 'left_index' and 'right_index' parameters are set
-        msg = "At least 'on' or 'left_index' and 'right_index' have to be set"
-        with self.assertRaises(ValueError, msg=msg):
-            left_kdf.merge(right_kdf)
-        msg = "Only 'on' or 'left_index' and 'right_index' can be set"
-        with self.assertRaises(ValueError, msg=msg):
-            left_kdf.merge(right_kdf, on='id', left_index=True)
-
-        # Assert a valid option for the 'how' parameter is used
-        msg = ("The 'how' parameter has to be amongst the following values: ['inner', 'left', " +
-               "'right', 'full', 'outer']")
-        with self.assertRaises(ValueError, msg=msg):
-            left_kdf.merge(right_kdf, how='foo', left_index=True, right_index=True)
-
-        # Assert inner join
-        res = left_kdf.merge(right_kdf, left_index=True, right_index=True)
-        self.assert_eq(res, pd.DataFrame({'A': [2], 'B': ['x']}))
-
-        # Assert inner join on non-default column
-        left_kdf_with_id = ks.DataFrame({'A': [1, 2], 'id': [0, 1]})
-        right_kdf_with_id = ks.DataFrame({'B': ['x', 'y'], 'id': [0, 1]}, index=[1, 2])
-        res = left_kdf_with_id.merge(right_kdf_with_id, on='id')
-        # Explicitly set columns to also assure their correct order with Python 3.5
-        self.assert_eq(res, pd.DataFrame({'A': [1, 2], 'id': [0, 1], 'B': ['x', 'y']},
-                                         columns=['A', 'id', 'B']))
-
-        # Assert left join
-        res = left_kdf.merge(right_kdf, left_index=True, right_index=True, how='left')
-        # FIXME Replace None with np.nan once #263 is solved
-        self.assert_eq(res, pd.DataFrame({'A': [1, 2], 'B': [None, 'x']}))
-
-        # Assert right join
-        res = left_kdf.merge(right_kdf, left_index=True, right_index=True, how='right')
-        self.assert_eq(res, pd.DataFrame({'A': [2, np.nan], 'B': ['x', 'y']}))
-
-        # Assert full outer join
-        res = left_kdf.merge(right_kdf, left_index=True, right_index=True, how='outer')
-        # FIXME Replace None with np.nan once #263 is solved
-        self.assert_eq(res, pd.DataFrame({'A': [1, 2, np.nan], 'B': [None, 'x', 'y']}))
-
-        # Assert full outer join also works with 'full' keyword
-        res = left_kdf.merge(right_kdf, left_index=True, right_index=True, how='full')
-        # FIXME Replace None with np.nan once #263 is solved
-        self.assert_eq(res, pd.DataFrame({'A': [1, 2, np.nan], 'B': [None, 'x', 'y']}))
-
-        # Assert suffixes create the expected column names
-        res = left_kdf.merge(ks.DataFrame({'A': [3, 4]}), left_index=True, right_index=True,
-                             suffixes=('_left', '_right'))
-        self.assert_eq(res, pd.DataFrame({'A_left': [1, 2], 'A_right': [3, 4]}))
+        left_pdf = pd.DataFrame({'lkey': ['foo', 'bar', 'baz', 'foo', 'bar', 'l'],
+                                 'value': [1, 2, 3, 5, 6, 7],
+                                 'x': list('abcdef')},
+                                columns=['lkey', 'value', 'x'])
+        right_pdf = pd.DataFrame({'rkey': ['baz', 'foo', 'bar', 'baz', 'foo', 'r'],
+                                  'value': [4, 5, 6, 7, 8, 9],
+                                  'y': list('efghij')},
+                                 columns=['rkey', 'value', 'y'])
+
+        left_kdf = ks.from_pandas(left_pdf)
+        right_kdf = ks.from_pandas(right_pdf)
+
+        def check(op):
+            k_res = op(left_kdf, right_kdf)
+            k_res = k_res.sort_values(by=list(k_res.columns))
+            p_res = op(left_pdf, right_pdf)
+            p_res = p_res.sort_values(by=list(p_res.columns))
+            p_res = p_res.reset_index(drop=True)
+            self.assert_eq(k_res, p_res)
+
+        check(lambda left, right: left.merge(right))
+        check(lambda left, right: left.merge(right, on='value'))
+        check(lambda left, right: left.merge(right, left_on='lkey', right_on='rkey'))
+        check(lambda left, right: left.set_index('lkey').merge(right.set_index('rkey')))
+        check(lambda left, right: left.set_index('lkey').merge(right,
+                                                               left_index=True, right_on='rkey'))
+        check(lambda left, right: left.merge(right.set_index('rkey'),
+                                             left_on='lkey', right_index=True))
+        check(lambda left, right: left.set_index('lkey').merge(right.set_index('rkey'),
+                                                               left_index=True, right_index=True))
+
+        # MultiIndex
+        check(lambda left, right: left.merge(right,
+                                             left_on=['lkey', 'value'], right_on=['rkey', 'value']))
+        check(lambda left, right: left.set_index(['lkey', 'value'])
+              .merge(right, left_index=True, right_on=['rkey', 'value']))
+        check(lambda left, right: left.merge(
+            right.set_index(['rkey', 'value']), left_on=['lkey', 'value'], right_index=True))
+        # TODO: when both left_index=True and right_index=True with multi-index
+        # check(lambda left, right: left.set_index(['lkey', 'value']).merge(
+        #     right.set_index(['rkey', 'value']), left_index=True, right_index=True))
+
+        # join types
+        for how in ['inner', 'left', 'right', 'outer']:
+            check(lambda left, right: left.merge(right, left_on='lkey', right_on='rkey', how=how))
+
+        # suffix
+        check(lambda left, right: left.merge(right, left_on='lkey', right_on='rkey',
+                                             suffixes=['_left', '_right']))
+
+    def test_merge_raises(self):
+        left = ks.DataFrame({'value': [1, 2, 3, 5, 6],
+                             'x': list('abcde')},
+                            columns=['value', 'x'],
+                            index=['foo', 'bar', 'baz', 'foo', 'bar'])
+        right = ks.DataFrame({'value': [4, 5, 6, 7, 8],
+                              'y': list('fghij')},
+                             columns=['value', 'y'],
+                             index=['baz', 'foo', 'bar', 'baz', 'foo'])
+
+        with self.assertRaisesRegex(ValueError,
+                                    'No common columns to perform merge on'):
+            left[['x']].merge(right[['y']])
+
+        with self.assertRaisesRegex(ValueError,
+                                    'not a combination of both'):
+            left.merge(right, on='value', left_on='x')
+
+        with self.assertRaisesRegex(ValueError,
+                                    'Must pass right_on or right_index=True'):
+            left.merge(right, left_on='x')
+
+        with self.assertRaisesRegex(ValueError,
+                                    'Must pass right_on or right_index=True'):
+            left.merge(right, left_index=True)
+
+        with self.assertRaisesRegex(ValueError,
+                                    'Must pass left_on or left_index=True'):
+            left.merge(right, right_on='y')
+
+        with self.assertRaisesRegex(ValueError,
+                                    'Must pass left_on or left_index=True'):
+            left.merge(right, right_index=True)
+
+        with self.assertRaisesRegex(ValueError,
+                                    'len\\(left_keys\\) must equal len\\(right_keys\\)'):
+            left.merge(right, left_on='value', right_on=['value', 'y'])
+
+        with self.assertRaisesRegex(ValueError,
+                                    'len\\(left_keys\\) must equal len\\(right_keys\\)'):
+            left.merge(right, left_on=['value', 'x'], right_on='value')
+
+        with self.assertRaisesRegex(ValueError,
+                                    "['inner', 'left', 'right', 'full', 'outer']"):
+            left.merge(right, left_index=True, right_index=True, how='foo')
+
+        with self.assertRaisesRegex(AnalysisException,
+                                    'Cannot resolve column name "id"'):
+            left.merge(right, on='id')
 
     def test_clip(self):
         pdf = pd.DataFrame({'A': [0, 2, 4]})