diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index febba91883bd5..0adb5f026411f 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -65,9 +65,14 @@ "positive": F.positive, "rad2deg": F.degrees, "radians": F.radians, - "reciprocal": pandas_udf( # type: ignore[call-overload] - lambda s: np.reciprocal(s), DoubleType() - ), + "reciprocal": lambda c: F.when(c.cast("double").isNull(), c.cast("double")) + .when( + c.cast("double") == 0, + F.when(c.cast("double").cast("string") == "-0.0", F.lit(float("-inf"))).otherwise( + F.lit(float("inf")) + ), + ) + .otherwise(F.lit(1.0) / c.cast("double")), "rint": pandas_udf(lambda s: np.rint(s), DoubleType()), # type: ignore[call-overload] "sign": F.signum, "signbit": lambda c: F.when(c < 0, True).otherwise(False), diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py b/python/pyspark/pandas/tests/test_numpy_compat.py index 432283cb5fe8a..8ebf020a3c78e 100644 --- a/python/pyspark/pandas/tests/test_numpy_compat.py +++ b/python/pyspark/pandas/tests/test_numpy_compat.py @@ -104,6 +104,10 @@ def test_np_math_functions(self): (np.negative, [-np.inf, -64.0, -2.0, 0.0, 2.0, 64.0, np.inf, np.nan]), (np.positive, [-np.inf, -64.0, -2.0, 0.0, 2.0, 64.0, np.inf, np.nan]), (np.rad2deg, [-np.inf, -64.0, -np.pi, 0.0, np.pi, 64.0, np.inf, np.nan]), + ( + np.reciprocal, + [-np.inf, -64.0, -2.0, -1.0, -0.0, 0.0, 1.0, 2.0, 64.0, np.inf, np.nan], + ), (np.sign, [-np.inf, -64.0, -2.0, -0.0, 0.0, 2.0, 64.0, np.inf, np.nan]), (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]),