Fix round to support multi-index columns. (#802)

Author: ueshin

databricks/koalas/frame.py

@@ -2531,17 +2531,21 @@ def round(self, decimals=0):
         third   0.9  0.0  0.49
         """
         if isinstance(decimals, ks.Series):
-            decimals_list = [kv for kv in decimals._to_internal_pandas().items()]
+            decimals_list = [(k if isinstance(k, tuple) else (k,), v)
+                             for k, v in decimals._to_internal_pandas().items()]
         elif isinstance(decimals, dict):
-            decimals_list = [(k, v) for k, v in decimals.items()]
+            decimals_list = [(k if isinstance(k, tuple) else (k,), v)
+                             for k, v in decimals.items()]
         elif isinstance(decimals, int):
-            decimals_list = [(v, decimals) for v in self._internal.data_columns]
+            decimals_list = [(k, decimals) for k in self._internal.column_index]
         else:
             raise ValueError("decimals must be an integer, a dict-like or a Series")
 
         sdf = self._sdf
-        for decimal in decimals_list:
-            sdf = sdf.withColumn(decimal[0], F.round(scol_for(sdf, decimal[0]), decimal[1]))
+        for idx, decimal in decimals_list:
+            if idx in self._internal.column_index:
+                col = self._internal.column_name_for(idx)
+                sdf = sdf.withColumn(col, F.round(scol_for(sdf, col), decimal))
         return DataFrame(self._internal.copy(sdf=sdf))
 
     def duplicated(self, subset=None, keep='first'):

databricks/koalas/tests/test_dataframe.py

@@ -1480,12 +1480,29 @@ def test_round(self):
                        kdf.round(2))
         self.assert_eq(pdf.round({'A': 1, 'C': 2}),
                        kdf.round({'A': 1, 'C': 2}))
+        self.assert_eq(pdf.round({'A': 1, 'D': 2}),
+                       kdf.round({'A': 1, 'D': 2}))
         self.assert_eq(pdf.round(pser),
                        kdf.round(kser))
         msg = "decimals must be an integer, a dict-like or a Series"
         with self.assertRaisesRegex(ValueError, msg):
             kdf.round(1.5)
 
+        # multi-index columns
+        columns = pd.MultiIndex.from_tuples([('X', 'A'), ('X', 'B'), ('Y', 'C')])
+        pdf.columns = columns
+        kdf.columns = columns
+        pser = pd.Series([1, 0, 2], index=columns)
+        kser = ks.Series([1, 0, 2], index=columns)
+        self.assert_eq(pdf.round(2),
+                       kdf.round(2))
+        self.assert_eq(pdf.round({('X', 'A'): 1, ('Y', 'C'): 2}),
+                       kdf.round({('X', 'A'): 1, ('Y', 'C'): 2}))
+        self.assert_eq(pdf.round({('X', 'A'): 1, 'Y': 2}),
+                       kdf.round({('X', 'A'): 1, 'Y': 2}))
+        self.assert_eq(pdf.round(pser),
+                       kdf.round(kser))
+
     def test_shift(self):
         pdf = pd.DataFrame({'Col1': [10, 20, 15, 30, 45],
                             'Col2': [13, 23, 18, 33, 48],