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 6ef2cdb955a5 [SPARK-58582][PS] Use native Spark function for NumPy fmod
6ef2cdb955a5 is described below
commit 6ef2cdb955a5d0d586c0792b01819cf7f1b5ad64
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Thu Aug 6 19:22:45 2026 +0800
[SPARK-58582][PS] Use native Spark function for NumPy fmod
### What changes were proposed in this pull request?
Replaces the pandas UDF used for `np.fmod` in the pandas API on Spark with
Spark’s native remainder expression.
Adds coverage for integral and floating inputs, including zero divisors,
infinities, and NaN.
### Why are the changes needed?
Using native Spark expressions avoids pandas UDF serialization and lets
Spark optimize and code-generate this operation while preserving
NumPy-compatible behavior.
### Does this PR introduce _any_ user-facing change?
No. This is a behavior-preserving implementation change.
### How was this patch tested?
`source ~/.zshrc && conda run -n spark-dev-313 python/run-tests --testnames
pyspark.pandas.tests.test_numpy_compat.NumPyCompatTests.test_np_fmod`
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57784 from zhengruifeng/pandas-native-fmod-dev2.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
python/pyspark/pandas/numpy_compat.py | 23 ++++++++++++++++++++--
python/pyspark/pandas/tests/test_numpy_compat.py | 25 ++++++++++++++++++++++++
2 files changed, 46 insertions(+), 2 deletions(-)
diff --git a/python/pyspark/pandas/numpy_compat.py
b/python/pyspark/pandas/numpy_compat.py
index 4894fbfbec7b..4ca221edd8ac 100644
--- a/python/pyspark/pandas/numpy_compat.py
+++ b/python/pyspark/pandas/numpy_compat.py
@@ -18,7 +18,7 @@ from typing import Any, Callable, no_type_check
import numpy as np
-from pyspark.sql import functions as F
+from pyspark.sql import Column, functions as F
from pyspark.sql.pandas.functions import pandas_udf
from pyspark.sql.types import DoubleType, BooleanType
from pyspark.pandas.base import IndexOpsMixin
@@ -99,6 +99,25 @@ unary_np_spark_mappings = {
),
}
+
+def _fmod_func(c1: Column, c2: Column) -> Column:
+ c1_double = c1.cast("double")
+ c2_double = c2.cast("double")
+
+ return F.when(
+ F.typeof(c1).isin("float", "double") | F.typeof(c2).isin("float",
"double"),
+ F.when(c1.isNull() | F.isnan(c1), c1_double)
+ .when(c2.isNull() | F.isnan(c2), c2_double)
+ .when(c2_double == 0, F.lit(float("nan")))
+ .otherwise(F.try_mod(c1_double, c2_double)),
+ ).otherwise(
+ F.when(c1.isNull() | F.isnan(c1), c1_double)
+ .when(c2.isNull() | F.isnan(c2), c2_double)
+ .when(c2_double == 0, F.lit(0.0))
+ .otherwise(F.try_mod(c1_double, c2_double))
+ )
+
+
binary_np_spark_mappings = {
"arctan2": F.atan2,
"bitwise_and": lambda c1, c2: c1.bitwiseAND(c2),
@@ -116,7 +135,7 @@ binary_np_spark_mappings = {
.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]
+ "fmod": _fmod_func,
"gcd": pandas_udf(lambda s1, s2: np.gcd(s1, s2), DoubleType()), # type:
ignore[call-overload]
"heaviside": lambda c1, c2: F.when(
c1.isNull() | F.isnan(c1.cast("double")),
diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py
b/python/pyspark/pandas/tests/test_numpy_compat.py
index b4324ea208b3..0c19b8bfe47b 100644
--- a/python/pyspark/pandas/tests/test_numpy_compat.py
+++ b/python/pyspark/pandas/tests/test_numpy_compat.py
@@ -196,6 +196,31 @@ class NumPyCompatTestsMixin:
self.assert_eq(np.ldexp(psdf.x, psdf.exp), np.ldexp(pdf.x, pdf.exp),
almost=True)
+ def test_np_fmod(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, -64.0, -2.0, -0.0, 0.0, 2.0, 64.0, np.inf,
np.nan, 1.0],
+ "x2": [2.0, 3.0, -2.0, -3.0, -3.0, 0.0, -np.inf, np.inf,
2.0, 0.0],
+ }
+ ),
+ pd.DataFrame(
+ {
+ "x1": pd.array([1, 2, None, None], dtype="Int64"),
+ "x2": pd.array([2, None, 2, 0], dtype="Int64"),
+ }
+ ),
+ ):
+ psdf = ps.from_pandas(pdf)
+
+ self.assert_eq(np.fmod(psdf.x1, psdf.x2), np.fmod(pdf.x1, pdf.x2),
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]}),
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
