Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Spenserrrr wants to merge 2 commits into from
Closed

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e6b73de
[SPARK-59186][PS] Reject NumPy ufunc operand types that NumPy itself …
Spenserrrr Sep 2, 2026
5142572
[SPARK-59186][PS] Check operand types for every NumPy ufunc that rest…
Spenserrrr Sep 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 122 additions & 2 deletions python/pyspark/pandas/numpy_compat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,18 +14,29 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
from typing import Any, Callable, Tuple, Union, no_type_check
from typing import Any, Callable, Dict, List, Optional, Tuple, Union, no_type_check

import numpy as np

from pyspark.loose_version import LooseVersion
from pyspark.pandas._typing import SeriesOrIndex
from pyspark.pandas.base import IndexOpsMixin
from pyspark.pandas.typedef.typehints import as_spark_type
from pyspark.pandas.utils import _floor_divide_func
from pyspark.sql import Column
from pyspark.sql import functions as F
from pyspark.sql.pandas.functions import pandas_udf
from pyspark.sql.types import BooleanType, DoubleType
from pyspark.sql.types import (
BooleanType,
DataType,
DoubleType,
FloatType,
IntegralType,
NullType,
NumericType,
TimestampNTZType,
TimestampType,
)

unary_np_spark_mappings = {
"abs": F.abs,
Expand Down Expand Up @@ -358,6 +369,111 @@ def _frexp_exponent_func(c: Column) -> Column:
}


# NumPy accepts a boolean wherever it accepts an integer.
_INTEGRAL_INPUT_TYPES = (BooleanType, IntegralType)
# Some ufuncs reject a decimal column, so the float types are spelled out rather than using
# NumericType, which includes DecimalType.
_NUMERIC_INPUT_TYPES = (BooleanType, IntegralType, FloatType, DoubleType)
# For the ufuncs that accept a decimal column too.
_NUMERIC_OR_DECIMAL_INPUT_TYPES = (BooleanType, NumericType)
# np.isnan and its neighbours have a datetime64 loop, which reports NaT.
_NUMERIC_OR_TIMESTAMP_INPUT_TYPES = _NUMERIC_INPUT_TYPES + (TimestampType, TimestampNTZType)
# np.sign is the one ufunc with no boolean loop, and NumericType excludes BooleanType.
_SIGN_INPUT_TYPES = (NumericType,)
# pandas reads an all-null column as False for the two-operand bitwise operators.
_BITWISE_INPUT_TYPES = _INTEGRAL_INPUT_TYPES + (NullType,)


# The operand types NumPy accepts, one tuple of Spark types per operand. Without this Spark
# casts the operand instead: np.fmod on a string column answered 1.0. A ufunc is absent when its
# accepted types cannot be listed per operand (np.fmax takes any pair it can compare), when it is
# still a pandas_udf, where NumPy runs per value, or when the entry is unreachable.
_np_spark_accepted_types: Dict[str, Tuple[Tuple[type, ...], ...]] = {
"absolute": (_NUMERIC_OR_DECIMAL_INPUT_TYPES,),
"arccos": (_NUMERIC_INPUT_TYPES,),
"arccosh": (_NUMERIC_INPUT_TYPES,),
"arcsin": (_NUMERIC_INPUT_TYPES,),
"arcsinh": (_NUMERIC_INPUT_TYPES,),
"arctan": (_NUMERIC_INPUT_TYPES,),
"arctan2": (_NUMERIC_INPUT_TYPES, _NUMERIC_INPUT_TYPES),
"arctanh": (_NUMERIC_INPUT_TYPES,),
"bitwise_and": (_BITWISE_INPUT_TYPES, _BITWISE_INPUT_TYPES),
"bitwise_or": (_BITWISE_INPUT_TYPES, _BITWISE_INPUT_TYPES),
"bitwise_xor": (_BITWISE_INPUT_TYPES, _BITWISE_INPUT_TYPES),
"cbrt": (_NUMERIC_INPUT_TYPES,),
"ceil": (_NUMERIC_OR_DECIMAL_INPUT_TYPES,),
"copysign": (_NUMERIC_INPUT_TYPES, _NUMERIC_INPUT_TYPES),
"cos": (_NUMERIC_INPUT_TYPES,),
"cosh": (_NUMERIC_INPUT_TYPES,),
"deg2rad": (_NUMERIC_INPUT_TYPES,),
"degrees": (_NUMERIC_INPUT_TYPES,),
"exp": (_NUMERIC_OR_DECIMAL_INPUT_TYPES,),
"exp2": (_NUMERIC_INPUT_TYPES,),
"expm1": (_NUMERIC_INPUT_TYPES,),
"fabs": (_NUMERIC_INPUT_TYPES,),
"float_power": (_NUMERIC_INPUT_TYPES, _NUMERIC_INPUT_TYPES),
"floor": (_NUMERIC_OR_DECIMAL_INPUT_TYPES,),
"fmod": (_NUMERIC_INPUT_TYPES, _NUMERIC_INPUT_TYPES),
"frexp": (_NUMERIC_INPUT_TYPES,),
"heaviside": (_NUMERIC_INPUT_TYPES, _NUMERIC_INPUT_TYPES),
"hypot": (_NUMERIC_INPUT_TYPES, _NUMERIC_INPUT_TYPES),
"invert": (_INTEGRAL_INPUT_TYPES,),
"isfinite": (_NUMERIC_OR_TIMESTAMP_INPUT_TYPES,),
"isinf": (_NUMERIC_OR_TIMESTAMP_INPUT_TYPES,),
"isnan": (_NUMERIC_OR_TIMESTAMP_INPUT_TYPES,),
# np.ldexp builds x * 2**exp and takes the exponent from an integer loop only.
"ldexp": (_NUMERIC_INPUT_TYPES, _INTEGRAL_INPUT_TYPES),
"left_shift": (_INTEGRAL_INPUT_TYPES, _INTEGRAL_INPUT_TYPES),
"log": (_NUMERIC_INPUT_TYPES,),
"log10": (_NUMERIC_OR_DECIMAL_INPUT_TYPES,),
"log1p": (_NUMERIC_INPUT_TYPES,),
"log2": (_NUMERIC_INPUT_TYPES,),
"logaddexp": (_NUMERIC_INPUT_TYPES, _NUMERIC_INPUT_TYPES),
"logaddexp2": (_NUMERIC_INPUT_TYPES, _NUMERIC_INPUT_TYPES),
"logical_xor": (_NUMERIC_OR_TIMESTAMP_INPUT_TYPES, _NUMERIC_OR_TIMESTAMP_INPUT_TYPES),
"modf": (_NUMERIC_INPUT_TYPES,),
"negative": (_NUMERIC_OR_DECIMAL_INPUT_TYPES,),
"positive": (_NUMERIC_OR_DECIMAL_INPUT_TYPES,),
"rad2deg": (_NUMERIC_INPUT_TYPES,),
"radians": (_NUMERIC_INPUT_TYPES,),
"reciprocal": (_NUMERIC_OR_DECIMAL_INPUT_TYPES,),
"right_shift": (_INTEGRAL_INPUT_TYPES, _INTEGRAL_INPUT_TYPES),
"rint": (_NUMERIC_INPUT_TYPES,),
"sign": (_SIGN_INPUT_TYPES,),
"signbit": (_NUMERIC_INPUT_TYPES,),
"sin": (_NUMERIC_INPUT_TYPES,),
"sinh": (_NUMERIC_INPUT_TYPES,),
"sqrt": (_NUMERIC_OR_DECIMAL_INPUT_TYPES,),
"square": (_NUMERIC_OR_DECIMAL_INPUT_TYPES,),
"tan": (_NUMERIC_INPUT_TYPES,),
"tanh": (_NUMERIC_INPUT_TYPES,),
"trunc": (_NUMERIC_OR_DECIMAL_INPUT_TYPES,),
}


def _check_operand_types(op_name: str, inputs: Tuple[Any, ...]) -> None:
accepted_per_operand = _np_spark_accepted_types.get(op_name)
if accepted_per_operand is None:
return

data_types: List[Optional[DataType]] = []
for inp in inputs:
if isinstance(inp, IndexOpsMixin):
data_types.append(inp.spark.data_type)
else:
# A scalar has no Spark type; an unmappable one gives None and goes unchecked.
data_types.append(as_spark_type(type(inp), raise_error=False))
for data_type, accepted in zip(data_types, accepted_per_operand):
if data_type is not None and not isinstance(data_type, accepted):
raise TypeError(
"ufunc '%s' is not supported for the input types (%s)."
% (
op_name,
", ".join("unknown" if dt is None else dt.simpleString() for dt in data_types),
)
)


# 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(
Expand Down Expand Up @@ -433,6 +549,10 @@ def maybe_dispatch_ufunc_to_spark_func(

op_name = ufunc.__name__

# Check before building the expression, so the error comes from the ufunc call itself.
if method == "__call__" and kwargs.get("out") is None:
_check_operand_types(op_name, inputs)

if (
method == "__call__"
and op_name in multi_output_np_spark_mappings
Expand Down
111 changes: 111 additions & 0 deletions python/pyspark/pandas/tests/test_numpy_compat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,117 @@ def test_np_unsupported_frame(self):
with self.assertRaisesRegex(ValueError, "cannot join with no overlapping index names"):
np.left_shift(psdf1, psdf2)

@property
def operand_type_pdf(self):
return pd.DataFrame(
{
"integer": [7, 8],
"double": [7.5, 8.5],
"decimal": [Decimal("7.5"), Decimal("8.5")],
"string": ["7", "8"],
"timestamp": pd.to_datetime(["2020-01-01", "2020-01-02"]),
"boolean": [True, False],
# All nulls, which Spark types as void.
"null": [None, None],
}
)

def test_np_unsupported_operand_types(self):
# Not at module scope: importing numpy_compat builds its pandas_udf entries, which
# bind the Column class of whichever session mode is active.
from pyspark.pandas.numpy_compat import (
_np_spark_accepted_types,
binary_np_spark_mappings,
multi_output_np_spark_mappings,
unary_np_spark_mappings,
)

psdf = ps.from_pandas(self.operand_type_pdf)

# No ufunc accepts a string column, so one loop covers the whole table.
for op_name, accepted_per_operand in _np_spark_accepted_types.items():
with self.subTest(name=op_name):
self.assertTrue(
op_name in unary_np_spark_mappings
or op_name in binary_np_spark_mappings
or op_name in multi_output_np_spark_mappings,
"%s has no mapping entry" % op_name,
)
with self.assertRaisesRegex(
TypeError,
"ufunc '%s' is not supported for the input types .*string" % op_name,
):
getattr(np, op_name)(*[psdf["string"]] * len(accepted_per_operand))

def test_np_unsupported_operand_types_by_ufunc(self):
# The types only some ufuncs reject, and which operand carries the rejected one.
psdf = ps.from_pandas(self.operand_type_pdf)

for np_func, columns, unsupported in (
(np.cosh, ["timestamp"], "timestamp"),
(np.cosh, ["null"], "void"),
(np.fmod, ["decimal", "decimal"], "decimal"),
# np.invert and the shifts have integer loops only.
(np.invert, ["double"], "double"),
# The rejected operand is the second one here, the first one below.
(np.left_shift, ["integer", "double"], "double"),
(np.copysign, ["double", "string"], "string"),
(np.logaddexp, ["timestamp", "double"], "timestamp"),
# np.ldexp takes its exponent from an integer loop.
(np.ldexp, ["double", "double"], "double"),
# np.sign is the only ufunc here with no boolean loop.
(np.sign, ["boolean"], "boolean"),
):
with self.subTest(np_func=np_func.__name__, unsupported=unsupported):
with self.assertRaisesRegex(
TypeError,
"ufunc '%s' is not supported for the input types .*%s"
% (np_func.__name__, unsupported),
):
np_func(*[psdf[column] for column in columns])

# An Index reaches the same dispatch as a Series.
with self.assertRaisesRegex(TypeError, "ufunc 'cosh' is not supported"):
np.cosh(ps.Index(["7", "8"]))

def test_np_unsupported_scalar_operand_types(self):
# A scalar operand is typed from its Python type, not from a Spark column.
psdf = ps.from_pandas(self.operand_type_pdf)

for np_func, args, unsupported in (
(np.fmod, (psdf["integer"], "8"), "string"),
(np.ldexp, (psdf["double"], 2.5), "double"),
(np.left_shift, (psdf["integer"], 1.5), "double"),
):
with self.subTest(np_func=np_func.__name__, unsupported=unsupported):
with self.assertRaisesRegex(
TypeError,
"ufunc '%s' is not supported for the input types .*%s"
% (np_func.__name__, unsupported),
):
np_func(*args)

def test_np_supported_operand_types(self):
pdf = self.operand_type_pdf
psdf = ps.from_pandas(pdf)

# The accepted cases the rest of this file does not reach: a decimal column, a scalar
# operand, and a ufunc with no table entry.
self.assert_eq(np.square(psdf["decimal"]), np.square(pdf["decimal"]), almost=True)
self.assert_eq(np.trunc(psdf["decimal"]), np.trunc(pdf["decimal"]), almost=True)
self.assert_eq(np.sqrt(psdf["decimal"]), np.sqrt(pdf["decimal"]), almost=True)
self.assert_eq(np.absolute(psdf["decimal"]), np.absolute(pdf["decimal"]), almost=True)
self.assert_eq(np.sign(psdf["decimal"]), np.sign(pdf["decimal"]), almost=True)
# pandas reads an all-null column as False for the bitwise operators, so the check must
# not reject void. The values still differ from pandas, which is a pre-existing gap.
self.assertIsNotNone(np.bitwise_and(psdf["null"], psdf["null"]))
self.assert_eq(np.ldexp(psdf["double"], 2), np.ldexp(pdf["double"], 2), almost=True)
self.assert_eq(np.fmod(psdf["integer"], 2), np.fmod(pdf["integer"], 2), almost=True)
self.assert_eq(np.left_shift(psdf["integer"], 1), np.left_shift(pdf["integer"], 1))
self.assert_eq(
np.fmax(psdf["string"], psdf["string"]), np.fmax(pdf["string"], pdf["string"])
)

def test_np_math_functions(self):
for np_func, values in (
(np.arccosh, [-np.inf, -1.0, 0.0, 1.0, 2.0, 64.0, np.inf, np.nan]),
Expand Down