From 06ead0734ff90e980e45ec80f8619ff775ba0126 Mon Sep 17 00:00:00 2001 From: Spenser Sun Date: Wed, 26 Aug 2026 07:10:03 +0000 Subject: [PATCH 1/4] [SPARK-58581][PS][FOLLOWUP] Fix NumPy floor_divide rounding and integer precision The native `floor_divide` mapping computed `floor(c1 / c2)`, which diverges from NumPy and pandas in two ways. Flooring the quotient is wrong when the division rounds up across an integer: `1.0 // 0.1` returned 10 where NumPy, pandas and Python return 9, and Spark's own `%` reports a remainder that no valid division leaves for that quotient. The floating branch now derives the quotient from the remainder, as npy_divmod does. Casting integral operands to double drops their low bits above 2**53, so `-9007199254740993 // 2` returned -4503599627370496 rather than -4503599627370497. Integral operands now divide in integer space, and only the result is cast back to double. The added test rows compare exactly rather than approximately, since a relative tolerance accepts an off-by-one at those magnitudes. --- python/pyspark/pandas/numpy_compat.py | 60 ++++++++++++++++++- .../pyspark/pandas/tests/test_numpy_compat.py | 39 +++++++++--- 2 files changed, 88 insertions(+), 11 deletions(-) diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index 33ec402a20128..b989a37407c78 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -177,10 +177,56 @@ 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. + # F.floor is unusable: a bigint cannot carry the infinities the caller's branches produce. + floor = quotient - F.pmod(quotient, F.lit(1.0)) + # Flooring goes one too low when the division landed just under the whole number. + return F.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((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"] + # Dispatched on type twice: floating operands need IEEE answers for infinities and signed + # zeros, and among the rest, at the end of the branch below, only integral operands can + # divide in integer space. return F.when( F.typeof(c1).isin("float", "double") | F.typeof(c2).isin("float", "double"), F.when(c1.isNull() | F.isnan(c1), c1_double) @@ -211,9 +257,11 @@ def _floor_divide_func(c1: Column, c2: Column) -> Column: .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(_floor_divide_floating(c1_double, c2_double)), ).otherwise( - # np.floor_divide on pandas Series returns IEEE values for an integral zero divisor. + # Non-floating operands. pandas masks a zero divisor and upcasts, so 1 // 0 is inf, + # -1 // 0 is -inf and 0 // 0 is nan. A nullable Int64 returns 0 instead, but both dtypes + # arrive as bigint and cannot be told apart, so this follows the default one. F.when(c1.isNull() | F.isnan(c1), c1_double) .when(c2.isNull() | F.isnan(c2), c2_double) .when( @@ -222,7 +270,13 @@ def _floor_divide_func(c1: Column, c2: Column) -> Column: .when(c1_double < 0, F.lit(float("-inf"))) .otherwise(F.lit(float("inf"))), ) - .otherwise((c1_double / c2_double) - F.pmod(c1_double / c2_double, F.lit(1.0))) + # A decimal column also lands here, and it cannot be named in a typeof test since the + # name carries its precision, so it falls through to the double casts. + .when( + F.typeof(c1).isin(integral_types) & F.typeof(c2).isin(integral_types), + _floor_divide_integral(c1, c2), + ) + .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..28ab3b6484dd5 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,28 @@ 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")) def test_np_logaddexp(self): for pdf in ( From 2773e485133cb91701c3ae7c6e748d8e2758ddd4 Mon Sep 17 00:00:00 2001 From: Spenser Sun Date: Wed, 26 Aug 2026 07:42:03 +0000 Subject: [PATCH 2/4] [SPARK-58581][PS][FOLLOWUP] Restore NumPy's quotient for the most negative long Integer division overflows for the most negative long divided by -1, the one quotient a long cannot hold, and Spark's `div` raises ARITHMETIC_OVERFLOW there. The mapping's earlier pandas UDF delegated to NumPy and so returned NumPy's wrapped value; the native conversion to a double divide returned it with the opposite sign instead. Guard that operand pair and return the wrapped value, which both restores the UDF's result and keeps the expression from raising. --- python/pyspark/pandas/numpy_compat.py | 6 +++++- python/pyspark/pandas/tests/test_numpy_compat.py | 5 +++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index b989a37407c78..fb60746d342cd 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -212,7 +212,11 @@ def _floor_divide_integral(c1: Column, c2: Column) -> Column: # 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 ( + 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") diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py b/python/pyspark/pandas/tests/test_numpy_compat.py index 28ab3b6484dd5..360aee8fb5064 100644 --- a/python/pyspark/pandas/tests/test_numpy_compat.py +++ b/python/pyspark/pandas/tests/test_numpy_compat.py @@ -425,6 +425,11 @@ def floor_divided(pdf): ) 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")) + def test_np_logaddexp(self): for pdf in ( pd.DataFrame( From 923c141d86359913accf783a3c03806e24c73936 Mon Sep 17 00:00:00 2001 From: Spenser Sun Date: Wed, 26 Aug 2026 19:46:52 +0000 Subject: [PATCH 3/4] [SPARK-58581][PS][FOLLOWUP] Keep an overflowing quotient in floor_divide A quotient that overflows to an infinity, such as `1e300 // 1e-300`, came back as nan. The floating branch emulates the floor with `q - pmod(q, 1.0)`, because `F.floor` returns a bigint and raises on an infinity, and pmod of an infinity is nan, so that subtraction left nan behind. An infinite quotient is its own floor, so it is now returned before that step. NumPy keeps the infinity here. The formula this PR replaced returned nan as well, so the case is not a regression from the earlier commits. --- python/pyspark/pandas/numpy_compat.py | 14 ++++++++++---- python/pyspark/pandas/tests/test_numpy_compat.py | 6 ++++++ 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index fb60746d342cd..3c0087d27226e 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -191,11 +191,17 @@ def _floor_divide_floating(c1: Column, c2: Column) -> Column: 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. - # F.floor is unusable: a bigint cannot carry the infinities the caller's branches produce. + # 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)) - # Flooring goes one too low when the division landed just under the whole number. - return F.when(quotient - floor > F.lit(0.5), floor + F.lit(1.0)).otherwise(floor) + 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: diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py b/python/pyspark/pandas/tests/test_numpy_compat.py index 360aee8fb5064..106664baecfeb 100644 --- a/python/pyspark/pandas/tests/test_numpy_compat.py +++ b/python/pyspark/pandas/tests/test_numpy_compat.py @@ -430,6 +430,12 @@ def floor_divided(pdf): 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 ( pd.DataFrame( From b954698bf7647645f9c8dee861b7e8a7dfa6059a Mon Sep 17 00:00:00 2001 From: Spenser Sun Date: Thu, 27 Aug 2026 22:00:55 +0000 Subject: [PATCH 4/4] [SPARK-58581][PS][FOLLOWUP] Dispatch floor_divide on the operand type Flatten the two-level dispatch in `_floor_divide_func` into one chain, as suggested in review. The null and zero-divisor cases answer the same way for every operand type, so they precede the type tests; the integral and floating branches then sit side by side, and a decimal falls through to the floating helper since `typeof` reports its precision and cannot be matched by name. Behaviour-preserving: the two spellings of the zero-divisor answer collapse into one, because a non-floating zero can never be negative. Verified against the previous version on 6361 rows covering every combination of double, float, bigint, int, tinyint, smallint and decimal operands, including negative zero divisors, infinities, values above 2**53 and -2**63 // -1. --- python/pyspark/pandas/numpy_compat.py | 65 ++++++++++----------------- 1 file changed, 24 insertions(+), 41 deletions(-) diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index 3c0087d27226e..59c111db8cfe3 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -234,58 +234,41 @@ def _floor_divide_func(c1: Column, c2: Column) -> Column: c2_double = c2.cast("double") integral_types = ["tinyint", "smallint", "int", "bigint"] - # Dispatched on type twice: floating operands need IEEE answers for infinities and signed - # zeros, and among the rest, at the end of the branch below, only integral operands can - # divide in integer space. - return F.when( - F.typeof(c1).isin("float", "double") | F.typeof(c2).isin("float", "double"), - F.when(c1.isNull() | F.isnan(c1), c1_double) - .when(c2.isNull() | F.isnan(c2), c2_double) - .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"))), - ) - .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)), - ) - .when( - 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"))), - ) - .when(c1_double == 0, c1_double / c2_double) - .otherwise(_floor_divide_floating(c1_double, c2_double)), - ).otherwise( - # Non-floating operands. pandas masks a zero divisor and upcasts, so 1 // 0 is inf, - # -1 // 0 is -inf and 0 // 0 is nan. A nullable Int64 returns 0 instead, but both dtypes - # arrive as bigint and cannot be told apart, so this follows the default one. + 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( c2_double == 0, F.when(c1_double == 0, F.lit(float("nan"))) - .when(c1_double < 0, F.lit(float("-inf"))) + .when((c1_double < 0) != (c2_double.cast("string") == "-0.0"), F.lit(float("-inf"))) .otherwise(F.lit(float("inf"))), ) - # A decimal column also lands here, and it cannot be named in a typeof test since the - # name carries its precision, so it falls through to the double casts. + # Integral operands divide as integers, so operands above 2**53 keep their low bits. .when( 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( + 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( + 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)), + ) + # 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)), + ) + # A decimal cannot be matched by name: typeof reports its precision, as in decimal(10,2). .otherwise(_floor_divide_floating(c1_double, c2_double)) )