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 9943a967242e [SPARK-58790][PS] Use native Spark function for NumPy modf
9943a967242e is described below
commit 9943a967242ed19ea4896d46d0f7159d125bf9e2
Author: Spenser Sun <[email protected]>
AuthorDate: Mon Aug 24 11:29:49 2026 +0800
[SPARK-58790][PS] Use native Spark function for NumPy modf
### What changes were proposed in this pull request?
This PR implements `np.modf` for the pandas API on Spark using native Spark
expressions.
`np.modf(x)` is a multi-output ufunc: it returns a tuple `(fractional_part,
integral_part)`. Previously the mapping for `modf` was a malformed
`pandas_udf(lambda s1, s2: np.modf(s1, s2), DoubleType())` — a unary,
two-output ufunc wired as a two-argument call with a single `DoubleType` return
type — so `np.modf` on a pandas-on-Spark object did not work, and `modf` was
excluded from the compatibility tests via a blacklist.
Multi-output ufuncs cannot be expressed in the existing
`unary_np_spark_mappings` / `binary_np_spark_mappings` tables, where each entry
is a single `Column -> Column` function that `column_op` applies once to
produce one Series. This PR adds:
- A dedicated `multi_output_np_spark_mappings` table whose entries are a
pair of `Column -> Column` functions (every NumPy multi-output ufunc has
exactly two outputs), and a branch in `maybe_dispatch_ufunc_to_spark_func` that
applies `column_op` once per output and returns the two results as a tuple. For
`modf`, the integral part reuses the existing native `trunc` mapping (the
integral part of `modf` is exactly truncation toward zero), and the fractional
part is `signum(x) * (abs(x) % [...]
- Support for the tuple return across all input types. NumPy's
`__array_ufunc__` protocol accepts a tuple return for a multi-output ufunc, so
`fractional, integral = np.modf(x)` unpacks. `Series` and `Index` share
`IndexOpsMixin.__array_ufunc__`, whose return type is widened to
`Union[SeriesOrIndex, Tuple[SeriesOrIndex, SeriesOrIndex]]`. A `DataFrame` has
its own `DataFrame.__array_ufunc__`, which applies the ufunc per column; it now
detects the per-column tuple and regroups the resul [...]
The two `maybe_dispatch_ufunc_to_*` helpers are also typed to return
`SeriesOrIndex` (matching `column_op` and the `data_type_ops` operators) rather
than the base `IndexOpsMixin`, so `__array_ufunc__` returns the dispatch result
directly without a cast.
### Why are the changes needed?
This is a sub-task of SPARK-58532, which replaces `pandas_udf`-based NumPy
ufunc implementations in the pandas API on Spark with native Spark SQL
expressions. Native expressions execute inside the JVM, code-generate, and are
visible to the optimizer, avoiding the per-batch JVM-to-Python-worker round
trip that the UDF fallback incurs. In addition, the previous `modf` mapping was
broken, so this also makes `np.modf` work for pandas-on-Spark.
### Does this PR introduce _any_ user-facing change?
Yes. Previously `np.modf` on a pandas-on-Spark `Series` / `Index` /
`DataFrame` raised an error (the mapping was malformed, and the `DataFrame`
path additionally assumed a single output per column). After this PR it returns
a tuple `(fractional, integral)` matching NumPy:
```python
>>> import numpy as np
>>> import pyspark.pandas as ps
>>> fractional, integral = np.modf(ps.Series([-3.5, 0.0, 2.7]))
>>> fractional.to_numpy()
array([-0.5, 0. , 0.7])
>>> integral.to_numpy()
array([-3., 0., 2.])
>>> fractional, integral = np.modf(ps.DataFrame({"a": [-3.5, 2.7], "b":
[1.5, -0.5]}))
>>> fractional.to_numpy()
array([[-0.5, 0.5],
[ 0.7, -0.5]])
```
### How was this patch tested?
Added `test_np_modf` to `NumPyCompatTestsMixin` (run by both the classic
`NumPyCompatTests` and the Spark Connect `NumPyCompatParityTests` suites). It
checks `Series` inputs (integer, float-edge with `±inf` / `nan` / signed zero /
whole numbers / fractions, and nullable `Int64`) with `np.signbit` comparisons
for signed-zero correctness, and `DataFrame` and `Index` inputs. `modf` was
removed from the test blacklist.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 4.8)
Closes #58021 from Spenserrrr/numpy-modf-two-output.
Authored-by: Spenser Sun <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 40c177b7f243fbba386cb6cecabddbc17734aa8c)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
python/pyspark/pandas/base.py | 12 +++---
python/pyspark/pandas/frame.py | 23 ++++++++---
python/pyspark/pandas/numpy_compat.py | 43 ++++++++++++++++++--
python/pyspark/pandas/tests/test_numpy_compat.py | 50 +++++++++++++++++++++++-
4 files changed, 112 insertions(+), 16 deletions(-)
diff --git a/python/pyspark/pandas/base.py b/python/pyspark/pandas/base.py
index fe14f2becc92..db3fb53ce855 100644
--- a/python/pyspark/pandas/base.py
+++ b/python/pyspark/pandas/base.py
@@ -542,12 +542,14 @@ class IndexOpsMixin(object, metaclass=ABCMeta):
# NDArray Compat
def __array_ufunc__(
self, ufunc: Callable, method: str, *inputs: Any, **kwargs: Any
- ) -> SeriesOrIndex:
+ ) -> Union[SeriesOrIndex, Tuple[SeriesOrIndex, SeriesOrIndex]]:
from pyspark.pandas import numpy_compat
- # Try dunder methods first.
- result = numpy_compat.maybe_dispatch_ufunc_to_dunder_op(
- self, ufunc, method, *inputs, **kwargs
+ # Try dunder methods first. A multi-output ufunc (for example np.modf)
yields a
+ # 2-tuple of results rather than a single Series or Index, so both
`result` and this
+ # method's return type must admit that tuple.
+ result: Union[SeriesOrIndex, Tuple[SeriesOrIndex, SeriesOrIndex]] = (
+ numpy_compat.maybe_dispatch_ufunc_to_dunder_op(self, ufunc,
method, *inputs, **kwargs)
)
# After that, we try with PySpark APIs.
@@ -557,7 +559,7 @@ class IndexOpsMixin(object, metaclass=ABCMeta):
)
if result is not NotImplemented:
- return cast(SeriesOrIndex, result)
+ return result
else:
# TODO: support more APIs?
raise NotImplementedError(
diff --git a/python/pyspark/pandas/frame.py b/python/pyspark/pandas/frame.py
index fdc7b6f7f49d..fc075ba32494 100644
--- a/python/pyspark/pandas/frame.py
+++ b/python/pyspark/pandas/frame.py
@@ -14394,7 +14394,7 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
# NDArray Compat
def __array_ufunc__(
self, ufunc: Callable, method: str, *inputs: Any, **kwargs: Any
- ) -> "DataFrame":
+ ) -> Union["DataFrame", Tuple["DataFrame", "DataFrame"]]:
# TODO: is it possible to deduplicate it with '_map_series_op'?
if all(isinstance(inp, DataFrame) for inp in inputs) and any(
not same_anchor(inp, inputs[0]) for inp in inputs
@@ -14423,14 +14423,25 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
# DataFrame and Series
this = inputs[0]
assert all(inp is this for inp in inputs if isinstance(inp,
DataFrame))
- applied = [
+ column_labels = this._internal.column_labels
+ outputs = [
ufunc(
*[inp[label] if isinstance(inp, DataFrame) else inp for
inp in inputs], **kwargs
- ).rename(label)
- for label in this._internal.column_labels
+ )
+ for label in column_labels
]
- internal = this._internal.with_new_columns(applied)
- return DataFrame(internal)
+ if outputs and isinstance(outputs[0], tuple):
+ # A multi-output ufunc (for example np.modf) yields a
two-element tuple per
+ # column; regroup into one DataFrame per output, matching
numpy/pandas.
+ first_cols = [out[0].rename(label) for out, label in
zip(outputs, column_labels)]
+ second_cols = [out[1].rename(label) for out, label in
zip(outputs, column_labels)]
+ first_df: DataFrame =
DataFrame(this._internal.with_new_columns(first_cols))
+ second_df: DataFrame =
DataFrame(this._internal.with_new_columns(second_cols))
+ return first_df, second_df
+ else:
+ applied = [out.rename(label) for out, label in zip(outputs,
column_labels)]
+ internal = this._internal.with_new_columns(applied)
+ return DataFrame(internal)
def __class_getitem__(cls, params: Any) -> object:
# See https://github.com/python/typing/issues/193
diff --git a/python/pyspark/pandas/numpy_compat.py
b/python/pyspark/pandas/numpy_compat.py
index 765ba4453444..c27aa0225e3f 100644
--- a/python/pyspark/pandas/numpy_compat.py
+++ b/python/pyspark/pandas/numpy_compat.py
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
-from typing import Any, Callable, no_type_check
+from typing import Any, Callable, Tuple, Union, no_type_check
import numpy as np
@@ -22,6 +22,7 @@ from pyspark.loose_version import LooseVersion
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._typing import SeriesOrIndex
from pyspark.pandas.base import IndexOpsMixin
@@ -272,7 +273,6 @@ binary_np_spark_mappings = {
),
"maximum": F.greatest,
"minimum": F.least,
- "modf": pandas_udf(lambda s1, s2: np.modf(s1, s2), DoubleType()), # type:
ignore[call-overload]
"nextafter": pandas_udf( # type: ignore[call-overload]
lambda s1, s2: np.nextafter(s1, s2), DoubleType()
),
@@ -284,11 +284,35 @@ binary_np_spark_mappings = {
}
+def _modf_fractional_func(c: Column) -> Column:
+ c_double = c.cast("double")
+ # signum * (abs % 1) keeps the fractional magnitude with the sign of the
input,
+ # including the signed zero of a whole number (for example -2.0 -> -0.0),
the same
+ # way the "trunc" mapping (reused below for the integral part) relies on
signum.
+ fractional = F.signum(c_double) * (F.abs(c_double) % F.lit(1.0))
+ return (
+ F.when(c.isNull() | F.isnan(c_double), c_double)
+ # +-inf has no fractional part; numpy returns a zero with the input's
sign.
+ .when(c_double == float("inf"), F.lit(0.0))
+ .when(c_double == float("-inf"), F.lit(-0.0))
+ .otherwise(fractional)
+ )
+
+
+# Every multi-output ufunc numpy ships (modf, frexp) has exactly two outputs,
so each entry
+# maps to a pair of Column->Column functions applied independently and
returned as a 2-tuple
+# that numpy's __array_ufunc__ unpacks (for example `fractional, integral =
np.modf(series)`).
+multi_output_np_spark_mappings = {
+ # np.modf(x) -> (fractional part, integral part); the integral part is
exactly trunc.
+ "modf": (_modf_fractional_func, unary_np_spark_mappings["trunc"]),
+}
+
+
# Copied from pandas.
# See also
https://docs.scipy.org/doc/numpy/reference/arrays.classes.html#standard-array-subclasses
def maybe_dispatch_ufunc_to_dunder_op(
ser_or_index: IndexOpsMixin, ufunc: Callable, method: str, *inputs: Any,
**kwargs: Any
-) -> IndexOpsMixin:
+) -> SeriesOrIndex:
special = {
"add",
"sub",
@@ -354,11 +378,22 @@ def maybe_dispatch_ufunc_to_dunder_op(
# See also
https://docs.scipy.org/doc/numpy/reference/arrays.classes.html#standard-array-subclasses
def maybe_dispatch_ufunc_to_spark_func(
ser_or_index: IndexOpsMixin, ufunc: Callable, method: str, *inputs: Any,
**kwargs: Any
-) -> IndexOpsMixin:
+) -> Union[SeriesOrIndex, Tuple[SeriesOrIndex, SeriesOrIndex]]:
from pyspark.pandas.base import column_op
op_name = ufunc.__name__
+ if (
+ method == "__call__"
+ and op_name in multi_output_np_spark_mappings
+ and kwargs.get("out") is None
+ ):
+ # These ufuncs are unary in their input, so the single input is always
a Series
+ # that column_op unwraps to a Column -- no literal wrapping needed.
Build one
+ # Series per output and return them as a 2-tuple (see the mapping's
docstring).
+ first_func, second_func = multi_output_np_spark_mappings[op_name]
+ return column_op(first_func)(*inputs), column_op(second_func)(*inputs)
+
if (
method == "__call__"
and (op_name in unary_np_spark_mappings or op_name in
binary_np_spark_mappings)
diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py
b/python/pyspark/pandas/tests/test_numpy_compat.py
index 36eb4444fc53..17201690cee2 100644
--- a/python/pyspark/pandas/tests/test_numpy_compat.py
+++ b/python/pyspark/pandas/tests/test_numpy_compat.py
@@ -42,7 +42,6 @@ class NumPyCompatTestsMixin:
"log", # flaky
"log10", # flaky
"log1p", # flaky
- "modf",
]
@property
@@ -264,6 +263,55 @@ class NumPyCompatTestsMixin:
self.assert_eq(np.fmod(psdf.x1, psdf.x2), np.fmod(pdf.x1, pdf.x2),
almost=True)
+ def test_np_modf(self):
+ # np.modf(x) returns a tuple (fractional part, integral part).
+ for pdf in (
+ pd.DataFrame({"a": [-64, -2, -1, 0, 1, 2, 64]}),
+ pd.DataFrame(
+ {"a": [-np.inf, -64.0, -2.0, -0.5, -0.0, 0.0, 0.5, 2.0, 64.0,
np.inf, np.nan]}
+ ),
+ pd.DataFrame({"a": pd.array([1, -2, None], dtype="Int64")}),
+ ):
+ psdf = ps.from_pandas(pdf)
+ ps_fractional, ps_integral = np.modf(psdf.a)
+ pd_fractional, pd_integral = np.modf(pdf.a)
+ self.assert_eq(ps_fractional, pd_fractional, almost=True)
+ self.assert_eq(ps_integral, pd_integral, almost=True)
+
+ # almost=True treats -0.0 and 0.0 as equal, so verify the sign of zero
explicitly:
+ # the fractional part of a negative whole number (-2.0 -> -0.0) and
the fractional
+ # part of -inf (-> -0.0) must keep the input's sign, as must the
integral part of a
+ # value in (-1, 0) (-0.5 -> -0.0).
+ pdf = pd.DataFrame({"a": [-2.0, -0.5, -0.0, 0.0, 0.5, 2.0, -np.inf,
np.inf]})
+ psdf = ps.from_pandas(pdf)
+ ps_fractional, ps_integral = np.modf(psdf.a)
+ pd_fractional, pd_integral = np.modf(pdf.a)
+ self.assert_eq(np.signbit(ps_fractional.to_pandas()),
np.signbit(pd_fractional))
+ self.assert_eq(np.signbit(ps_integral.to_pandas()),
np.signbit(pd_integral))
+
+ # DataFrame input: np.modf returns a tuple of DataFrames, one per
output.
+ pdf = pd.DataFrame(
+ {
+ "a": [-3.5, -2.0, -0.5, 0.0, 2.7],
+ "b": [1.5, -0.0, np.inf, -np.inf, np.nan],
+ }
+ )
+ psdf = ps.from_pandas(pdf)
+ ps_fractional, ps_integral = np.modf(psdf)
+ pd_fractional, pd_integral = np.modf(pdf)
+ self.assert_eq(ps_fractional, pd_fractional, almost=True)
+ self.assert_eq(ps_integral, pd_integral, almost=True)
+ self.assert_eq(np.signbit(ps_fractional.to_pandas()),
np.signbit(pd_fractional))
+ self.assert_eq(np.signbit(ps_integral.to_pandas()),
np.signbit(pd_integral))
+
+ # Index input: np.modf returns a tuple of Index objects.
+ pidx = pd.Index([-3.5, -2.0, -0.5, 0.0, 2.7])
+ psidx = ps.from_pandas(pidx)
+ ps_fractional, ps_integral = np.modf(psidx)
+ pd_fractional, pd_integral = np.modf(pidx)
+ self.assert_eq(ps_fractional, pd_fractional, almost=True)
+ self.assert_eq(ps_integral, pd_integral, almost=True)
+
def test_floor_divide_func(self):
from pyspark.pandas.numpy_compat import _floor_divide_func
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
