Fix where to support multi-index columns. (#1249)

Author: ueshin

databricks/koalas/frame.py

@@ -2200,26 +2200,27 @@ def where(self, cond, other=np.nan):
         >>> reset_option("compute.ops_on_diff_frames")
         """
         from databricks.koalas.series import Series
-        tmp_cond_col_name = '__tmp_cond_col_{}__'
-        tmp_other_col_name = '__tmp_other_col_{}__'
+
+        tmp_cond_col_name = '__tmp_cond_col_{}__'.format
+        tmp_other_col_name = '__tmp_other_col_{}__'.format
+
         kdf = self.copy()
         if isinstance(cond, DataFrame):
-            for column in self._internal.data_columns:
-                kdf[tmp_cond_col_name.format(column)] = cond.get(column, False)
+            for idx in self._internal.column_index:
+                kdf[tmp_cond_col_name(name_like_string(idx))] = cond.get(idx, False)
         elif isinstance(cond, Series):
-            for column in self._internal.data_columns:
-                kdf[tmp_cond_col_name.format(column)] = cond
+            for idx in self._internal.column_index:
+                kdf[tmp_cond_col_name(name_like_string(idx))] = cond
         else:
             raise ValueError("type of cond must be a DataFrame or Series")
 
         if isinstance(other, DataFrame):
-            for column in self._internal.data_columns:
-                kdf[tmp_other_col_name.format(column)] = other.get(column, np.nan)
+            for idx in self._internal.column_index:
+                kdf[tmp_other_col_name(name_like_string(idx))] = other.get(idx, np.nan)
         else:
-            for column in self._internal.data_columns:
-                kdf[tmp_other_col_name.format(column)] = other
+            for idx in self._internal.column_index:
+                kdf[tmp_other_col_name(name_like_string(idx))] = other
 
-        sdf = kdf._sdf
         # above logic make spark dataframe looks like below:
         # +-----------------+---+---+------------------+-------------------+------------------+--...
         # |__index_level_0__|  A|  B|__tmp_cond_col_A__|__tmp_other_col_A__|__tmp_cond_col_B__|__...
@@ -2231,22 +2232,18 @@ def where(self, cond, other=np.nan):
         # |                4|  4|500|             false|                 -4|             false|  ...
         # +-----------------+---+---+------------------+-------------------+------------------+--...
 
-        output = []
-        for column in self._internal.data_columns:
-            data_col_name = self._internal.column_name_for(column)
-            output.append(
+        column_scols = []
+        for idx in self._internal.column_index:
+            column_scols.append(
                 F.when(
-                    scol_for(sdf, tmp_cond_col_name.format(column)), scol_for(sdf, data_col_name)
+                    kdf[tmp_cond_col_name(name_like_string(idx))]._scol,
+                    kdf[idx]._scol
                 ).otherwise(
-                    scol_for(sdf, tmp_other_col_name.format(column))
-                ).alias(data_col_name))
+                    kdf[tmp_other_col_name(name_like_string(idx))]._scol
+                ).alias(kdf._internal.column_name_for(idx)))
 
-        index_scols = kdf._internal.index_scols
-        sdf = sdf.select(index_scols + output + list(HIDDEN_COLUMNS))
-
-        return DataFrame(self._internal.copy(
-            sdf=sdf,
-            column_scols=[scol_for(sdf, column) for column in self._internal.data_columns]))
+        return DataFrame(kdf._internal.with_new_columns(column_scols,
+                                                        column_index=self._internal.column_index))
 
     def mask(self, cond, other=np.nan):
         """

databricks/koalas/series.py

@@ -3886,8 +3886,6 @@ def where(self, cond, other=np.nan):
 
         >>> reset_option("compute.ops_on_diff_frames")
         """
-        data_col_name = self._internal.column_name_for(self._internal.column_index[0])
-
         assert isinstance(cond, Series)
 
         # We should check the DataFrame from both `cond` and `other`.
@@ -3901,7 +3899,6 @@ def where(self, cond, other=np.nan):
             kdf['__tmp_cond_col__'] = cond
             kdf['__tmp_other_col__'] = other
 
-            sdf = kdf._sdf
             # above logic makes a Spark DataFrame looks like below:
             # +-----------------+---+----------------+-----------------+
             # |__index_level_0__|  0|__tmp_cond_col__|__tmp_other_col__|
@@ -3913,21 +3910,18 @@ def where(self, cond, other=np.nan):
             # |                4|  4|            true|              500|
             # +-----------------+---+----------------+-----------------+
             condition = F.when(
-                sdf['__tmp_cond_col__'], sdf[data_col_name]
-            ).otherwise(sdf['__tmp_other_col__']).alias(data_col_name)
+                kdf['__tmp_cond_col__']._scol, kdf[self._internal.column_index[0]]._scol
+            ).otherwise(kdf['__tmp_other_col__']._scol).alias(self._internal.data_columns[0])
 
-            sdf = sdf.select(*self._internal.index_columns + [condition])
-            return _col(ks.DataFrame(_InternalFrame(
-                sdf=sdf,
-                index_map=self._internal.index_map,
-                column_index=self._internal.column_index,
-                column_index_names=self._internal.column_index_names)))
+            internal = kdf._internal.with_new_columns([condition],
+                                                      column_index=self._internal.column_index)
+            return _col(DataFrame(internal))
         else:
             if isinstance(other, Series):
                 other = other._scol
             condition = F.when(
                 cond._scol, self._scol
-            ).otherwise(other).alias(data_col_name)
+            ).otherwise(other).alias(self._internal.data_columns[0])
             return self._with_new_scol(condition)
 
     def mask(self, cond, other=np.nan):

databricks/koalas/tests/test_ops_on_diff_frames.py

@@ -571,6 +571,17 @@ def test_where(self):
         self.assert_eq(repr(pdf1.where(pdf2 < -250)),
                        repr(kdf1.where(kdf2 < -250).sort_index()))
 
+        # multi-index columns
+        pdf1 = pd.DataFrame({('X', 'A'): [0, 1, 2, 3, 4],
+                             ('X', 'B'): [100, 200, 300, 400, 500]})
+        pdf2 = pd.DataFrame({('X', 'A'): [0, -1, -2, -3, -4],
+                             ('X', 'B'): [-100, -200, -300, -400, -500]})
+        kdf1 = ks.from_pandas(pdf1)
+        kdf2 = ks.from_pandas(pdf2)
+
+        self.assert_eq(repr(pdf1.where(pdf2 > 100)),
+                       repr(kdf1.where(kdf2 > 100).sort_index()))
+
     def test_mask(self):
         pdf1 = pd.DataFrame({'A': [0, 1, 2, 3, 4], 'B': [100, 200, 300, 400, 500]})
         pdf2 = pd.DataFrame({'A': [0, -1, -2, -3, -4], 'B': [-100, -200, -300, -400, -500]})
@@ -588,6 +599,17 @@ def test_mask(self):
         self.assert_eq(repr(pdf1.mask(pdf2 > -250)),
                        repr(kdf1.mask(kdf2 > -250).sort_index()))
 
+        # multi-index columns
+        pdf1 = pd.DataFrame({('X', 'A'): [0, 1, 2, 3, 4],
+                             ('X', 'B'): [100, 200, 300, 400, 500]})
+        pdf2 = pd.DataFrame({('X', 'A'): [0, -1, -2, -3, -4],
+                             ('X', 'B'): [-100, -200, -300, -400, -500]})
+        kdf1 = ks.from_pandas(pdf1)
+        kdf2 = ks.from_pandas(pdf2)
+
+        self.assert_eq(repr(pdf1.mask(pdf2 < 100)),
+                       repr(kdf1.mask(kdf2 < 100).sort_index()))
+
     def test_multi_index_column_assignment_frame(self):
         pdf = pd.DataFrame({'a': [1, 2, 3, 2], 'b': [4.0, 2.0, 3.0, 1.0]})
         pdf.columns = pd.MultiIndex.from_tuples([('a', 'x'), ('a', 'y')])