diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index febba91883bd5..316e2e9e3afd1 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -20,7 +20,7 @@ from pyspark.sql import functions as F from pyspark.sql.pandas.functions import pandas_udf -from pyspark.sql.types import DoubleType, LongType, BooleanType +from pyspark.sql.types import DoubleType, BooleanType from pyspark.pandas.base import IndexOpsMixin unary_np_spark_mappings = { @@ -107,8 +107,10 @@ "ldexp": pandas_udf( # type: ignore[call-overload] lambda s1, s2: np.ldexp(s1, s2), DoubleType() ), - "left_shift": pandas_udf( # type: ignore[call-overload] - lambda s1, s2: np.left_shift(s1, s2), LongType() + # F.shiftleft accepts literal counts only; call_function also accepts a column. + # NumPy returns zero for counts outside an int64's bit width, unlike JVM shifts. + "left_shift": lambda c1, c2: F.when((c2 < 0) | (c2 >= 64), F.lit(0)).otherwise( + F.call_function("shiftleft", c1, c2) ), "logaddexp": pandas_udf( # type: ignore[call-overload] lambda s1, s2: np.logaddexp(s1, s2), DoubleType() @@ -129,9 +131,11 @@ "nextafter": pandas_udf( # type: ignore[call-overload] lambda s1, s2: np.nextafter(s1, s2), DoubleType() ), - "right_shift": pandas_udf( # type: ignore[call-overload] - lambda s1, s2: np.right_shift(s1, s2), LongType() - ), + # F.shiftright accepts literal counts only; call_function also accepts a column. + # NumPy sign-extends counts outside an int64's bit width, unlike JVM shifts. + "right_shift": lambda c1, c2: F.when( + (c2 < 0) | (c2 >= 64), F.call_function("shiftright", c1, F.lit(63)) + ).otherwise(F.call_function("shiftright", c1, c2)), } diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py b/python/pyspark/pandas/tests/test_numpy_compat.py index 432283cb5fe8a..bd7ba4c3e2103 100644 --- a/python/pyspark/pandas/tests/test_numpy_compat.py +++ b/python/pyspark/pandas/tests/test_numpy_compat.py @@ -115,6 +115,21 @@ def test_np_math_functions(self): self.assert_eq(np_func(psdf.a), np_func(pdf.a), almost=True) + def test_np_bitwise_shift_functions(self): + pdf = pd.DataFrame( + { + "value": [np.iinfo(np.int64).min, -2, -1, 0, 1, 2, np.iinfo(np.int64).max], + "bits": [-1, 0, 1, 63, 64, 65, 2], + } + ) + psdf = ps.from_pandas(pdf) + + for np_func in (np.left_shift, np.right_shift): + with self.subTest(name=np_func.__name__): + self.assert_eq( + np_func(psdf.value, psdf.bits), np_func(pdf.value, pdf.bits), almost=True + ) + def test_np_spark_compat_series(self): from pyspark.pandas.numpy_compat import unary_np_spark_mappings, binary_np_spark_mappings