diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index 20e319654ddb4..906f5512d6374 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -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, @@ -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( @@ -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 diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py b/python/pyspark/pandas/tests/test_numpy_compat.py index e2c2927ac2241..69f7eb30a11ad 100644 --- a/python/pyspark/pandas/tests/test_numpy_compat.py +++ b/python/pyspark/pandas/tests/test_numpy_compat.py @@ -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]),