Repository: databricks/koalas

Issue:
index changed after merge
```python
left_kdf = ks.DataFrame({'A': [1, 2]}, index = ['zero', 'one'])
left_kdf
         A
zero   1
one    2

right_kdf = ks.DataFrame({'B': ['x', 'y']}, index=['one', 'two'])
right_kdf
         B
one    x
two    y

left_kdf.merge(right_kdf, left_index=True, right_index=True)
       A     B
0     2     x
```

Index is changed. Expected answer is 
```python
left_kdf.merge(right_kdf, left_index=True, right_index=True)
       A     B
one     2     x

```






Discussion:
@ueshin @thunterdb. Do we need fix it?
cc @floscha thoughts? 
Looks like something worth fixing to me. I can create a PR for this later today.
> Looks like something worth fixing to me. I can create a PR for this later today.

Hello floscha. If you fix this, i can add some simple function like `join`,`reindex` based on `merge` easily.

Resolving patch:
diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
--- a/databricks/koalas/frame.py
+++ b/databricks/koalas/frame.py
@@ -2705,6 +2705,14 @@ def merge(left, right: 'DataFrame', how: str = 'inner',
         """
         Merge DataFrame objects with a database-style join.
 
+        The index of the resulting DataFrame will be one of the following:
+            - 0...n if no index is used for merging
+            - Index of the left DataFrame if merged only on the index of the right DataFrame
+            - Index of the right DataFrame if merged only on the index of the left DataFrame
+            - All involved indices if merged using the indices of both DataFrames
+                e.g. if `left` with indices (a, x) and `right` with indices (b, x), the result will
+                be an index (x, a, b)
+
         Parameters
         ----------
         right: Object to merge with.
@@ -2780,7 +2788,7 @@ def merge(left, right: 'DataFrame', how: str = 'inner',
 
         >>> left_kdf.merge(right_kdf, left_index=True, right_index=True)
            A  B
-        0  2  x
+        1  2  x
 
         >>> left_kdf.merge(right_kdf, left_index=True, right_index=True, how='left')
            A     B
@@ -2789,8 +2797,8 @@ def merge(left, right: 'DataFrame', how: str = 'inner',
 
         >>> left_kdf.merge(right_kdf, left_index=True, right_index=True, how='right')
              A  B
-        0  2.0  x
-        1  NaN  y
+        1  2.0  x
+        2  NaN  y
 
         >>> left_kdf.merge(right_kdf, left_index=True, right_index=True, how='outer')
              A     B
@@ -2870,7 +2878,6 @@ def merge(left, right: 'DataFrame', how: str = 'inner',
         left_index_columns = set(left._metadata.index_columns)
         right_index_columns = set(right._metadata.index_columns)
 
-        # TODO: in some case, we can keep indexes.
         exprs = []
         for col in left_table.columns:
             if col in left_index_columns:
@@ -2895,7 +2902,48 @@ def merge(left, right: 'DataFrame', how: str = 'inner',
                     scol = scol.alias(col)
             exprs.append(scol)
 
-        return DataFrame(joined_table.select(*exprs))
+        # Retain indices if they are used for joining
+        if left_index:
+            if right_index:
+                exprs.extend(['left_table.%s' % col for col in left_index_columns])
+                exprs.extend(['right_table.%s' % col for col in right_index_columns])
+                index_map = left._metadata.index_map + [idx for idx in right._metadata.index_map
+                                                        if idx not in left._metadata.index_map]
+            else:
+                exprs.extend(['right_table.%s' % col for col in right_index_columns])
+                index_map = right._metadata.index_map
+        elif right_index:
+            exprs.extend(['left_table.%s' % col for col in left_index_columns])
+            index_map = left._metadata.index_map
+        else:
+            index_map = []
+
+        selected_columns = joined_table.select(*exprs)
+
+        # Merge left and right indices after the join by replacing missing values in the left index
+        # with values from the right index and dropping
+        if (how == 'right' or how == 'full') and right_index:
+            for left_index_col, right_index_col in zip(left._metadata.index_columns,
+                                                       right._metadata.index_columns):
+                selected_columns = selected_columns.withColumn(
+                    'left_table.' + left_index_col,
+                    F.when(F.col('left_table.%s' % left_index_col).isNotNull(),
+                           F.col('left_table.%s' % left_index_col))
+                    .otherwise(F.col('right_table.%s' % right_index_col))
+                ).withColumnRenamed(
+                    'left_table.%s' % left_index_col, left_index_col
+                ).drop(F.col('left_table.%s' % left_index_col))
+        if not(left_index and not right_index):
+            selected_columns = selected_columns.drop(*[F.col('right_table.%s' % right_index_col)
+                                                       for right_index_col in right_index_columns
+                                                       if right_index_col in left_index_columns])
+
+        if index_map:
+            data_columns = [c for c in selected_columns.columns
+                            if c not in [idx[0] for idx in index_map]]
+            return DataFrame(selected_columns, index=Metadata(data_columns, index_map))
+        else:
+            return DataFrame(selected_columns)
 
     def append(self, other: 'DataFrame', ignore_index: bool = False,
                verify_integrity: bool = False, sort: bool = False) -> '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
@@ -471,7 +471,9 @@ def test_merge(self):
 
         def check(op):
             k_res = op(left_kdf, right_kdf)
+            k_res = k_res.to_pandas()
             k_res = k_res.sort_values(by=list(k_res.columns))
+            k_res = k_res.reset_index(drop=True)
             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)
@@ -507,6 +509,36 @@ def check(op):
         check(lambda left, right: left.merge(right, left_on='lkey', right_on='rkey',
                                              suffixes=['_left', '_right']))
 
+    def test_merge_retains_indices(self):
+        left_pdf = pd.DataFrame({'A': [0, 1]})
+        right_pdf = pd.DataFrame({'B': [1, 2]}, index=[1, 2])
+        left_kdf = ks.from_pandas(left_pdf)
+        right_kdf = ks.from_pandas(right_pdf)
+
+        self.assert_eq(left_kdf.merge(right_kdf, left_index=True, right_index=True),
+                       left_pdf.merge(right_pdf, left_index=True, right_index=True))
+        self.assert_eq(left_kdf.merge(right_kdf, left_on='A', right_index=True),
+                       left_pdf.merge(right_pdf, left_on='A', right_index=True))
+        self.assert_eq(left_kdf.merge(right_kdf, left_index=True, right_on='B'),
+                       left_pdf.merge(right_pdf, left_index=True, right_on='B'))
+        self.assert_eq(left_kdf.merge(right_kdf, left_on='A', right_on='B'),
+                       left_pdf.merge(right_pdf, left_on='A', right_on='B'))
+
+    def test_merge_how_parameter(self):
+        left_pdf = pd.DataFrame({'A': [1, 2]})
+        right_pdf = pd.DataFrame({'B': ['x', 'y']}, index=[1, 2])
+        left_kdf = ks.from_pandas(left_pdf)
+        right_kdf = ks.from_pandas(right_pdf)
+
+        self.assert_eq(left_kdf.merge(right_kdf, left_index=True, right_index=True),
+                       left_pdf.merge(right_pdf, left_index=True, right_index=True))
+        self.assert_eq(left_kdf.merge(right_kdf, left_index=True, right_index=True, how='left'),
+                       left_pdf.merge(right_pdf, left_index=True, right_index=True, how='left'))
+        self.assert_eq(left_kdf.merge(right_kdf, left_index=True, right_index=True, how='right'),
+                       left_pdf.merge(right_pdf, left_index=True, right_index=True, how='right'))
+        self.assert_eq(left_kdf.merge(right_kdf, left_index=True, right_index=True, how='outer'),
+                       left_pdf.merge(right_pdf, left_index=True, right_index=True, how='outer'))
+
     def test_merge_raises(self):
         left = ks.DataFrame({'value': [1, 2, 3, 5, 6],
                              'x': list('abcde')},