Repository: databricks/koalas

Issue:
DataFrame.reindex with a koalas Index parameter is not implemented
In pandas, we have :

```python
>>> df = pd.DataFrame({"a": [1, 2, 3]}, index=["A", "B", "C"])
>>> df2 = pd.DataFrame({"a": [4, 5, 6]}, index=["A", "D", "B"])
>>> df.reindex(index=df2.index)
     a
A  1.0
D  NaN
B  2.0
```

But it is not implemented in koalas :

```python
>>> df = ks.from_pandas(df)
>>> df2 = ks.from_pandas(df2)
>>> df = df.reindex(df2.index)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "D:\Dev\koalas\databricks\koalas\frame.py", line 7805, in reindex
    df = df._reindex_index(index, fill_value)
  File "D:\Dev\koalas\databricks\koalas\frame.py", line 7824, in _reindex_index
    kser = ks.Series(list(index))
  File "D:\Dev\koalas\databricks\koalas\indexes.py", line 1989, in __iter__
    return MissingPandasLikeIndex.__iter__(self)
  File "D:\Dev\koalas\databricks\koalas\missing\__init__.py", line 23, in unsupported_function
    class_name=class_name, method_name=method_name, reason=reason
databricks.koalas.exceptions.PandasNotImplementedError: The method `pd.Index.__iter__()` is not implemented. If you want to collect your data as an NumPy array, use 'to_numpy()' instead.
```

It could be interesting to have this feature, I will submit a PR for this.


Discussion:

Resolving patch:
diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
--- a/databricks/koalas/frame.py
+++ b/databricks/koalas/frame.py
@@ -4142,7 +4142,7 @@ def to_delta(
             format="delta",
             partition_cols=partition_cols,
             index_col=index_col,
-            **options
+            **options,
         )
 
     def to_parquet(
@@ -7830,43 +7830,85 @@ def reindex(
 
     def _reindex_index(self, index, fill_value):
         # When axis is index, we can mimic pandas' by a right outer join.
-        assert (
-            len(self._internal.index_spark_column_names) <= 1
-        ), "Index should be single column or not set."
+        nlevels = len(self._internal.index_spark_column_names)
+        assert nlevels <= 1 or (
+            isinstance(index, ks.MultiIndex) and nlevels == index.nlevels
+        ), "MultiIndex DataFrame can only be reindexed with a similar Koalas MultiIndex."
 
-        index_column = self._internal.index_spark_column_names[0]
-
-        kser = ks.Series(list(index))
-        labels = kser._internal.spark_frame.select(kser.spark.column.alias(index_column))
+        index_columns = self._internal.index_spark_column_names
         frame = self._internal.resolved_copy.spark_frame.drop(NATURAL_ORDER_COLUMN_NAME)
 
+        if isinstance(index, ks.Index):
+            if nlevels != index.nlevels:
+                return DataFrame(index._internal.with_new_columns([], column_labels=[])).reindex(
+                    columns=self.columns, fill_value=fill_value
+                )
+
+            index_map = OrderedDict(zip(index_columns, index._internal.index_names))
+            scols = index._internal.index_spark_columns
+            labels = index._internal.spark_frame.select(
+                [scol.alias(index_column) for scol, index_column in zip(scols, index_columns)]
+            )
+        else:
+            kser = ks.Series(list(index))
+            labels = kser._internal.spark_frame.select(kser.spark.column.alias(index_columns[0]))
+            index_map = self._internal._index_map
+
         if fill_value is not None:
-            frame_index_column = verify_temp_column_name(frame, "__frame_index_column__")
-            frame = frame.withColumnRenamed(index_column, frame_index_column)
+            frame_index_columns = [
+                verify_temp_column_name(frame, "__frame_index_column_{}__".format(i))
+                for i in range(nlevels)
+            ]
+            index_scols = [
+                scol_for(frame, index_col).alias(frame_index_col)
+                for index_col, frame_index_col in zip(index_columns, frame_index_columns)
+            ]
+            scols = self._internal.resolved_copy.data_spark_columns
+            frame = frame.select(index_scols + scols)
 
             temp_fill_value = verify_temp_column_name(frame, "__fill_value__")
             labels = labels.withColumn(temp_fill_value, F.lit(fill_value))
 
-            frame_index_scol = scol_for(frame, frame_index_column)
-            labels_index_scol = scol_for(labels, index_column)
+            frame_index_scols = [scol_for(frame, col) for col in frame_index_columns]
+            labels_index_scols = [scol_for(labels, col) for col in index_columns]
+
+            joined_df = frame.join(
+                labels,
+                on=[fcol == lcol for fcol, lcol in zip(frame_index_scols, labels_index_scols)],
+                how="right",
+            )
 
-            joined_df = frame.join(labels, on=[frame_index_scol == labels_index_scol], how="right")
             joined_df = joined_df.select(
-                labels_index_scol,
+                *labels_index_scols,
                 *[
                     F.when(
-                        frame_index_scol.isNull() & labels_index_scol.isNotNull(),
+                        reduce(
+                            lambda c1, c2: c1 & c2,
+                            [
+                                fcol.isNull() & lcol.isNotNull()
+                                for fcol, lcol in zip(frame_index_scols, labels_index_scols)
+                            ],
+                        ),
                         scol_for(joined_df, temp_fill_value),
                     )
                     .otherwise(scol_for(joined_df, col))
                     .alias(col)
                     for col in self._internal.data_spark_column_names
-                ]
+                ],
             )
         else:
-            joined_df = frame.join(labels, on=index_column, how="right")
+            joined_df = frame.join(labels, on=index_columns, how="right")
 
-        internal = self._internal.with_new_sdf(joined_df)
+        sdf = joined_df.drop(NATURAL_ORDER_COLUMN_NAME)
+        internal = InternalFrame(
+            sdf,
+            index_map=index_map,
+            column_labels=self._internal.column_labels,
+            data_spark_columns=[
+                scol_for(sdf, col) for col in self._internal.data_spark_column_names
+            ],
+            column_label_names=self._internal.column_label_names,
+        )
         return DataFrame(internal)
 
     def _reindex_columns(self, columns, fill_value):


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
@@ -2370,7 +2370,8 @@ def test_drop_duplicates(self):
             self.assert_eq(kser.sort_index(), pser.sort_index())
 
     def test_reindex(self):
-        index = ["A", "B", "C", "D", "E"]
+        index = pd.Index(["A", "B", "C", "D", "E"])
+
         pdf = pd.DataFrame({"numbers": [1.0, 2.0, 3.0, 4.0, None]}, index=index)
         kdf = ks.from_pandas(pdf)
 
@@ -2409,29 +2410,101 @@ def test_reindex(self):
             kdf.reindex(columns=["numbers", "2", "3"], fill_value=0.0).sort_index(),
         )
 
+        # Reindexing single Index on single Index
+        pindex2 = pd.Index(["A", "C", "D", "E", "0"], name="index2")
+        kindex2 = ks.from_pandas(pindex2)
+
+        for fill_value in [None, 0]:
+            self.assert_eq(
+                pdf.reindex(index=pindex2, fill_value=fill_value).sort_index(),
+                kdf.reindex(index=kindex2, fill_value=fill_value).sort_index(),
+            )
+
+        pindex2 = pd.DataFrame({"index2": ["A", "C", "D", "E", "0"]}).set_index("index2").index
+        kindex2 = ks.from_pandas(pindex2)
+
+        for fill_value in [None, 0]:
+            self.assert_eq(
+                pdf.reindex(index=pindex2, fill_value=fill_value).sort_index(),
+                kdf.reindex(index=kindex2, fill_value=fill_value).sort_index(),
+            )
+
+        # Reindexing MultiIndex on single Index
+        pindex = pd.MultiIndex.from_tuples(
+            [("A", "B"), ("C", "D"), ("F", "G")], names=["name1", "name2"]
+        )
+        kindex = ks.from_pandas(pindex)
+
+        self.assert_eq(
+            pdf.reindex(index=pindex, fill_value=0.0).sort_index(),
+            kdf.reindex(index=kindex, fill_value=0.0).sort_index(),
+        )
+
         self.assertRaises(TypeError, lambda: kdf.reindex(columns=["numbers", "2", "3"], axis=1))
         self.assertRaises(TypeError, lambda: kdf.reindex(columns=["numbers", "2", "3"], axis=2))
         self.assertRaises(TypeError, lambda: kdf.reindex(index=["A", "B", "C"], axis=1))
         self.assertRaises(TypeError, lambda: kdf.reindex(index=123))
 
+        # Reindexing MultiIndex on MultiIndex
+        pdf = pd.DataFrame({"numbers": [1.0, 2.0, None]}, index=pindex)
+        kdf = ks.from_pandas(pdf)
+        pindex2 = pd.MultiIndex.from_tuples(
+            [("A", "G"), ("C", "D"), ("I", "J")], names=["name1", "name2"]
+        )
+        kindex2 = ks.from_pandas(pindex2)
+
+        for fill_value in [None, 0.0]:
+            self.assert_eq(
+                pdf.reindex(index=pindex2, fill_value=fill_value).sort_index(),
+                kdf.reindex(index=kindex2, fill_value=fill_value).sort_index(),
+            )
+
+        pindex2 = (
+            pd.DataFrame({"index_level_1": ["A", "C", "I"], "index_level_2": ["G", "D", "J"]})
+            .set_index(["index_level_1", "index_level_2"])
+            .index
+        )
+        kindex2 = ks.from_pandas(pindex2)
+
+        for fill_value in [None, 0.0]:
+            self.assert_eq(
+                pdf.reindex(index=pindex2, fill_value=fill_value).sort_index(),
+                kdf.reindex(index=kindex2, fill_value=fill_value).sort_index(),
+            )
+
         columns = pd.MultiIndex.from_tuples([("X", "numbers")])
         pdf.columns = columns
         kdf.columns = columns
 
-        self.assert_eq(
-            pdf.reindex(columns=[("X", "numbers"), ("Y", "2"), ("Y", "3")]).sort_index(),
-            kdf.reindex(columns=[("X", "numbers"), ("Y", "2"), ("Y", "3")]).sort_index(),
-        )
+        # Reindexing MultiIndex index on MultiIndex columns and MultiIndex index
+        for fill_value in [None, 0.0]:
+            self.assert_eq(
+                pdf.reindex(index=pindex2, fill_value=fill_value).sort_index(),
+                kdf.reindex(index=kindex2, fill_value=fill_value).sort_index(),
+            )
 
-        # Using float as fill_value to avoid int64/32 clash
-        self.assert_eq(
-            pdf.reindex(
-                columns=[("X", "numbers"), ("Y", "2"), ("Y", "3")], fill_value=0.0
-            ).sort_index(),
-            kdf.reindex(
-                columns=[("X", "numbers"), ("Y", "2"), ("Y", "3")], fill_value=0.0
-            ).sort_index(),
-        )
+        index = pd.Index(["A", "B", "C", "D", "E"])
+        pdf = pd.DataFrame(data=[1.0, 2.0, 3.0, 4.0, None], index=index, columns=columns)
+        kdf = ks.from_pandas(pdf)
+        pindex2 = pd.Index(["A", "C", "D", "E", "0"], name="index2")
+        kindex2 = ks.from_pandas(pindex2)
+
+        # Reindexing single Index on MultiIndex columns and single Index
+        for fill_value in [None, 0.0]:
+            self.assert_eq(
+                pdf.reindex(index=pindex2, fill_value=fill_value).sort_index(),
+                kdf.reindex(index=kindex2, fill_value=fill_value).sort_index(),
+            )
+
+        for fill_value in [None, 0.0]:
+            self.assert_eq(
+                pdf.reindex(
+                    columns=[("X", "numbers"), ("Y", "2"), ("Y", "3")], fill_value=fill_value
+                ).sort_index(),
+                kdf.reindex(
+                    columns=[("X", "numbers"), ("Y", "2"), ("Y", "3")], fill_value=fill_value
+                ).sort_index(),
+            )
 
         self.assertRaises(TypeError, lambda: kdf.reindex(columns=["X"]))
         self.assertRaises(ValueError, lambda: kdf.reindex(columns=[("X",)]))