Fix column access with named column index. (#629)

Fix an issue on the comment #621 (review).

Author: ueshin

databricks/koalas/frame.py

@@ -3685,13 +3685,12 @@ def pivot(self, index=None, columns=None, values=None):
     def columns(self):
         """The column labels of the DataFrame."""
         if self._internal.column_index is not None:
-            if self._internal.column_index_names is not None:
-                return pd.MultiIndex.from_tuples(self._internal.column_index,
-                                                 names=self._internal.column_index_names)
-            else:
-                return pd.MultiIndex.from_tuples(self._internal.column_index)
+            columns = pd.MultiIndex.from_tuples(self._internal.column_index)
         else:
-            return pd.Index(self._internal.data_columns)
+            columns = pd.Index(self._internal.data_columns)
+        if self._internal.column_index_names is not None:
+            columns.names = self._internal.column_index_names
+        return columns
 
     @columns.setter
     def columns(self, columns):
@@ -3702,17 +3701,24 @@ def columns(self, columns):
                 raise ValueError(
                     "Length mismatch: Expected axis has %d elements, new values have %d elements"
                     % (len(old_names), len(column_index)))
-            self._internal = self._internal.copy(column_index=column_index)
+            column_index_names = columns.names
+            self._internal = self._internal.copy(column_index=column_index,
+                                                 column_index_names=column_index_names)
         else:
             old_names = self._internal.data_columns
             if len(old_names) != len(columns):
                 raise ValueError(
                     "Length mismatch: Expected axis has %d elements, new values have %d elements"
                     % (len(old_names), len(columns)))
+            if isinstance(columns, pd.Index):
+                column_index_names = columns.names
+            else:
+                column_index_names = None
             sdf = self._sdf.select(self._internal.index_scols +
                                    [self._internal.scol_for(old_name).alias(new_name)
                                     for (old_name, new_name) in zip(old_names, columns)])
-            self._internal = self._internal.copy(sdf=sdf, data_columns=columns, column_index=None)
+            self._internal = self._internal.copy(sdf=sdf, data_columns=columns, column_index=None,
+                                                 column_index_names=column_index_names)
 
     @property
     def dtypes(self):
@@ -6226,6 +6232,13 @@ def _get_from_multiindex_column(self, key):
             recursive = True
             for i, (col, idx) in enumerate(columns):
                 columns[i] = (col, tuple([str(key), *idx[1:]]))
+        column_index_names = None
+        if self._internal.column_index_names is not None:
+            # Manage column index names
+            column_index_level = set(len(idx) for _, idx in columns)
+            assert len(column_index_level) == 1
+            column_index_level = list(column_index_level)[0]
+            column_index_names = self._internal.column_index_names[-column_index_level:]
         if all(len(idx) == 1 for _, idx in columns):
             # If len(idx) == 1, then the result is not MultiIndex anymore
             sdf = self._sdf.select(self._internal.index_scols +
@@ -6234,15 +6247,17 @@ def _get_from_multiindex_column(self, key):
             kdf_or_ser = DataFrame(self._internal.copy(
                 sdf=sdf,
                 data_columns=[idx[0] for _, idx in columns],
-                column_index=None))
+                column_index=None,
+                column_index_names=column_index_names))
         else:
             # Otherwise, the result is still MultiIndex and need to manage column_index.
             sdf = self._sdf.select(self._internal.index_scols +
                                    [self._internal.scol_for(col) for col, _ in columns])
             kdf_or_ser = DataFrame(self._internal.copy(
                 sdf=sdf,
                 data_columns=[col for col, _ in columns],
-                column_index=[idx for _, idx in columns]))
+                column_index=[idx for _, idx in columns],
+                column_index_names=column_index_names))
         if recursive:
             kdf_or_ser = kdf_or_ser._pd_getitem(str(key))
         if isinstance(kdf_or_ser, Series):

databricks/koalas/internal.py

@@ -491,11 +491,9 @@ def pandas_df(self):
             pdf = pdf[self.data_columns]
 
         if self._column_index is not None:
-            if self.column_index_names is not None:
-                pdf.columns = pd.MultiIndex.from_tuples(self._column_index,
-                                                        names=self.column_index_names)
-            else:
-                pdf.columns = pd.MultiIndex.from_tuples(self._column_index)
+            pdf.columns = pd.MultiIndex.from_tuples(self._column_index)
+        if self._column_index_names is not None:
+            pdf.columns.names = self._column_index_names
 
         index_names = self.index_names
         if len(index_names) > 0:
@@ -547,13 +545,9 @@ def from_pandas(pdf: pd.DataFrame) -> '_InternalFrame':
         data_columns = [str(col) for col in columns]
         if isinstance(columns, pd.MultiIndex):
             column_index = columns.tolist()
-            if columns.names is not None:
-                column_index_names = columns.names
-            else:
-                column_index_names = None
         else:
             column_index = None
-            column_index_names = None
+        column_index_names = columns.names
 
         index = pdf.index
 

databricks/koalas/tests/test_dataframe.py

@@ -100,20 +100,25 @@ def test_dataframe_multiindex_columns(self):
         self.assert_eq(kdf[('x', 'a')], pdf[('x', 'a')])
         self.assert_eq(kdf[('x', 'a', '1')], pdf[('x', 'a', '1')])
 
-    def test_dataframe_multiindex_names_level(self):
-        columns = pd.MultiIndex.from_tuples([('X', 'A'), ('X', 'B'), ('Y', 'C'), ('Y', 'D')],
-                                            names=['lvl_1', 'lvl_2'])
-        kdf = ks.DataFrame([[1, 2, 3, 4],
-                            [5, 6, 7, 8],
-                            [9, 10, 11, 12],
-                            [13, 14, 15, 16],
-                            [17, 18, 19, 20]], columns=columns)
+    def test_dataframe_column_level_name(self):
+        column = pd.Index(['A', 'B', 'C'], name='X')
+        pdf = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=column)
+        kdf = ks.from_pandas(pdf)
 
+        self.assert_eq(kdf, pdf)
+        self.assert_eq(kdf.columns.names, pdf.columns.names)
+        self.assert_eq(kdf.to_pandas().columns.names, pdf.columns.names)
+
+    def test_dataframe_multiindex_names_level(self):
+        columns = pd.MultiIndex.from_tuples([('X', 'A', 'Z'), ('X', 'B', 'Z'),
+                                             ('Y', 'C', 'Z'), ('Y', 'D', 'Z')],
+                                            names=['lvl_1', 'lvl_2', 'lv_3'])
         pdf = pd.DataFrame([[1, 2, 3, 4],
                             [5, 6, 7, 8],
                             [9, 10, 11, 12],
                             [13, 14, 15, 16],
                             [17, 18, 19, 20]], columns=columns)
+        kdf = ks.from_pandas(pdf)
 
         self.assert_eq(kdf.columns.names, pdf.columns.names)
         self.assert_eq(kdf.to_pandas().columns.names, pdf.columns.names)
@@ -125,6 +130,17 @@ def test_dataframe_multiindex_names_level(self):
                                                 'be list-like or None for a MultiIndex'):
             ks.DataFrame(kdf1._internal.copy(column_index_names='level'))
 
+        self.assert_eq(kdf['X'], pdf['X'])
+        self.assert_eq(kdf['X'].columns.names, pdf['X'].columns.names)
+        self.assert_eq(kdf['X'].to_pandas().columns.names, pdf['X'].columns.names)
+        self.assert_eq(kdf['X']['A'], pdf['X']['A'])
+        self.assert_eq(kdf['X']['A'].columns.names, pdf['X']['A'].columns.names)
+        self.assert_eq(kdf['X']['A'].to_pandas().columns.names, pdf['X']['A'].columns.names)
+        self.assert_eq(kdf[('X', 'A')], pdf[('X', 'A')])
+        self.assert_eq(kdf[('X', 'A')].columns.names, pdf[('X', 'A')].columns.names)
+        self.assert_eq(kdf[('X', 'A')].to_pandas().columns.names, pdf[('X', 'A')].columns.names)
+        self.assert_eq(kdf[('X', 'A', 'Z')], pdf[('X', 'A', 'Z')])
+
     def test_reset_index_with_multiindex_columns(self):
         index = pd.MultiIndex.from_tuples([('bird', 'falcon'),
                                            ('bird', 'parrot'),
@@ -290,6 +306,13 @@ def test_rename_columns(self):
         self.assert_eq(kdf.columns, pd.Index(['x', 'y']))
         self.assert_eq(kdf, pdf)
 
+        columns = pdf.columns
+        columns.name = 'lvl_1'
+
+        kdf.columns = columns
+        self.assert_eq(kdf.columns.names, ['lvl_1'])
+        self.assert_eq(kdf, pdf)
+
         msg = "Length mismatch: Expected axis has 2 elements, new values have 4 elements"
         with self.assertRaisesRegex(ValueError, msg):
             kdf.columns = [1, 2, 3, 4]
@@ -300,6 +323,7 @@ def test_rename_columns(self):
 
         columns = pdf.columns
         self.assert_eq(kdf.columns, columns)
+        self.assert_eq(kdf, pdf)
 
         pdf.columns = ['x', 'y']
         kdf.columns = ['x', 'y']
@@ -311,6 +335,12 @@ def test_rename_columns(self):
         self.assert_eq(kdf.columns, columns)
         self.assert_eq(kdf, pdf)
 
+        columns.names = ['lvl_1', 'lvl_2']
+
+        kdf.columns = columns
+        self.assert_eq(kdf.columns.names, ['lvl_1', 'lvl_2'])
+        self.assert_eq(kdf, pdf)
+
     def test_dot_in_column_name(self):
         self.assert_eq(
             ks.DataFrame(ks.range(1)._sdf.selectExpr("1 as `a.b`"))['a.b'],