Implements multi-index support in Dataframe.filter (#1512)

Implementing multi-index support in `DataFrame.filter`

Author: beobest2

databricks/koalas/frame.py

@@ -8945,34 +8945,58 @@ def filter(self, items=None, like=None, regex=None, axis=None):
             else:
                 raise ValueError("items should be a list-like object.")
             if axis == 0:
-                # TODO: support multi-index here
-                if len(index_scols) != 1:
-                    raise ValueError("Single index must be specified.")
-                col = None
-                for item in items:
-                    if col is None:
-                        col = index_scols[0] == F.lit(item)
-                    else:
-                        col = col | (index_scols[0] == F.lit(item))
+                if len(index_scols) == 1:
+                    col = None
+                    for item in items:
+                        if col is None:
+                            col = index_scols[0] == F.lit(item)
+                        else:
+                            col = col | (index_scols[0] == F.lit(item))
+                elif len(index_scols) > 1:
+                    # for multi-index
+                    col = None
+                    for item in items:
+                        if not isinstance(item, (tuple)):
+                            raise TypeError("Unsupported type {}".format(type(item)))
+                        if not item:
+                            raise ValueError("The item should not be empty.")
+                        midx_col = None
+                        for i, element in enumerate(item):
+                            if midx_col is None:
+                                midx_col = index_scols[i] == F.lit(element)
+                            else:
+                                midx_col = midx_col & (index_scols[i] == F.lit(element))
+                        if col is None:
+                            col = midx_col
+                        else:
+                            col = col | midx_col
+                else:
+                    raise ValueError("Single or multi index must be specified.")
                 return DataFrame(self._internal.with_filter(col))
             elif axis == 1:
                 return self[items]
         elif like is not None:
             if axis == 0:
-                # TODO: support multi-index here
-                if len(index_scols) != 1:
-                    raise ValueError("Single index must be specified.")
-                return DataFrame(self._internal.with_filter(index_scols[0].contains(like)))
+                col = None
+                for index_scol in index_scols:
+                    if col is None:
+                        col = index_scol.contains(like)
+                    else:
+                        col = col | index_scol.contains(like)
+                return DataFrame(self._internal.with_filter(col))
             elif axis == 1:
                 column_labels = self._internal.column_labels
                 output_labels = [label for label in column_labels if any(like in i for i in label)]
                 return self[output_labels]
         elif regex is not None:
             if axis == 0:
-                # TODO: support multi-index here
-                if len(index_scols) != 1:
-                    raise ValueError("Single index must be specified.")
-                return DataFrame(self._internal.with_filter(index_scols[0].rlike(regex)))
+                col = None
+                for index_scol in index_scols:
+                    if col is None:
+                        col = index_scol.rlike(regex)
+                    else:
+                        col = col | index_scol.rlike(regex)
+                return DataFrame(self._internal.with_filter(col))
             elif axis == 1:
                 column_labels = self._internal.column_labels
                 matcher = re.compile(regex)

databricks/koalas/tests/test_dataframe.py

@@ -2685,17 +2685,23 @@ def test_filter(self):
         pdf = pdf.set_index("ba", append=True)
         kdf = ks.from_pandas(pdf)
 
-        with self.assertRaisesRegex(ValueError, "items should be a list-like object"):
-            kdf.filter(items="b")
+        self.assert_eq(
+            kdf.filter(items=[("aa", 1), ("bd", 2)], axis=0).sort_index(),
+            pdf.filter(items=[("aa", 1), ("bd", 2)], axis=0).sort_index(),
+        )
 
-        with self.assertRaisesRegex(ValueError, "Single index must be specified."):
-            kdf.filter(items=["b"], axis=0)
+        with self.assertRaisesRegex(TypeError, "Unsupported type <class 'list'>"):
+            kdf.filter(items=[["aa", 1], ("bd", 2)], axis=0)
 
-        with self.assertRaisesRegex(ValueError, "Single index must be specified."):
-            kdf.filter(like="b", axis="index")
+        with self.assertRaisesRegex(ValueError, "The item should not be empty."):
+            kdf.filter(items=[(), ("bd", 2)], axis=0)
 
-        with self.assertRaisesRegex(ValueError, "Single index must be specified."):
-            kdf.filter(regex="b.*", axis="index")
+        self.assert_eq(kdf.filter(like="b", axis=0), pdf.filter(like="b", axis=0))
+
+        self.assert_eq(kdf.filter(regex="b.*", axis=0), pdf.filter(regex="b.*", axis=0))
+
+        with self.assertRaisesRegex(ValueError, "items should be a list-like object"):
+            kdf.filter(items="b")
 
         with self.assertRaisesRegex(ValueError, "No axis named"):
             kdf.filter(regex="b.*", axis=123)