From 2e280aea1189a1db98087d37b8313c39d86b223c Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Tue, 4 Aug 2026 09:32:43 +0000 Subject: [PATCH] [SPARK-58553][PS] Use native Spark functions for NumPy fmax and fmin --- python/pyspark/pandas/numpy_compat.py | 7 +++++-- python/pyspark/pandas/tests/test_numpy_compat.py | 16 ++++++++++++++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index f333594fd08d..f349b96411a2 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -111,8 +111,11 @@ "floor_divide": pandas_udf( # type: ignore[call-overload] lambda s1, s2: np.floor_divide(s1, s2), DoubleType() ), - "fmax": pandas_udf(lambda s1, s2: np.fmax(s1, s2), DoubleType()), # type: ignore[call-overload] - "fmin": pandas_udf(lambda s1, s2: np.fmin(s1, s2), DoubleType()), # type: ignore[call-overload] + "fmax": lambda c1, c2: F.when(F.isnan(c1.cast("double")), c2) + .when(F.isnan(c2.cast("double")), c1) + .otherwise(F.greatest(c1, c2)) + .cast("double"), + "fmin": lambda c1, c2: F.least(c1, c2).cast("double"), "fmod": pandas_udf(lambda s1, s2: np.fmod(s1, s2), DoubleType()), # type: ignore[call-overload] "gcd": pandas_udf(lambda s1, s2: np.gcd(s1, s2), DoubleType()), # type: ignore[call-overload] "heaviside": pandas_udf( # type: ignore[call-overload] diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py b/python/pyspark/pandas/tests/test_numpy_compat.py index 53bdcc58e36c..ff74b792434c 100644 --- a/python/pyspark/pandas/tests/test_numpy_compat.py +++ b/python/pyspark/pandas/tests/test_numpy_compat.py @@ -171,6 +171,22 @@ def test_np_float_power(self): almost=True, ) + def test_np_fmax_fmin(self): + for pdf in ( + pd.DataFrame({"x1": [-2, -1, 0, 1, 2], "x2": [2, 1, 0, -1, -2]}), + pd.DataFrame( + { + "x1": [np.nan, 2.0, np.nan, -np.inf, -2.0, -0.0, 0.0, 2.0, np.inf], + "x2": [2.0, np.nan, np.nan, np.inf, -np.inf, 0.0, -0.0, np.inf, -np.inf], + } + ), + ): + psdf = ps.from_pandas(pdf) + for np_func in (np.fmax, np.fmin): + result = np_func(psdf.x1, psdf.x2) + expected = np_func(pdf.x1, pdf.x2) + self.assert_eq(result, expected, almost=True) + def test_np_spark_compat_series(self): from pyspark.pandas.numpy_compat import unary_np_spark_mappings, binary_np_spark_mappings