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 d06e4faa76ce [SPARK-59134][PS] Use native Spark function for NumPy
frexp
d06e4faa76ce is described below
commit d06e4faa76ce15305ccd2d684ac3a260aba2e65c
Author: Spenser Sun <[email protected]>
AuthorDate: Tue Sep 1 09:11:46 2026 +0800
[SPARK-59134][PS] Use native Spark function for NumPy frexp
### What changes were proposed in this pull request?
Implements `np.frexp` for pandas-on-Spark with native Spark expressions,
replacing its `NotImplemented` placeholder. It reuses the two-output ufunc
support added for `np.modf` (SPARK-58790), so `mantissa, exponent =
np.frexp(psser)` works on Series, Index and DataFrame.
### Why are the changes needed?
`np.frexp` on a pandas-on-Spark object raises today, while pandas returns a
`(mantissa, exponent)` pair, so the two APIs diverge here. Supporting it
natively continues the NumPy coverage work under SPARK-58532.
### Does this PR introduce _any_ user-facing change?
Yes. `np.frexp` on a pandas-on-Spark Series, Index or DataFrame returns a
`(mantissa, exponent)` tuple instead of raising.
### How was this patch tested?
New `test_np_frexp` in `NumPyCompatTestsMixin` (classic and Connect
parity), which fails without the change; the outputs were also checked
bit-exact against NumPy with a probe covering every power of two ±1 ulp and
both ends of the double range.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5)
Closes #58434 from Spenserrrr/numpy-frexp-two-output.
Authored-by: Spenser Sun <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
python/pyspark/pandas/numpy_compat.py | 62 +++++++++++++++++++++-
python/pyspark/pandas/tests/test_numpy_compat.py | 67 +++++++++++++++++++++++-
2 files changed, 126 insertions(+), 3 deletions(-)
diff --git a/python/pyspark/pandas/numpy_compat.py
b/python/pyspark/pandas/numpy_compat.py
index 0c93be229904..20e319654ddb 100644
--- a/python/pyspark/pandas/numpy_compat.py
+++ b/python/pyspark/pandas/numpy_compat.py
@@ -51,8 +51,6 @@ unary_np_spark_mappings = {
"expm1": F.expm1,
"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.
"invert": F.bitwise_not,
"isfinite": lambda c: F.coalesce(
~(F.isnan(c) | (c == float("inf")) | (c == float("-inf"))),
F.lit(False)
@@ -291,10 +289,70 @@ def _modf_fractional_func(c: Column) -> Column:
)
+def _frexp_scale(c_double: Column, exponent: Column) -> Column:
+ # x * 2**-exponent, in two halves because a single 2**exponent is not
finite across
+ # frexp's range (-1073 for the smallest subnormal, 1024 for the largest
double).
+ half = F.floor(exponent / F.lit(2.0))
+ return c_double / F.pow(F.lit(2.0), half) / F.pow(F.lit(2.0), exponent -
half)
+
+
+def _frexp_finite_exponent(c_double: Column) -> Column:
+ # floor(log2|x|) + 1, which the logarithm can put on the wrong side of an
integer next
+ # to a power of two, but never by more than one: check the scaled
magnitude against
+ # frexp's [0.5, 1) range and step the estimate where it falls outside.
+ estimate = F.floor(F.log2(F.abs(c_double))) + F.lit(1)
+ magnitude = F.abs(_frexp_scale(c_double, estimate))
+ return (
+ F.when(magnitude >= F.lit(1.0), estimate + F.lit(1))
+ .when(magnitude < F.lit(0.5), estimate - F.lit(1))
+ .otherwise(estimate)
+ )
+
+
+def _frexp_is_special_value(c: Column) -> Column:
+ # frexp returns 0, +-inf and nan unchanged as the mantissa, keeping the
sign of a zero,
+ # and pairs them with a zero exponent. from_pandas delivers a NaN as a
null; the sign of a
+ # NaN never survives to here, so numpy's -nan mantissa reads +nan.
+ c_double = c.cast("double")
+ return (
+ c.isNull()
+ | F.isnan(c_double)
+ | (c_double == 0)
+ | c_double.isin(float("-inf"), float("inf"))
+ )
+
+
+def _frexp_mantissa_func(c: Column) -> Column:
+ c_double = c.cast("double")
+ return F.when(_frexp_is_special_value(c), c_double).otherwise(
+ _frexp_scale(c_double, _frexp_finite_exponent(c_double))
+ )
+
+
+def _frexp_exponent_func(c: Column) -> Column:
+ return (
+ F.when(
+ # A genuine <NA> from a nullable dtype (e.g. Int64) arrives as a
non-floating null
+ # and must propagate; a NaN arrives as a floating null and takes
the zero branch
+ # below. A nullable float dtype's <NA> (Float32 or Float64) is a
floating null too,
+ # indistinguishable from a NaN after from_pandas, so it reads 0
instead.
+ c.isNull() & ~F.typeof(c).isin("float", "double"),
+ F.lit(None),
+ )
+ .when(_frexp_is_special_value(c), F.lit(0))
+ .otherwise(_frexp_finite_exponent(c.cast("double")))
+ # numpy returns the exponent as an int32.
+ .cast("int")
+ )
+
+
# 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.frexp(x) -> (mantissa, exponent) with x == mantissa * 2**exponent,
the mantissa
+ # keeping x's sign at a magnitude in [0.5, 1).
+ "frexp": (_frexp_mantissa_func, _frexp_exponent_func),
# np.modf(x) -> (fractional part, integral part); the integral part is
exactly trunc.
"modf": (_modf_fractional_func, unary_np_spark_mappings["trunc"]),
}
diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py
b/python/pyspark/pandas/tests/test_numpy_compat.py
index da4ed34aa121..e2c2927ac224 100644
--- a/python/pyspark/pandas/tests/test_numpy_compat.py
+++ b/python/pyspark/pandas/tests/test_numpy_compat.py
@@ -58,7 +58,6 @@ class NumPyCompatTestsMixin:
"conjugate",
"isnat",
"matmul",
- "frexp",
# Values are close enough but tests failed.
"log", # flaky
"log10", # flaky
@@ -390,6 +389,72 @@ class NumPyCompatTestsMixin:
self.assert_eq(ps_fractional, pd_fractional, almost=True)
self.assert_eq(ps_integral, pd_integral, almost=True)
+ def test_np_frexp(self):
+ # np.frexp(x) returns a tuple (mantissa, exponent), where x ==
mantissa * 2**exponent.
+ for pdf in (
+ pd.DataFrame({"a": [-64, -3, -1, 0, 1, 3, 64]}),
+ pd.DataFrame(
+ {"a": [-np.inf, -64.0, -1.5, -0.5, -0.0, 0.0, 0.5, 1.5, 64.0,
np.inf, np.nan]}
+ ),
+ pd.DataFrame({"a": pd.array([1, -2, None], dtype="Int64")}),
+ ):
+ psdf = ps.from_pandas(pdf)
+ ps_mantissa, ps_exponent = np.frexp(psdf.a)
+ pd_mantissa, pd_exponent = np.frexp(pdf.a)
+ self.assert_eq(ps_mantissa, pd_mantissa, almost=True)
+ self.assert_eq(ps_exponent, pd_exponent, almost=True)
+
+ # Values next to a power of two, where the logarithm behind the
exponent needs its
+ # correction, and both ends of the double range. Compared exactly, not
with almost=True.
+ pdf = pd.DataFrame(
+ {
+ "a": [
+ np.nextafter(2.0, 0.0),
+ 2.0,
+ np.nextafter(2.0, np.inf),
+ np.nextafter(-2.0, 0.0),
+ np.finfo(np.float64).max,
+ np.finfo(np.float64).tiny,
+ np.nextafter(0.0, 1.0), # the smallest subnormal
+ ]
+ }
+ )
+ psdf = ps.from_pandas(pdf)
+ ps_mantissa, ps_exponent = np.frexp(psdf.a)
+ pd_mantissa, pd_exponent = np.frexp(pdf.a)
+ self.assert_eq(ps_mantissa, pd_mantissa)
+ self.assert_eq(ps_exponent, pd_exponent)
+
+ # almost=True treats -0.0 and 0.0 as equal, so check the sign of zero
explicitly:
+ # the mantissa of a zero is that zero, and of +-inf that infinity.
+ 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_mantissa, _ = np.frexp(psdf.a)
+ pd_mantissa, _ = np.frexp(pdf.a)
+ self.assert_eq(np.signbit(ps_mantissa.to_pandas()),
np.signbit(pd_mantissa))
+
+ # DataFrame input: np.frexp returns a tuple of DataFrames, one per
output.
+ pdf = pd.DataFrame(
+ {
+ "a": [-3.5, -1.0, -0.5, 0.0, 6.0],
+ "b": [1.5, -0.0, np.inf, -np.inf, np.nan],
+ }
+ )
+ psdf = ps.from_pandas(pdf)
+ ps_mantissa, ps_exponent = np.frexp(psdf)
+ pd_mantissa, pd_exponent = np.frexp(pdf)
+ self.assert_eq(ps_mantissa, pd_mantissa, almost=True)
+ self.assert_eq(ps_exponent, pd_exponent, almost=True)
+ self.assert_eq(np.signbit(ps_mantissa.to_pandas()),
np.signbit(pd_mantissa))
+
+ # Index input: np.frexp returns a tuple of Index objects.
+ pidx = pd.Index([-3.5, -1.0, -0.5, 0.0, 6.0])
+ psidx = ps.from_pandas(pidx)
+ ps_mantissa, ps_exponent = np.frexp(psidx)
+ pd_mantissa, pd_exponent = np.frexp(pidx)
+ self.assert_eq(ps_mantissa, pd_mantissa, almost=True)
+ self.assert_eq(ps_exponent, pd_exponent, almost=True)
+
def test_floor_divide_func(self):
from pyspark.pandas.utils import _floor_divide_func
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
