diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index 33538ef9de87f..5682f55e5df89 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -100,6 +100,20 @@ } +def _copysign_func(c1: Column, c2: Column) -> Column: + # Sign of y is taken from its IEEE-754 sign bit, so -0.0 counts as negative. + # c2 < 0 misses -0.0, so detect it via the string cast, the same way the + # 'reciprocal' mapping distinguishes -0.0 from 0.0. NaN's sign bit is positive + # and c2 < 0 is already false for NaN, so it correctly falls through to +1.0. + sign = F.when((c2 < 0) | (c2.cast("string") == "-0.0"), F.lit(-1.0)).otherwise(F.lit(1.0)) + # An integer y column's NULL is a genuine missing value and propagates. A + # float/double column instead stores its missing value as NaN (surfaced as a + # Spark NULL by pandas-on-Spark), for which copysign(x, NaN) returns |x|. + return F.when( + c2.isNull() & ~F.typeof(c2).isin("float", "double"), F.lit(None).cast("double") + ).otherwise(F.abs(c1.cast("double")) * sign) + + def _fmod_func(c1: Column, c2: Column) -> Column: c1_double = c1.cast("double") c2_double = c2.cast("double") @@ -123,9 +137,7 @@ def _fmod_func(c1: Column, c2: Column) -> Column: "bitwise_and": lambda c1, c2: c1.bitwiseAND(c2), "bitwise_or": lambda c1, c2: c1.bitwiseOR(c2), "bitwise_xor": lambda c1, c2: c1.bitwiseXOR(c2), - "copysign": pandas_udf( # type: ignore[call-overload] - lambda s1, s2: np.copysign(s1, s2), DoubleType() - ), + "copysign": _copysign_func, "float_power": lambda c1, c2: F.pow(c1.cast("double"), c2.cast("double")), "floor_divide": pandas_udf( # type: ignore[call-overload] lambda s1, s2: np.floor_divide(s1, s2), DoubleType() diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py b/python/pyspark/pandas/tests/test_numpy_compat.py index d95474f7fc071..6d2ac6a8e6c55 100644 --- a/python/pyspark/pandas/tests/test_numpy_compat.py +++ b/python/pyspark/pandas/tests/test_numpy_compat.py @@ -262,6 +262,50 @@ def test_np_fmax_fmin(self): expected = np_func(pdf.x1, pdf.x2) self.assert_eq(result, expected, almost=True) + def test_np_copysign(self): + for pdf in ( + pd.DataFrame( + { + "x1": [-64, -2, -1, 0, 1, 2, 64], + "x2": [2, -3, -2, -3, 3, -1, 2], + } + ), + pd.DataFrame( + { + "x1": [-np.inf, -64.0, -2.0, -0.0, 0.0, 2.0, 64.0, np.inf, np.nan, 1.0], + "x2": [2.0, -3.0, -2.0, 0.0, -0.0, -1.0, np.inf, -np.inf, 2.0, np.nan], + } + ), + pd.DataFrame( + { + "x1": pd.array([1, -2, 3, None, None], dtype="Int64"), + "x2": pd.array([-2, 3, None, 2, None], dtype="Int64"), + } + ), + ): + psdf = ps.from_pandas(pdf) + result = np.copysign(psdf.x1, psdf.x2) + expected = np.copysign(pdf.x1, pdf.x2) + self.assert_eq(result, expected, almost=True) + # copysign only differs from |x| in the sign bit, so assert on signbit + # explicitly -- 0.0 == -0.0 numerically and would hide a wrong sign. + self.assert_eq(np.signbit(result.to_pandas()), np.signbit(expected)) + + def test_np_copysign_signed_zero(self): + # np.copysign takes the sign from y's IEEE-754 sign bit, not from y < 0: + # copysign(1.0, -0.0) == -1.0 and copysign(1.0, 0.0) == 1.0. + pdf = pd.DataFrame( + { + "x1": [1.0, 1.0, -0.0, -0.0, 3.0], + "x2": [0.0, -0.0, 0.0, -0.0, -0.0], + } + ) + psdf = ps.from_pandas(pdf) + result = np.copysign(psdf.x1, psdf.x2).to_pandas() + expected = np.copysign(pdf.x1, pdf.x2) + self.assert_eq(result, expected) + self.assert_eq(np.signbit(result), np.signbit(expected)) + def test_np_heaviside(self): for pdf in ( pd.DataFrame({"x1": [-2, -1, 0, 1, 2], "x2": [-2, -1, 0, 1, 2]}),