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 e0c8b1283601 [SPARK-58249][PS] Use native functions for NumPy math
ufuncs
e0c8b1283601 is described below
commit e0c8b1283601aa707ae824dcb949b950ad9483d0
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Wed Jul 22 19:49:04 2026 +0800
[SPARK-58249][PS] Use native functions for NumPy math ufuncs
### What changes were proposed in this pull request?
Replace the scalar pandas UDF mappings or column expressions for NumPy
`cosh`, `deg2rad`, `exp2`, `fabs`, `negative`, `positive`, `rad2deg`, `sign`,
`sinh`, `square`, and `tanh` on pandas-on-Spark objects with native Spark SQL
functions or expressions.
`np.square` casts its input to `double` before multiplication, avoiding
fixed-width integer overflow under ANSI mode. The end-to-end pandas-on-Spark
coverage compares the mapped ufunc results with pandas; it includes the `int64`
minimum value for `fabs` and `NaN` and signed-zero values for `sign`.
### Why are the changes needed?
These mappings originated before the corresponding native Spark SQL
functions and expressions were available. Current Spark provides native
equivalents, including Spark Connect support, so this removes the Python worker
boundary while preserving NumPy-compatible results.
### Does this PR introduce _any_ user-facing change?
Yes. `np.square` on overflowing integer values now produces a double-valued
result instead of NumPy fixed-width integer overflow. The other mapped ufuncs
preserve their existing NumPy-compatible results.
### How was this patch tested?
- Added end-to-end pandas-on-Spark coverage for the mapped ufuncs.
- Ran `build/sbt -Phive package`.
- Ran `SPARK_TESTING=1 SPARK_PREPEND_CLASSES=1
PYSPARK_PYTHON=.venv/bin/python PYSPARK_DRIVER_PYTHON=.venv/bin/python
bin/pyspark pyspark.pandas.tests.test_numpy_compat`.
- Ran `git diff --check`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57417 from zhengruifeng/pandas-native-math-ufuncs-dev2.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 31f58da0bca5f9a6340d7fab0f2565f0f8f2140a)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
python/pyspark/pandas/numpy_compat.py | 23 +++++++++++------------
python/pyspark/pandas/tests/test_numpy_compat.py | 13 ++++++++++++-
2 files changed, 23 insertions(+), 13 deletions(-)
diff --git a/python/pyspark/pandas/numpy_compat.py
b/python/pyspark/pandas/numpy_compat.py
index e43dd8a1094e..f2b676045724 100644
--- a/python/pyspark/pandas/numpy_compat.py
+++ b/python/pyspark/pandas/numpy_compat.py
@@ -39,13 +39,13 @@ unary_np_spark_mappings = {
"conj": lambda _: NotImplemented,
"conjugate": lambda _: NotImplemented, # It requires complex type
"cos": F.cos,
- "cosh": pandas_udf(lambda s: np.cosh(s), DoubleType()), # type:
ignore[call-overload]
- "deg2rad": pandas_udf(lambda s: np.deg2rad(s), DoubleType()), # type:
ignore[call-overload]
+ "cosh": F.cosh,
+ "deg2rad": F.radians,
"degrees": F.degrees,
"exp": F.exp,
- "exp2": pandas_udf(lambda s: np.exp2(s), DoubleType()), # type:
ignore[call-overload]
+ "exp2": lambda c: F.pow(F.lit(2.0), c),
"expm1": F.expm1,
- "fabs": pandas_udf(lambda s: np.fabs(s), DoubleType()), # type:
ignore[call-overload]
+ "fabs": lambda c: F.abs(c.cast("double")),
"floor": F.floor,
"frexp": lambda _: NotImplemented, # 'frexp' output lengths become
different
# and it cannot be supported via pandas UDF.
@@ -57,26 +57,25 @@ unary_np_spark_mappings = {
"log": F.log,
"log10": F.log10,
"log1p": F.log1p,
- "log2": pandas_udf(lambda s: np.log2(s), DoubleType()), # type:
ignore[call-overload]
"logical_not": lambda c: ~(c.cast(BooleanType())),
"matmul": lambda _: NotImplemented, # Can return a NumPy array in pandas.
- "negative": lambda c: c * -1,
- "positive": lambda c: c,
- "rad2deg": pandas_udf(lambda s: np.rad2deg(s), DoubleType()), # type:
ignore[call-overload]
+ "negative": F.negative,
+ "positive": F.positive,
+ "rad2deg": F.degrees,
"radians": F.radians,
"reciprocal": pandas_udf( # type: ignore[call-overload]
lambda s: np.reciprocal(s), DoubleType()
),
"rint": pandas_udf(lambda s: np.rint(s), DoubleType()), # type:
ignore[call-overload]
- "sign": lambda c: F.when(c == 0, 0).when(c < 0, -1).otherwise(1),
+ "sign": F.signum,
"signbit": lambda c: F.when(c < 0, True).otherwise(False),
"sin": F.sin,
- "sinh": pandas_udf(lambda s: np.sinh(s), DoubleType()), # type:
ignore[call-overload]
+ "sinh": F.sinh,
"spacing": pandas_udf(lambda s: np.spacing(s), DoubleType()), # type:
ignore[call-overload]
"sqrt": F.sqrt,
- "square": pandas_udf(lambda s: np.square(s), DoubleType()), # type:
ignore[call-overload]
+ "square": lambda c: c.cast("double") * c,
"tan": F.tan,
- "tanh": pandas_udf(lambda s: np.tanh(s), DoubleType()), # type:
ignore[call-overload]
+ "tanh": F.tanh,
"trunc": pandas_udf(lambda s: np.trunc(s), DoubleType()), # 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 d04c2c135a5d..2864f8637796 100644
--- a/python/pyspark/pandas/tests/test_numpy_compat.py
+++ b/python/pyspark/pandas/tests/test_numpy_compat.py
@@ -85,11 +85,22 @@ class NumPyCompatTestsMixin:
with self.assertRaisesRegex(ValueError, "cannot join with no
overlapping index names"):
np.left_shift(psdf1, psdf2)
- def test_np_inverse_hyperbolic_series_uses_native_functions(self):
+ def test_np_math_functions(self):
for np_func, values in (
(np.arccosh, [1.0, 2.0, 4.0]),
(np.arcsinh, [-2.0, 0.0, 2.0]),
(np.arctanh, [-0.5, 0.0, 0.5]),
+ (np.cosh, [-2.0, 0.0, 2.0]),
+ (np.deg2rad, [-180.0, 0.0, 180.0]),
+ (np.exp2, [-2.0, 0.0, 2.0]),
+ (np.fabs, [np.iinfo(np.int64).min, -2, 0, 2]),
+ (np.negative, [-2.0, 0.0, 2.0]),
+ (np.positive, [-2.0, 0.0, 2.0]),
+ (np.rad2deg, [-np.pi, 0.0, np.pi]),
+ (np.sign, [-2.0, -0.0, 0.0, 2.0, np.nan]),
+ (np.sinh, [-2.0, 0.0, 2.0]),
+ (np.square, [-2.0, 0.0, 2.0]),
+ (np.tanh, [-2.0, 0.0, 2.0]),
):
with self.subTest(name=np_func.__name__):
pdf = pd.DataFrame({"a": values})
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
