Fix filter for multi-index columns support. (#859)

Author: ueshin

databricks/koalas/frame.py

@@ -6844,9 +6844,9 @@ def filter(self, items=None, like=None, regex=None, axis=None):
                 sdf = sdf.filter(index_scols[0].contains(like))
                 return DataFrame(self._internal.copy(sdf=sdf))
             elif axis in ('columns', 1, None):
-                data_columns = self._internal.data_columns
-                output_columns = [c for c in data_columns if like in c]
-                return self[output_columns]
+                column_index = self._internal.column_index
+                output_idx = [idx for idx in column_index if any(like in i for i in idx)]
+                return self[output_idx]
         elif regex is not None:
             if axis in ('index', 0):
                 # TODO: support multi-index here
@@ -6855,10 +6855,11 @@ def filter(self, items=None, like=None, regex=None, axis=None):
                 sdf = sdf.filter(index_scols[0].rlike(regex))
                 return DataFrame(self._internal.copy(sdf=sdf))
             elif axis in ('columns', 1, None):
-                data_columns = self._internal.data_columns
+                column_index = self._internal.column_index
                 matcher = re.compile(regex)
-                output_columns = [c for c in data_columns if matcher.search(c) is not None]
-                return self[output_columns]
+                output_idx = [idx for idx in column_index
+                              if any(matcher.search(i) is not None for i in idx)]
+                return self[output_idx]
         else:
             raise TypeError("Must pass either `items`, `like`, or `regex`")
 

databricks/koalas/internal.py

@@ -720,14 +720,14 @@ def from_pandas(pdf: pd.DataFrame) -> '_InternalFrame':
                              for i, name in enumerate(index.names)]
         else:
             name = index.name
-            index_map = [(name if name is not None else '__index_level_0__',
+            index_map = [(str(name) if name is not None else '__index_level_0__',
                           name if name is None or isinstance(name, tuple) else (name,))]
 
         index_columns = [index_column for index_column, _ in index_map]
 
         reset_index = pdf.reset_index()
         reset_index.columns = index_columns + data_columns
-        schema = StructType([StructField(name, infer_pd_series_spark_type(col),
+        schema = StructType([StructField(str(name), infer_pd_series_spark_type(col),
                                          nullable=bool(col.isnull().any()))
                              for name, col in reset_index.iteritems()])
         for name, col in reset_index.iteritems():

databricks/koalas/tests/test_dataframe.py

@@ -1863,10 +1863,10 @@ def test_filter(self):
         self.assert_eq(kdf.filter(like='b', axis='index'), pdf.filter(like='b', axis='index'))
         self.assert_eq(kdf.filter(like='c', axis='columns'), pdf.filter(like='c', axis='columns'))
 
-        self.assert_eq(
-            kdf.filter(regex='b.*', axis='index'), pdf.filter(regex='b.*', axis='index'))
-        self.assert_eq(
-            kdf.filter(regex='b.*', axis='columns'), pdf.filter(regex='b.*', axis='columns'))
+        self.assert_eq(kdf.filter(regex='b.*', axis='index'),
+                       pdf.filter(regex='b.*', axis='index'))
+        self.assert_eq(kdf.filter(regex='b.*', axis='columns'),
+                       pdf.filter(regex='b.*', axis='columns'))
 
         pdf = pdf.set_index('ba', append=True)
         kdf = ks.from_pandas(pdf)
@@ -1892,6 +1892,33 @@ def test_filter(self):
         with self.assertRaisesRegex(TypeError, "mutually exclusive"):
             kdf.filter(regex='b.*', like="aaa")
 
+        # multi-index columns
+        pdf = pd.DataFrame({
+            ('x', 'aa'): ['aa', 'ab', 'bc', 'bd', 'ce'],
+            ('x', 'ba'): [1, 2, 3, 4, 5],
+            ('y', 'cb'): [1., 2., 3., 4., 5.],
+            ('z', 'db'): [1., np.nan, 3., np.nan, 5.],
+        })
+        pdf = pdf.set_index(('x', 'aa'))
+        kdf = ks.from_pandas(pdf)
+
+        self.assert_eq(
+            kdf.filter(items=['ab', 'aa'], axis=0).sort_index(),
+            pdf.filter(items=['ab', 'aa'], axis=0).sort_index())
+        self.assert_eq(
+            kdf.filter(items=[('x', 'ba'), ('z', 'db')], axis=1).sort_index(),
+            pdf.filter(items=[('x', 'ba'), ('z', 'db')], axis=1).sort_index())
+
+        self.assert_eq(kdf.filter(like='b', axis='index'),
+                       pdf.filter(like='b', axis='index'))
+        self.assert_eq(kdf.filter(like='c', axis='columns'),
+                       pdf.filter(like='c', axis='columns'))
+
+        self.assert_eq(kdf.filter(regex='b.*', axis='index'),
+                       pdf.filter(regex='b.*', axis='index'))
+        self.assert_eq(kdf.filter(regex='b.*', axis='columns'),
+                       pdf.filter(regex='b.*', axis='columns'))
+
     def test_pipe(self):
         kdf = ks.DataFrame({'category': ['A', 'A', 'B'],
                             'col1': [1, 2, 3],