diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index 1653121d3124..ef45e5fe4e22 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -98,9 +98,7 @@ "copysign": pandas_udf( # type: ignore[call-overload] lambda s1, s2: np.copysign(s1, s2), DoubleType() ), - "float_power": pandas_udf( # type: ignore[call-overload] - lambda s1, s2: np.float_power(s1, s2), DoubleType() - ), + "float_power": lambda c1, c2: F.pow(c1.cast("double"), c2.cast("double")), "floor_divide": pandas_udf( # type: ignore[call-overload] lambda s1, s2: np.floor_divide(s1, s2), DoubleType() ), diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py b/python/pyspark/pandas/tests/test_numpy_compat.py index 2530d6b04a80..781ef1b0ac9c 100644 --- a/python/pyspark/pandas/tests/test_numpy_compat.py +++ b/python/pyspark/pandas/tests/test_numpy_compat.py @@ -149,6 +149,23 @@ def test_np_bitwise_shift_functions(self): np_func(psdf.value, psdf.bits), np_func(pdf.value, pdf.bits), almost=True ) + def test_np_float_power(self): + for pdf in ( + pd.DataFrame({"base": [-64, -2, -1, 0, 1, 2, 64], "exponent": [-2, -1, 0, 1, 2, 3, 2]}), + pd.DataFrame( + { + "base": [-np.inf, -64.0, -2.0, -0.0, 0.0, 2.0, 64.0, np.inf, np.nan], + "exponent": [2.0, 3.0, -2.0, -3.0, -3.0, 0.5, -2.0, 2.0, 2.0], + } + ), + ): + psdf = ps.from_pandas(pdf) + self.assert_eq( + np.float_power(psdf.base, psdf.exponent), + np.float_power(pdf.base, pdf.exponent), + almost=True, + ) + def test_np_spark_compat_series(self): from pyspark.pandas.numpy_compat import unary_np_spark_mappings, binary_np_spark_mappings