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 ff9a90ef835d [SPARK-58581][PS] Use native Spark function for NumPy
floor_divide
ff9a90ef835d is described below
commit ff9a90ef835da7f37273b31eed8289400452b059
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Wed Aug 12 10:21:20 2026 +0800
[SPARK-58581][PS] Use native Spark function for NumPy floor_divide
### What changes were proposed in this pull request?
This PR replaces the pandas UDF registered for `floor_divide` with
`_floor_divide_func`, a native Spark SQL expression. It retains the existing
NumPy-to-dunder dispatch and tests the native helper directly on Spark columns.
### Why are the changes needed?
The native expression avoids Python and Arrow execution when the function
mapping is used, while keeping the established `np.floor_divide` dispatch
behavior unchanged.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- Added direct coverage for `_floor_divide_func` using integer, nullable,
NaN, infinity, and signed-zero inputs.
- `source ~/.zshrc && conda run -n spark-dev-313 ruff check
python/pyspark/pandas/numpy_compat.py
python/pyspark/pandas/tests/test_numpy_compat.py`
- `source ~/.zshrc && conda run -n spark-dev-313 ruff format --check
python/pyspark/pandas/numpy_compat.py
python/pyspark/pandas/tests/test_numpy_compat.py`
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57783 from zhengruifeng/pandas-native-floor-divide-dev2.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
python/pyspark/pandas/numpy_compat.py | 55 +++++++++++++++++--
python/pyspark/pandas/tests/test_numpy_compat.py | 67 ++++++++++++++++++++++++
2 files changed, 119 insertions(+), 3 deletions(-)
diff --git a/python/pyspark/pandas/numpy_compat.py
b/python/pyspark/pandas/numpy_compat.py
index d2e5fb01597c..dea76a4cd9ec 100644
--- a/python/pyspark/pandas/numpy_compat.py
+++ b/python/pyspark/pandas/numpy_compat.py
@@ -124,6 +124,55 @@ def _fmod_func(c1: Column, c2: Column) -> Column:
)
+def _floor_divide_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(
+ c1_double.isin(float("-inf"), float("inf")),
+ F.when(
+ c2_double == 0,
+ F.when(
+ (c1_double < 0) != (c2_double.cast("string") == "-0.0"),
+ F.lit(float("-inf")),
+ ).otherwise(F.lit(float("inf"))),
+ ).otherwise(F.lit(float("nan"))),
+ )
+ .when(
+ c2_double.isin(float("-inf"), float("inf")),
+ F.when(c1_double == 0, c1_double / c2_double)
+ .when((c1_double < 0) != (c2_double < 0), F.lit(-1.0))
+ .otherwise(F.lit(0.0)),
+ )
+ .when(
+ c2_double == 0,
+ F.when(c1_double == 0, F.lit(float("nan")))
+ .when(
+ (c1_double < 0) != (c2_double.cast("string") == "-0.0"),
+ F.lit(float("-inf")),
+ )
+ .otherwise(F.lit(float("inf"))),
+ )
+ .when(c1_double == 0, c1_double / c2_double)
+ .otherwise((c1_double / c2_double) - F.pmod(c1_double / c2_double,
F.lit(1.0))),
+ ).otherwise(
+ # np.floor_divide on pandas Series returns IEEE values for an integral
zero divisor.
+ F.when(c1.isNull() | F.isnan(c1), c1_double)
+ .when(c2.isNull() | F.isnan(c2), c2_double)
+ .when(
+ c2_double == 0,
+ F.when(c1_double == 0, F.lit(float("nan")))
+ .when(c1_double < 0, F.lit(float("-inf")))
+ .otherwise(F.lit(float("inf"))),
+ )
+ .otherwise((c1_double / c2_double) - F.pmod(c1_double / c2_double,
F.lit(1.0)))
+ )
+
+
binary_np_spark_mappings = {
"arctan2": F.atan2,
"bitwise_and": lambda c1, c2: c1.bitwiseAND(c2),
@@ -133,9 +182,9 @@ binary_np_spark_mappings = {
lambda s1, s2: np.copysign(s1, s2), DoubleType()
),
"float_power": lambda c1, c2: F.pow(c1.cast("double"), c2.cast("double")),
- "floor_divide": pandas_udf( # type: ignore[call-overload]
- lambda s1, s2: np.floor_divide(s1, s2), DoubleType()
- ),
+ # np.floor_divide dispatches to the pandas-on-Spark floordiv dunder
operation
+ # before this registry is consulted, so this mapping is not used for that
case.
+ "floor_divide": _floor_divide_func,
"fmax": lambda c1, c2: F.when(F.isnan(c1.cast("double")), c2)
.when(F.isnan(c2.cast("double")), c1)
.when(c1 == c2, c1)
diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py
b/python/pyspark/pandas/tests/test_numpy_compat.py
index 2573ff586c87..0e5aaf47df23 100644
--- a/python/pyspark/pandas/tests/test_numpy_compat.py
+++ b/python/pyspark/pandas/tests/test_numpy_compat.py
@@ -20,6 +20,7 @@ import pandas as pd
from pyspark import pandas as ps
from pyspark.pandas import set_option, reset_option
+from pyspark.sql import functions as F
from pyspark.testing.pandasutils import PandasOnSparkTestCase
@@ -263,6 +264,72 @@ class NumPyCompatTestsMixin:
self.assert_eq(np.fmod(psdf.x1, psdf.x2), np.fmod(pdf.x1, pdf.x2),
almost=True)
+ def test_floor_divide_func(self):
+ from pyspark.pandas.numpy_compat import _floor_divide_func
+
+ for pdf in (
+ pd.DataFrame(
+ {
+ "x1": [-64, -2, -1, 0, 1, 2, 64, -1, 0],
+ "x2": [2, 3, -2, -3, -3, 0, 2, 0, 0],
+ }
+ ),
+ pd.DataFrame(
+ {
+ "x1": [
+ -np.inf,
+ -64.0,
+ -2.0,
+ -0.0,
+ 0.0,
+ 2.0,
+ 64.0,
+ np.inf,
+ np.nan,
+ 1.0,
+ -1.0,
+ np.inf,
+ -np.inf,
+ np.inf,
+ -np.inf,
+ 1.0,
+ ],
+ "x2": [
+ 2.0,
+ 3.0,
+ -2.0,
+ -3.0,
+ -3.0,
+ 0.0,
+ -np.inf,
+ np.inf,
+ 2.0,
+ 0.0,
+ 0.0,
+ 0.0,
+ 0.0,
+ -0.0,
+ -0.0,
+ np.nan,
+ ],
+ }
+ ),
+ pd.DataFrame(
+ {
+ "x1": pd.array([1, None, None], dtype="Int64"),
+ "x2": pd.array([None, 2, 0], dtype="Int64"),
+ }
+ ),
+ ):
+ psdf = ps.from_pandas(pdf)
+ result = (
+ psdf.spark.frame()
+ .select(_floor_divide_func(F.col("x1"),
F.col("x2")).alias("result"))
+ .toPandas()["result"]
+ .rename(None)
+ )
+ self.assert_eq(result, np.floor_divide(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]
