This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new c5e7e1c4adcd [SPARK-58548][PS] Use native Spark function for NumPy
heaviside
c5e7e1c4adcd is described below
commit c5e7e1c4adcd4542b8844038392d965141f01964
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Wed Aug 5 13:08:01 2026 +0800
[SPARK-58548][PS] Use native Spark function for NumPy heaviside
### What changes were proposed in this pull request?
This PR replaces the pandas UDF implementation of `np.heaviside` in pandas
API on Spark with a native Spark `when` expression. It returns `0.0` for
negative values, the second argument at either signed zero, and `1.0` for
positive values; null and NaN first arguments propagate.
The new test covers integral inputs, signed zero, NaN, infinities, and NaN
as the second argument at zero.
### Why are the changes needed?
Using native Spark expressions avoids pandas UDF and Arrow overhead while
preserving NumPy `heaviside` semantics.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- Added `NumPyCompatTests.test_np_heaviside`.
- `ruff check python/pyspark/pandas/numpy_compat.py
python/pyspark/pandas/tests/test_numpy_compat.py`
- `ruff format --check python/pyspark/pandas/numpy_compat.py
python/pyspark/pandas/tests/test_numpy_compat.py`
- `python -m unittest
pyspark.pandas.tests.test_numpy_compat.NumPyCompatTests.test_np_heaviside`
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex GPT-5
Closes #57749 from zhengruifeng/pandas-native-heaviside-dev2.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 781cc206ec7d058480ca037376847b7056afab20)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
python/pyspark/pandas/numpy_compat.py | 10 +++++++---
python/pyspark/pandas/tests/test_numpy_compat.py | 15 +++++++++++++++
2 files changed, 22 insertions(+), 3 deletions(-)
diff --git a/python/pyspark/pandas/numpy_compat.py
b/python/pyspark/pandas/numpy_compat.py
index f349b96411a2..c109df5e9b82 100644
--- a/python/pyspark/pandas/numpy_compat.py
+++ b/python/pyspark/pandas/numpy_compat.py
@@ -118,9 +118,13 @@ binary_np_spark_mappings = {
"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 @@ class NumPyCompatTestsMixin:
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
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
