Spenserrrr opened a new pull request, #58021:
URL: https://github.com/apache/spark/pull/58021

   ### 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 two-argument 
lambda with a single `DoubleType` return type for what is a unary, two-output 
ufunc — 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: 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 tuple of one `Column -> Column` function per output, and a branch 
in `maybe_dispatch_ufunc_to_spark_func` that applies `column_op` once per 
output function and returns the results as a tuple. NumPy's `__array_ufunc__` 
protocol accepts a tuple return for a multi-output ufunc, so `fractional, 
integral = np.modf(psser)` unpacks as expected.
   
   For `modf`:
   - the integral part reuses the existing native `trunc` mapping (the integral 
part of `modf` is exactly truncation toward zero);
   - the fractional part is `signum(x) * (abs(x) % 1)`, with edge handling so 
that a null/NaN input propagates and `±inf` yields a zero carrying the input's 
sign — matching NumPy, including signed zero (for example `np.modf(-2.0) == 
(-0.0, -2.0)` and `np.modf(-inf) == (-0.0, -inf)`).
   
   `frexp`, the other two-output ufunc, is intentionally left unsupported: its 
exponent output is not straightforward to express with native functions.
   
   The `base.py` change is a type annotation on `__array_ufunc__`'s local 
`result` so it admits the multi-output tuple; there is no behavior change there.
   
   ### 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). 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.])
   ```
   
   ### 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 integer, float-edge (`±inf`, `nan`, signed zero, whole numbers, 
fractions), and nullable `Int64` inputs, and includes an explicit `np.signbit` 
comparison for signed-zero correctness (which the approximate value comparison 
cannot detect). `modf` was removed from the test blacklist. The classic and 
Connect parity suites pass locally.
   
   ### Was this patch authored or co-authored using generative AI tooling?
   
   Generated-by: Claude Code (Claude Opus 4.8)
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to