diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index febba91883bd5..9eae2a52bc5a0 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -23,6 +23,7 @@ from pyspark.sql.types import DoubleType, LongType, BooleanType from pyspark.pandas.base import IndexOpsMixin + unary_np_spark_mappings = { "abs": F.abs, "absolute": F.abs, @@ -78,7 +79,15 @@ "square": lambda c: c.cast("double") * c, "tan": F.tan, "tanh": F.tanh, - "trunc": pandas_udf(lambda s: np.trunc(s), DoubleType()), # type: ignore[call-overload] + "trunc": lambda c: F.when( + c.cast("double").isNull() + | F.isnan(c.cast("double")) + | c.cast("double").isin(float("-inf"), float("inf")), + c.cast("double"), + ).otherwise( + F.signum(c.cast("double")) + * (F.abs(c.cast("double")) - (F.abs(c.cast("double")) % F.lit(1.0))) + ), } binary_np_spark_mappings = { diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py b/python/pyspark/pandas/tests/test_numpy_compat.py index 432283cb5fe8a..dd6ff2cd8d0e7 100644 --- a/python/pyspark/pandas/tests/test_numpy_compat.py +++ b/python/pyspark/pandas/tests/test_numpy_compat.py @@ -108,6 +108,24 @@ def test_np_math_functions(self): (np.sinh, [-np.inf, -64.0, -2.0, 0.0, 2.0, 64.0, np.inf, np.nan]), (np.square, [-np.inf, -64.0, -2.0, 0.0, 2.0, 64.0, np.inf, np.nan]), (np.tanh, [-np.inf, -64.0, -2.0, 0.0, 2.0, 64.0, np.inf, np.nan]), + ( + np.trunc, + [ + -np.inf, + -64.0, + -2.0, + -1.5, + -0.5, + -0.0, + 0.0, + 0.5, + 1.5, + 2.0, + 64.0, + np.inf, + np.nan, + ], + ), ): with self.subTest(name=np_func.__name__, values=values): pdf = pd.DataFrame({"a": values})