This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 417b37cc5de6 [SPARK-58694][PS] Use native Spark functions for NumPy
logaddexp and logaddexp2
417b37cc5de6 is described below
commit 417b37cc5de6fed8587b0ec16e84364ed0a266c2
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Wed Aug 12 18:27:51 2026 +0800
[SPARK-58694][PS] Use native Spark functions for NumPy logaddexp and
logaddexp2
### What changes were proposed in this pull request?
This PR replaces the pandas UDF implementations of `np.logaddexp` and
`np.logaddexp2` in the pandas API on Spark with native Spark SQL expressions.
The expressions use a numerically stable max-plus-log1p form and explicitly
handle null, NaN, and infinite inputs.
### Why are the changes needed?
Using native Spark expressions avoids Python and Arrow execution overhead
while preserving NumPy-compatible results.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- Added `test_np_logaddexp`, covering finite values, signed zero, NaN,
infinities, and extreme finite magnitudes for both ufuncs.
- `conda run -n spark-dev-313 ruff check
python/pyspark/pandas/numpy_compat.py
python/pyspark/pandas/tests/test_numpy_compat.py`
- `conda run -n spark-dev-313 ruff format --check
python/pyspark/pandas/numpy_compat.py
python/pyspark/pandas/tests/test_numpy_compat.py`
The PySpark test could not run locally because the available Java 11
runtime cannot load the locally built Java 17 Spark artifacts.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57898 from zhengruifeng/pandas-native-logaddexp-logaddexp2-dev2.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
python/pyspark/pandas/numpy_compat.py | 28 ++++++++++---
python/pyspark/pandas/tests/test_numpy_compat.py | 50 ++++++++++++++++++++++++
2 files changed, 72 insertions(+), 6 deletions(-)
diff --git a/python/pyspark/pandas/numpy_compat.py
b/python/pyspark/pandas/numpy_compat.py
index f361e02712f9..6611abb3b44c 100644
--- a/python/pyspark/pandas/numpy_compat.py
+++ b/python/pyspark/pandas/numpy_compat.py
@@ -134,6 +134,26 @@ def _fmod_func(c1: Column, c2: Column) -> Column:
)
+def _logaddexp_func(c1: Column, c2: Column, base2: bool = False) -> Column:
+ c1_double = c1.cast("double")
+ c2_double = c2.cast("double")
+ difference = F.abs(c1_double - c2_double)
+ maximum = F.greatest(c1_double, c2_double)
+ if base2:
+ log_term = F.log1p(F.pow(F.lit(2.0), -difference)) / F.log(F.lit(2.0))
+ else:
+ log_term = F.log1p(F.exp(-difference))
+
+ return (
+ F.when(c1_double.isNull() | F.isnan(c1_double), c1_double)
+ .when(c2_double.isNull() | F.isnan(c2_double), c2_double)
+ .when((c1_double == float("inf")) | (c2_double == float("inf")),
F.lit(float("inf")))
+ .when(c1_double == float("-inf"), c2_double + F.lit(0.0))
+ .when(c2_double == float("-inf"), c1_double + F.lit(0.0))
+ .otherwise(maximum + log_term)
+ )
+
+
def _floor_divide_func(c1: Column, c2: Column) -> Column:
c1_double = c1.cast("double")
c2_double = c2.cast("double")
@@ -221,12 +241,8 @@ binary_np_spark_mappings = {
"left_shift": lambda c1, c2: F.when((c2 < 0) | (c2 >= 64),
F.lit(0)).otherwise(
F.call_function("shiftleft", c1, c2)
),
- "logaddexp": pandas_udf( # type: ignore[call-overload]
- lambda s1, s2: np.logaddexp(s1, s2), DoubleType()
- ),
- "logaddexp2": pandas_udf( # type: ignore[call-overload]
- lambda s1, s2: np.logaddexp2(s1, s2), DoubleType()
- ),
+ "logaddexp": _logaddexp_func,
+ "logaddexp2": lambda c1, c2: _logaddexp_func(c1, c2, base2=True),
"logical_and": lambda c1, c2: c1.cast(BooleanType()) &
c2.cast(BooleanType()),
"logical_or": lambda c1, c2: c1.cast(BooleanType()) |
c2.cast(BooleanType()),
"logical_xor": lambda c1, c2: (
diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py
b/python/pyspark/pandas/tests/test_numpy_compat.py
index 79c1fb1da7e6..36eb4444fc53 100644
--- a/python/pyspark/pandas/tests/test_numpy_compat.py
+++ b/python/pyspark/pandas/tests/test_numpy_compat.py
@@ -330,6 +330,56 @@ class NumPyCompatTestsMixin:
)
self.assert_eq(result, np.floor_divide(pdf.x1, pdf.x2),
almost=True)
+ def test_np_logaddexp(self):
+ for pdf in (
+ pd.DataFrame(
+ {
+ "x1": [-64, -2, -1, 0, 1, 2, 64],
+ "x2": [2, 3, -2, -3, -3, 0, 2],
+ }
+ ),
+ pd.DataFrame(
+ {
+ "x1": [
+ -np.inf,
+ -np.inf,
+ -2.0,
+ -2.0,
+ -0.0,
+ 0.0,
+ 2.0,
+ np.inf,
+ np.inf,
+ np.nan,
+ -1000.0,
+ -np.inf,
+ -0.0,
+ ],
+ "x2": [
+ -np.inf,
+ 3.0,
+ -np.inf,
+ 2.0,
+ 0.0,
+ -0.0,
+ np.inf,
+ 2.0,
+ np.inf,
+ 2.0,
+ 1000.0,
+ -0.0,
+ -np.inf,
+ ],
+ }
+ ),
+ ):
+ psdf = ps.from_pandas(pdf)
+ for np_func in (np.logaddexp, np.logaddexp2):
+ result = np_func(psdf.x1, psdf.x2)
+ expected = np_func(pdf.x1, pdf.x2)
+ self.assert_eq(result, expected, almost=True)
+ self.assert_eq(np.signbit(result.to_pandas()),
np.signbit(expected))
+
def test_np_fmax_fmin(self):
for pdf in (
pd.DataFrame({"x1": [-2, -1, 0, 1, 2], "x2": [2, 1, 0, -1, -2]}),
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
