Fix

databricks · Aug 18, 2019 · 68433c7 · 68433c7
1 parent c13ee18
commit 68433c7
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
@@ -3660,7 +3660,11 @@ def pivot_table(self, values=None, index=None, columns=None,
         if columns not in self.columns.values:
             raise ValueError("Wrong columns {}.".format(columns))
 
-        if not all(isinstance(self._internal.spark_type_for(col), NumericType) for col in values):
+        if isinstance(values, list):
+            if not all(isinstance(self._internal.spark_type_for(col), NumericType)
+                       for col in values):
+                raise TypeError('values should be a numeric type.')
+        elif not isinstance(self._internal.spark_type_for(values), NumericType):
             raise TypeError('values should be a numeric type.')
 
         if isinstance(aggfunc, str):

diff --git a/databricks/koalas/tests/test_dataframe.py b/databricks/koalas/tests/test_dataframe.py
@@ -1148,6 +1148,11 @@ def test_pivot_table_errors(self):
             kdf.pivot_table(index=['C'], columns="A", values=['B', 'E'],
                             aggfunc={'B': 'mean', 'E': 'sum'})
 
+        msg = "values should be a numeric type."
+        with self.assertRaisesRegex(TypeError, msg):
+            kdf.pivot_table(index=['C'], columns="A", values='B',
+                            aggfunc={'B': 'mean'})
+
     def test_transpose(self):
         pdf1 = pd.DataFrame(
             data={'col1': [1, 2], 'col2': [3, 4]},