diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index 33ec402a20128..59c111db8cfe3 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -177,52 +177,99 @@ def _logaddexp_func(c1: Column, c2: Column, base2: bool = False) -> Column: ) +def _floor_divide_floating(c1: Column, c2: Column) -> Column: + """Return floor(c1 / c2) for finite non-zero double operands, derived from the remainder. + + Flooring the quotient is wrong when the division rounds up across an integer: 1.0 / 0.1 + rounds to exactly 10.0, so its floor is 10 where NumPy, pandas and Python return 9. A + remainder is exact, so NumPy's npy_divmod derives the quotient from it, as this does. + """ + remainder = F.try_mod(c1, c2) + # The remainder carries the dividend's sign, so this is the truncating quotient. + truncated = (c1 - remainder) / c2 + # Truncating and flooring differ by one on opposite signs with a remainder left over. + quotient = F.when( + (remainder != 0) & ((remainder < 0) != (c2 < 0)), truncated - F.lit(1.0) + ).otherwise(truncated) + # The quotient is whole in exact arithmetic, but the division can leave it a few bits off, so + # round it back. F.floor cannot do this: it returns a bigint, which raises on an infinity. + floor = quotient - F.pmod(quotient, F.lit(1.0)) + return ( + # An infinite quotient is its own floor, and has to be returned before the line above + # is used, since pmod of an infinity is nan and leaves `floor` nan. + F.when(quotient.isin(float("inf"), float("-inf")), quotient) + # Flooring goes one too low when the division landed just under the whole number. + .when(quotient - floor > F.lit(0.5), floor + F.lit(1.0)) + .otherwise(floor) + ) + + +def _floor_divide_integral(c1: Column, c2: Column) -> Column: + """Return floor(c1 / c2) for integral operands, keeping the quotient in integer space. + + Casting an operand above 2**53 to double drops its low bits, turning 9007199254740993 into + 9007199254740992, and Spark's `/` always divides as double. The long casts are no-ops for + the integral types the caller admits; they are there because `div` rejects a double even in + a branch the guard turns off. + """ + c1_long = c1.cast("long") + c2_long = c2.cast("long") + # `div` is integer division, truncating toward zero, so it needs the same flooring + # correction as the floating helper. Integer arithmetic cannot round, so nothing more. + truncated = F.call_function("div", c1_long, c2_long) + remainder = F.try_mod(c1_long, c2_long) + return F.when( + # The one quotient a long cannot hold, where NumPy wraps around and `div` would raise. + (c1_long == F.lit(-(2**63))) & (c2_long == F.lit(-1)), + F.lit(float(-(2**63))), + ).otherwise( + F.when((remainder != 0) & ((remainder < 0) != (c2_long < 0)), truncated - F.lit(1)) + .otherwise(truncated) + .cast("double") + ) + + def _floor_divide_func(c1: Column, c2: Column) -> Column: c1_double = c1.cast("double") c2_double = c2.cast("double") + integral_types = ["tinyint", "smallint", "int", "bigint"] - return F.when( - F.typeof(c1).isin("float", "double") | F.typeof(c2).isin("float", "double"), + return ( + # Null, nan and a zero divisor are handled the same way for every operand type. F.when(c1.isNull() | F.isnan(c1), c1_double) .when(c2.isNull() | F.isnan(c2), c2_double) + # pandas upcasts a zero divisor instead of raising. A negative zero divisor negates the + # result, and no comparison can see that sign, so the string form is used. A nullable Int64 + # returns 0 instead, but arrives as bigint like a default int64, whose answer this follows. .when( - c1_double.isin(float("-inf"), float("inf")), - F.when( - c2_double == 0, - F.when( - (c1_double < 0) != (c2_double.cast("string") == "-0.0"), - F.lit(float("-inf")), - ).otherwise(F.lit(float("inf"))), - ).otherwise(F.lit(float("nan"))), + c2_double == 0, + F.when(c1_double == 0, F.lit(float("nan"))) + .when((c1_double < 0) != (c2_double.cast("string") == "-0.0"), F.lit(float("-inf"))) + .otherwise(F.lit(float("inf"))), ) + # Integral operands divide as integers, so operands above 2**53 keep their low bits. .when( - c2_double.isin(float("-inf"), float("inf")), - F.when(c1_double == 0, c1_double / c2_double) - .when((c1_double < 0) != (c2_double < 0), F.lit(-1.0)) - .otherwise(F.lit(0.0)), + F.typeof(c1).isin(integral_types) & F.typeof(c2).isin(integral_types), + _floor_divide_integral(c1, c2), ) + # Only floating operands can be infinite or a negative zero, handled before the division. .when( - c2_double == 0, - F.when(c1_double == 0, F.lit(float("nan"))) + F.typeof(c1).isin("float", "double") | F.typeof(c2).isin("float", "double"), + # An infinite dividend has no remainder, so NumPy's quotient is nan for any divisor. + F.when(c1_double.isin(float("-inf"), float("inf")), F.lit(float("nan"))) + # An infinite divisor gives a quotient between -1 and 1, so the floor is 0 or -1. .when( - (c1_double < 0) != (c2_double.cast("string") == "-0.0"), - F.lit(float("-inf")), + c2_double.isin(float("-inf"), float("inf")), + F.when(c1_double == 0, c1_double / c2_double) + .when((c1_double < 0) != (c2_double < 0), F.lit(-1.0)) + .otherwise(F.lit(0.0)), ) - .otherwise(F.lit(float("inf"))), - ) - .when(c1_double == 0, c1_double / c2_double) - .otherwise((c1_double / c2_double) - F.pmod(c1_double / c2_double, F.lit(1.0))), - ).otherwise( - # np.floor_divide on pandas Series returns IEEE values for an integral zero divisor. - F.when(c1.isNull() | F.isnan(c1), c1_double) - .when(c2.isNull() | F.isnan(c2), c2_double) - .when( - c2_double == 0, - F.when(c1_double == 0, F.lit(float("nan"))) - .when(c1_double < 0, F.lit(float("-inf"))) - .otherwise(F.lit(float("inf"))), + # Dividing a zero dividend keeps its sign, so -0.0 // 3.0 is -0.0. + .when(c1_double == 0, c1_double / c2_double) + .otherwise(_floor_divide_floating(c1_double, c2_double)), ) - .otherwise((c1_double / c2_double) - F.pmod(c1_double / c2_double, F.lit(1.0))) + # A decimal cannot be matched by name: typeof reports its precision, as in decimal(10,2). + .otherwise(_floor_divide_floating(c1_double, c2_double)) ) diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py b/python/pyspark/pandas/tests/test_numpy_compat.py index a3f92ef720583..106664baecfeb 100644 --- a/python/pyspark/pandas/tests/test_numpy_compat.py +++ b/python/pyspark/pandas/tests/test_numpy_compat.py @@ -339,6 +339,15 @@ def test_np_modf(self): def test_floor_divide_func(self): from pyspark.pandas.numpy_compat import _floor_divide_func + def floor_divided(pdf): + psdf = ps.from_pandas(pdf) + return ( + psdf.spark.frame() + .select(_floor_divide_func(F.col("x1"), F.col("x2")).alias("result")) + .toPandas()["result"] + .rename(None) + ) + for pdf in ( pd.DataFrame( { @@ -393,14 +402,39 @@ def test_floor_divide_func(self): } ), ): - psdf = ps.from_pandas(pdf) - result = ( - psdf.spark.frame() - .select(_floor_divide_func(F.col("x1"), F.col("x2")).alias("result")) - .toPandas()["result"] - .rename(None) - ) - self.assert_eq(result, np.floor_divide(pdf.x1, pdf.x2), almost=True) + self.assert_eq(floor_divided(pdf), np.floor_divide(pdf.x1, pdf.x2), almost=True) + + # Divisors binary cannot represent exactly, where the quotient rounds up across an + # integer: 1.0 / 0.1 rounds to 10.0, so flooring it gives 10 instead of 9. Compared + # exactly, since almost=True would accept an off-by-one on the large values below. + pdf = pd.DataFrame( + { + "x1": [1.0, 10.0, 2.0, 0.5, 7.0, -1.0, -10.0, 3.0], + "x2": [0.1, 0.1, 0.2, 0.1, 0.7, 0.1, 0.1, 7.0], + } + ) + self.assert_eq(floor_divided(pdf), np.floor_divide(pdf.x1, pdf.x2)) + + # Integral operands above 2**53, where casting an operand to double would drop its + # low bits: -9007199254740993 // 2 is -4503599627370497, not -4503599627370496. + pdf = pd.DataFrame( + { + "x1": [9007199254740993, -9007199254740993, 4611686018427387905, 7, -7], + "x2": [1, 2, 3, 3, 3], + } + ) + self.assert_eq(floor_divided(pdf), np.floor_divide(pdf.x1, pdf.x2).astype("float64")) + + # The most negative long divided by -1, whose quotient a long cannot hold. NumPy wraps + # around, while Spark's integer division raises. + pdf = pd.DataFrame({"x1": [-(2**63), -(2**63)], "x2": [-1, 2]}) + self.assert_eq(floor_divided(pdf), np.floor_divide(pdf.x1, pdf.x2).astype("float64")) + + # Finite operands whose quotient overflows to an infinity, which is its own floor. + pdf = pd.DataFrame( + {"x1": [1e300, -1e300, 1e300, -1e300], "x2": [1e-300, 1e-300, -1e-300, -1e-300]} + ) + self.assert_eq(floor_divided(pdf), np.floor_divide(pdf.x1, pdf.x2)) def test_np_logaddexp(self): for pdf in (