[SPARK-59232][PS] Support boolean operands in NumPy ufuncs

python/pyspark/pandas/numpy_compat.py

@@ -28,6 +28,7 @@
 from pyspark.sql.pandas.functions import pandas_udf
 from pyspark.sql.types import (
     BooleanType,
+    ByteType,
     DataType,
     DoubleType,
     FloatType,
@@ -546,12 +547,26 @@ def maybe_dispatch_ufunc_to_spark_func(
     ser_or_index: IndexOpsMixin, ufunc: Callable, method: str, *inputs: Any, **kwargs: Any
 ) -> Union[SeriesOrIndex, Tuple[SeriesOrIndex, SeriesOrIndex]]:
     from pyspark.pandas.base import column_op
+    from pyspark.pandas.data_type_ops.base import transform_boolean_operand_to_numeric
 
     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 op_name in ("invert", "negative") and isinstance(
+            ser_or_index.spark.data_type, BooleanType
+        ):
+            # np.invert on a boolean is a logical not, and pandas reads np.negative the same way;
+            # Spark keeps that meaning in the logical_not entry instead.
+            op_name = "logical_not"
+        elif op_name in _np_spark_accepted_types:
+            # Spark's functions reject a boolean, so cast it to int8, the narrowest integer NumPy
+            # promotes one to. Only the gated ufuncs are cast, since the rest either return an
+            # operand unchanged, as np.fmax does, or handle a boolean already.
+            inputs = tuple(
+                transform_boolean_operand_to_numeric(inp, spark_type=ByteType()) for inp in inputs
+            )
 
     if (
         method == "__call__"

python/pyspark/pandas/tests/test_numpy_compat.py

@@ -219,6 +219,59 @@ def test_np_supported_operand_types(self):
             np.fmax(psdf["string"], psdf["string"]), np.fmax(pdf["string"], pdf["string"])
         )
 
+    def test_np_boolean_operand(self):
+        from pyspark.pandas.numpy_compat import _np_spark_accepted_types
+
+        pdf = pd.DataFrame({"b": [True, False, True], "i": [7, 8, 9]})
+        psdf = ps.from_pandas(pdf)
+        # The reference is pandas on the same values as int64, where its float loops are float64 as
+        # Spark's math always is; on a boolean or an int8 pandas runs coarser float16 loops.
+        as_int64 = pdf.b.astype("int64")
+        # Skipped: np.sign rejects a boolean; np.log and np.log10 answer nan for a zero where pandas
+        # answers -inf; np.reciprocal divides in int8; np.invert and np.negative are logical.
+        skip = {"sign", "log", "log10", "reciprocal", "invert", "negative"}
+
+        for op_name, accepted_per_operand in _np_spark_accepted_types.items():
+            if op_name in skip:
+                continue
+            np_func = getattr(np, op_name)
+            with self.subTest(name=op_name):
+                result = np_func(*[psdf.b] * len(accepted_per_operand))
+                expected = np_func(*[as_int64] * len(accepted_per_operand))
+                if isinstance(expected, tuple):
+                    for one_result, one_expected in zip(result, expected):
+                        self.assert_eq(one_result, one_expected, almost=True)
+                else:
+                    self.assert_eq(result, expected, almost=True)
+
+        # A boolean scalar operand takes the same promotion as a column.
+        self.assert_eq(np.ldexp(psdf.i, True), np.ldexp(pdf.i, True), almost=True)
+
+        # A nullable boolean keeps an extension dtype and propagates <NA>.
+        nullable_pser = pd.Series(pd.array([True, None, False], dtype="boolean"))
+        nullable = ps.from_pandas(nullable_pser)
+        self.assertIsInstance(np.sqrt(nullable).dtype, pd.Float64Dtype)
+        self.assert_eq(np.sqrt(nullable), np.sqrt(nullable_pser.astype("Int64")), almost=True)
+
+    def test_np_boolean_operand_logical(self):
+        # NumPy applies these logically rather than promoting: np.invert on a boolean is
+        # np.logical_not and pandas reads np.negative the same way; promoting would answer -2, -1.
+        pdf = pd.DataFrame(
+            {
+                "b": [True, False, True],
+                "nullable": pd.array([True, None, False], dtype="boolean"),
+                "i": [-2, 0, 3],
+            }
+        )
+        psdf = ps.from_pandas(pdf)
+
+        for np_func in (np.invert, np.negative):
+            with self.subTest(name=np_func.__name__):
+                self.assert_eq(np_func(psdf.b), np_func(pdf.b))
+                self.assert_eq(np_func(psdf.nullable), np_func(pdf.nullable))
+                # An integer column keeps the arithmetic meaning.
+                self.assert_eq(np_func(psdf.i), np_func(pdf.i))
+
     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]),