Fix assign to support multi-index columns. (#811)

Author: ueshin

databricks/koalas/frame.py

@@ -3118,6 +3118,10 @@ def assign(self, **kwargs):
         feature is supported in pandas for Python 3.6 and later but not in
         Koalas. In Koalas, all items are computed first, and then assigned.
         """
+        return self._assign(kwargs)
+
+    def _assign(self, kwargs):
+        assert isinstance(kwargs, dict)
         from databricks.koalas.series import Series
         for k, v in kwargs.items():
             if not (isinstance(v, (Series, spark.Column)) or
@@ -3127,23 +3131,39 @@ def assign(self, **kwargs):
             if callable(v):
                 kwargs[k] = v(self)
 
-        pairs = list(kwargs.items())
-        sdf = self._sdf
-        for (name, c) in pairs:
-            if isinstance(c, Series):
-                sdf = sdf.withColumn(name, c._scol)
-            elif isinstance(c, Column):
-                sdf = sdf.withColumn(name, c)
-            else:
-                sdf = sdf.withColumn(name, F.lit(c))
+        pairs = {(k if isinstance(k, tuple) else (k,)):
+                 (v._scol if isinstance(v, Series)
+                  else v if isinstance(v, spark.Column)
+                  else F.lit(v))
+                 for k, v in kwargs.items()}
 
-        data_columns = set(self._internal.data_columns)
-        adding_columns = [name for name, _ in pairs if name not in data_columns]
+        scols = []
+        for idx in self._internal.column_index:
+            for i in range(len(idx)):
+                if idx[:len(idx)-i] in pairs:
+                    name = self._internal.column_name_for(idx)
+                    scol = pairs[idx[:len(idx)-i]].alias(name)
+                    break
+            else:
+                scol = self._internal.scol_for(idx)
+            scols.append(scol)
+
+        adding_data_columns = []
+        adding_column_index = []
+        for idx, scol in pairs.items():
+            if idx not in set(i[:len(idx)] for i in self._internal.column_index):
+                name = str(idx) if len(idx) > 1 else idx[0]
+                scols.append(scol.alias(name))
+                adding_data_columns.append(name)
+                adding_column_index.append(idx)
+
+        sdf = self._sdf.select(self._internal.index_scols + scols)
         level = self._internal.column_index_level
-        adding_column_index = [tuple([col, *([''] * (level - 1))]) for col in adding_columns]
+        adding_column_index = [tuple(list(idx) + ([''] * (level - len(idx))))
+                               for idx in adding_column_index]
         internal = self._internal.copy(
             sdf=sdf,
-            data_columns=(self._internal.data_columns + adding_columns),
+            data_columns=(self._internal.data_columns + adding_data_columns),
             column_index=(self._internal.column_index + adding_column_index))
         return DataFrame(internal)
 
@@ -6991,14 +7011,14 @@ def assign_columns(kdf, this_column_index, that_column_index):
                         yield (kdf[this_idx], this_idx)
 
             kdf = align_diff_frames(assign_columns, self, value, fillna=False, how="left")
-        elif isinstance(key, (tuple, list)):
+        elif isinstance(key, list):
             assert isinstance(value, DataFrame)
             # Same DataFrames.
             field_names = value.columns
-            kdf = self.assign(**{k: value[c] for k, c in zip(key, field_names)})
+            kdf = self._assign({k: value[c] for k, c in zip(key, field_names)})
         else:
             # Same Series.
-            kdf = self.assign(**{key: value})
+            kdf = self._assign({key: value})
 
         self._internal = kdf._internal
 

databricks/koalas/tests/test_dataframe.py

@@ -265,25 +265,36 @@ def test_nullable_object(self):
             self.assert_eq(kdf, pdf)
 
     def test_assign(self):
-        kdf = self.kdf.copy()
-        pdf = self.pdf.copy()
+        kdf = self.kdf
+        pdf = self.pdf
 
         kdf['w'] = 1.0
         pdf['w'] = 1.0
 
         self.assert_eq(kdf, pdf)
 
-        kdf['a'] = 'abc'
-        pdf['a'] = 'abc'
+        kdf = kdf.assign(a=kdf['a'] * 2)
+        pdf = pdf.assign(a=pdf['a'] * 2)
 
         self.assert_eq(kdf, pdf)
 
+        # multi-index columns
         columns = pd.MultiIndex.from_tuples([('x', 'a'), ('x', 'b'), ('y', 'w')])
         pdf.columns = columns
         kdf.columns = columns
 
-        pdf['Z'] = 'ZZ'
-        kdf['Z'] = 'ZZ'
+        kdf[('a', 'c')] = 'def'
+        pdf[('a', 'c')] = 'def'
+
+        self.assert_eq(kdf, pdf)
+
+        kdf = kdf.assign(Z='ZZ')
+        pdf = pdf.assign(Z='ZZ')
+
+        self.assert_eq(kdf, pdf)
+
+        kdf['x'] = 'ghi'
+        pdf['x'] = 'ghi'
 
         self.assert_eq(kdf, pdf)