[SPARK-59186][PS] Validate NumPy ufunc operand types instead of silently casting them

[Spenserrrr]

What changes were proposed in this pull request?

numpy_compat.py maps each NumPy ufunc to a Spark expression, and that expression accepts any operand Spark can cast. This PR records the operand types NumPy itself accepts and checks them in the dispatch:

• _np_spark_accepted_types, mapping each ufunc name to the Spark types it accepts, listed once per operand. np.ldexp shows why per operand: it accepts a float first operand, but its exponent must be an integer.
• _check_operand_types, called from maybe_dispatch_ufunc_to_spark_func. It raises TypeError at the ufunc call itself, rather than when the result is later collected. Series, Index and Python scalar operands are all checked, so np.fmod(psser, "8") is caught too.
• Not in the table: np.fmax, np.fmin, np.maximum and np.minimum, which accept two integers or two strings but not one of each, so a list per operand cannot describe them; the four mappings still implemented as a pandas_udf, where NumPy checks the types itself; np.floor_divide, which the floordiv operator handles before this table is consulted; and np.abs and np.bitwise_not, which arrive under the name of the ufunc they alias.

Why two commits. The first covers the 27 ufuncs whose native expression is unreleased — they were pandas_udfs until recently, and the UDF rejected these operand types, so the check restores what they did before the conversion. The second covers the 31 ufuncs that were already native in released versions, so that every ufunc behaves the same way; there, adding the check is a behaviour change. Happy to drop the second commit if you would rather ship only the first.

Boolean columns are left to a follow-up. Most of these ufuncs raise an AnalysisException on one today, where pandas returns a value, because Spark does not implicitly cast a boolean to a numeric type. The follow-up would cast the boolean operands, which was reverted before because it also admitted timestamps and computed a result for them; behind this check it becomes safe. NumPy returns float16 or int8 for a boolean input, so the values will match pandas and the dtype will not.

Why are the changes needed?

np.fmod on a string column returns 1.0, and np.ldexp(float_col, 2.5) returns 42.4, where pandas raises TypeError.

Does this PR introduce any user-facing change?

Yes, for an operand type NumPy does not accept:

>>> np.fmod(psdf.a, psdf.a)  # a is a string column
TypeError: ufunc 'fmod' is not supported for the input types (string, string).

For the 27 ufuncs in the first commit there is no change from any released version, since the pandas_udf they replaced raised as well. For the 31 in the second commit it is a change from released versions. The new behaviour matches pandas in both cases.

How was this patch tested?

Four tests in NumPyCompatTestsMixin, all failing on a pristine numpy_compat.py; one loops the table so later entries need no new row. Every accepted set was checked against pandas and ufunc.types.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Claude Opus 5)

[Spenserrrr]

Hi @zhengruifeng! This PR adds a per-ufunc table of the operand types NumPy accepts and checks it in the dispatch, so an unsupported type raises TypeError instead of being cast and computed. There are three things I want to point out:

• The two commits are split by release status, and the second one has a behavior change in the released version. It can be dropped if you would rather ship only the first one.
• One alternative was putting the check inline in each mapping function, as a leading when(~typeof(c).isin(...), raise_error(...)). That rejects the same types, but it raises lazily: np.fmod(str_col, str_col) returns a Series, and the error only appears at to_pandas() as a SparkRuntimeException. In contrast, pandas raises TypeError at the call itself. Also, a mapping function receives Columns, so it cannot raise in Python, while the dispatch still has the Series and its data type. Happy to switch to the inline form if you prefer it.
• The boolean gap noted in the description is the follow-up I plan to file; the operand cast that fixes it was reverted earlier for letting timestamps through, which this check prevents. Please let me know if you don't want this change.

Could you take a look when you have time? Thanks!

[zhengruifeng]

shall we add a upstream test to monitor numpy's type coersion?

[zhengruifeng]

Merge Summary:

• merged into master 0300805
• merged into branch-4.x eafc987

Posted by merge_spark_pr.py

[Spenserrrr]

shall we add a upstream test to monitor numpy's type coersion?

Yeah I think we can add one. Will create a new PR for this.