From 46e07d7f2ffee05d8bc1c1a2684f8abbd5396716 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Tue, 4 Aug 2026 08:41:29 +0000 Subject: [PATCH 1/3] [SPARK-58548][PS] Use native Spark function for NumPy heaviside --- python/pyspark/pandas/numpy_compat.py | 16 +++++++++++++--- python/pyspark/pandas/tests/test_numpy_compat.py | 15 +++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index f349b96411a2..59a3d2222952 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -118,9 +118,19 @@ "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] - lambda s1, s2: np.heaviside(s1, s2), DoubleType() - ), + "heaviside": lambda c1, c2: F.when( + c1.isNull() | F.isnan(c1.cast("double")), + c1.cast("double"), + ) + .when( + c1 < 0, + F.lit(0.0), + ) + .when( + c1 == 0, + c2.cast("double"), + ) + .otherwise(F.lit(1.0)), "hypot": F.hypot, "lcm": pandas_udf(lambda s1, s2: np.lcm(s1, s2), DoubleType()), # type: ignore[call-overload] "ldexp": 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 ff74b792434c..5797337945e6 100644 --- a/python/pyspark/pandas/tests/test_numpy_compat.py +++ b/python/pyspark/pandas/tests/test_numpy_compat.py @@ -187,6 +187,21 @@ def test_np_fmax_fmin(self): expected = np_func(pdf.x1, pdf.x2) self.assert_eq(result, expected, almost=True) + def test_np_heaviside(self): + for pdf in ( + pd.DataFrame({"x1": [-2, -1, 0, 1, 2], "x2": [-2, -1, 0, 1, 2]}), + pd.DataFrame( + { + "x1": [-np.inf, -2.0, -0.0, 0.0, 0.0, 2.0, np.inf, np.nan], + "x2": [2.0, -2.0, -0.0, 0.5, np.nan, np.nan, -0.0, 2.0], + } + ), + ): + psdf = ps.from_pandas(pdf) + self.assert_eq( + np.heaviside(psdf.x1, psdf.x2), np.heaviside(pdf.x1, pdf.x2), almost=True + ) + def test_np_spark_compat_series(self): from pyspark.pandas.numpy_compat import unary_np_spark_mappings, binary_np_spark_mappings From a9b5136bc2f5f0d967eded07f1202027695cf85b Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Tue, 4 Aug 2026 21:41:34 +0800 Subject: [PATCH 2/3] nit --- python/pyspark/pandas/numpy_compat.py | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index 59a3d2222952..feb411a21de8 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -121,16 +121,9 @@ "heaviside": lambda c1, c2: F.when( c1.isNull() | F.isnan(c1.cast("double")), c1.cast("double"), - ) - .when( - c1 < 0, - F.lit(0.0), - ) - .when( - c1 == 0, - c2.cast("double"), - ) - .otherwise(F.lit(1.0)), + ).when(c1 < 0, F.lit(0.0) + ).when(c1 == 0, c2.cast("double") + ).otherwise(F.lit(1.0)), "hypot": F.hypot, "lcm": pandas_udf(lambda s1, s2: np.lcm(s1, s2), DoubleType()), # type: ignore[call-overload] "ldexp": pandas_udf( # type: ignore[call-overload] From 5fa0b3e9734154c4951d5dc1d27fa46043d105cf Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Wed, 5 Aug 2026 01:17:25 +0000 Subject: [PATCH 3/3] [SPARK-58548][PS] Fix NumPy heaviside formatting --- python/pyspark/pandas/numpy_compat.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index feb411a21de8..c109df5e9b82 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -121,9 +121,10 @@ "heaviside": lambda c1, c2: F.when( c1.isNull() | F.isnan(c1.cast("double")), c1.cast("double"), - ).when(c1 < 0, F.lit(0.0) - ).when(c1 == 0, c2.cast("double") - ).otherwise(F.lit(1.0)), + ) + .when(c1 < 0, F.lit(0.0)) + .when(c1 == 0, c2.cast("double")) + .otherwise(F.lit(1.0)), "hypot": F.hypot, "lcm": pandas_udf(lambda s1, s2: np.lcm(s1, s2), DoubleType()), # type: ignore[call-overload] "ldexp": pandas_udf( # type: ignore[call-overload]