Fix assign for multi-index columns support. (#657)

Author: ueshin

databricks/koalas/frame.py

@@ -2785,10 +2785,13 @@ def assign(self, **kwargs):
                 sdf = sdf.withColumn(name, F.lit(c))
 
         data_columns = set(self._internal.data_columns)
+        adding_columns = [name for name, _ in pairs if name not in data_columns]
+        level = self._internal.column_index_level
+        adding_column_index = [tuple([col, *([''] * (level - 1))]) for col in adding_columns]
         internal = self._internal.copy(
             sdf=sdf,
-            data_columns=(self._internal.data_columns +
-                          [name for name, _ in pairs if name not in data_columns]))
+            data_columns=(self._internal.data_columns + adding_columns),
+            column_index=(self._internal.column_index + adding_column_index))
         return DataFrame(internal)
 
     @staticmethod

databricks/koalas/tests/test_dataframe.py

@@ -267,6 +267,15 @@ def test_assign(self):
 
         self.assert_eq(kdf, pdf)
 
+        columns = pd.MultiIndex.from_tuples([('x', 'a'), ('x', 'b'), ('y', 'w')])
+        pdf.columns = columns
+        kdf.columns = columns
+
+        pdf['Z'] = 'ZZ'
+        kdf['Z'] = 'ZZ'
+
+        self.assert_eq(kdf, pdf)
+
     def test_head_tail(self):
         kdf = self.kdf
         pdf = self.pdf
@@ -848,6 +857,12 @@ def test_append(self):
         self.assert_eq(kdf.append(multi_index_kdf, ignore_index=True),
                        pdf.append(multi_index_pdf, ignore_index=True))
 
+        columns = pd.MultiIndex.from_tuples([('A', 'X'), ('A', 'Y')])
+        pdf.columns = columns
+        kdf.columns = columns
+
+        self.assert_eq(kdf.append(kdf), pdf.append(pdf))
+
     def test_clip(self):
         pdf = pd.DataFrame({'A': [0, 2, 4]})
         kdf = ks.from_pandas(pdf)