Fix Frame.abs to support bool type and disallow non-numeric types. (#…

…1980) Fixes `Frame.abs` to support bool type and disallow non-numeric types.
databricks · Dec 23, 2020 · 4c86f3c · 4c86f3c
1 parent f9465aa
commit 4c86f3c
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 5 deletions.
diff --git a/databricks/koalas/generic.py b/databricks/koalas/generic.py
@@ -1748,7 +1748,18 @@ def abs(self) -> Union["DataFrame", "Series"]:
         2  6  30   30
         3  7  40   50
         """
-        return self._apply_series_op(lambda kser: kser._with_new_scol(F.abs(kser.spark.column)))
+
+        def abs(kser):
+            if isinstance(kser.spark.data_type, BooleanType):
+                return kser
+            elif isinstance(kser.spark.data_type, NumericType):
+                return kser.spark.transform(F.abs)
+            else:
+                raise TypeError(
+                    "bad operand type for abs(): {}".format(kser.spark.data_type.simpleString())
+                )
+
+        return self._apply_series_op(abs)
 
     # TODO: by argument only support the grouping name and as_index only for now. Documentation
     # should be updated when it's supported.

diff --git a/databricks/koalas/tests/test_stats.py b/databricks/koalas/tests/test_stats.py
@@ -94,16 +94,26 @@ def test_sum(self):
     def test_abs(self):
         pdf = pd.DataFrame(
             {
-                "A": [1, -2, 3, -4, 5],
-                "B": [1.0, -2, 3, -4, 5],
-                "C": [-6.0, -7, -8, -9, 10],
-                "D": ["a", "b", "c", "d", "e"],
+                "A": [1, -2, np.nan, -4, 5],
+                "B": [1.0, -2, np.nan, -4, 5],
+                "C": [-6.0, -7, -8, np.nan, 10],
+                "D": ["a", "b", "c", "d", np.nan],
+                "E": [True, np.nan, False, True, True],
             }
         )
         kdf = ks.from_pandas(pdf)
         self.assert_eq(kdf.A.abs(), pdf.A.abs())
         self.assert_eq(kdf.B.abs(), pdf.B.abs())
+        self.assert_eq(kdf.E.abs(), pdf.E.abs())
+        # pandas' bug?
+        # self.assert_eq(kdf[["B", "C", "E"]].abs(), pdf[["B", "C", "E"]].abs())
         self.assert_eq(kdf[["B", "C"]].abs(), pdf[["B", "C"]].abs())
+        self.assert_eq(kdf[["E"]].abs(), pdf[["E"]].abs())
+
+        with self.assertRaisesRegex(TypeError, "bad operand type for abs\\(\\): string"):
+            kdf.abs()
+        with self.assertRaisesRegex(TypeError, "bad operand type for abs\\(\\): string"):
+            kdf.D.abs()
 
     def test_axis_on_dataframe(self):
         # The number of each count is intentionally big