[SPARK-59093][PS] Fix incorrect results and arithmetic errors in floor division

python/pyspark/pandas/data_type_ops/num_ops.py

@@ -17,7 +17,7 @@
 
 import decimal
 import numbers
-from typing import Any, Callable, Union, cast
+from typing import Any, Union, cast
 
 import numpy as np
 import pandas as pd
@@ -52,7 +52,7 @@
     handle_dtype_as_extension_dtype,
     pandas_on_spark_type,
 )
-from pyspark.pandas.utils import is_ansi_mode_enabled
+from pyspark.pandas.utils import _floor_divide_func, is_ansi_mode_enabled
 from pyspark.sql import Column as PySparkColumn
 from pyspark.sql import functions as F
 from pyspark.sql.types import (
@@ -373,23 +373,9 @@ def floordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
         _sanitize_list_like(right)
         if not is_valid_operand_for_numeric_arithmetic(right):
             raise TypeError("Floor division can not be applied to given types.")
-        spark_session = left._internal.spark_frame.sparkSession
-        use_try_divide = is_ansi_mode_enabled(spark_session)
-
-        def fallback_div(x: PySparkColumn, y: PySparkColumn) -> PySparkColumn:
-            return x.__div__(y)
-
-        safe_div: Callable[[PySparkColumn, PySparkColumn], PySparkColumn] = (
-            F.try_divide if use_try_divide else fallback_div
-        )
 
         def floordiv(left: PySparkColumn, right: Any) -> PySparkColumn:
-            return F.when(F.lit(right is np.nan), np.nan).otherwise(
-                F.when(
-                    F.lit(right != 0) | F.lit(right).isNull(),
-                    F.floor(left.__div__(right)),
-                ).otherwise(safe_div(F.lit(np.inf), left))
-            )
+            return _floor_divide_func(left, F.lit(right))
 
         right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type)
         return numpy_column_op(floordiv)(left, right)
@@ -413,9 +399,7 @@ def rfloordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
             raise TypeError("Floor division can not be applied to given types.")
 
         def rfloordiv(left: PySparkColumn, right: Any) -> PySparkColumn:
-            return F.when(F.lit(left == 0), F.lit(np.inf).__div__(right)).otherwise(
-                F.floor(F.lit(right).__div__(left))
-            )
+            return _floor_divide_func(F.lit(right), left)
 
         right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type)
         return numpy_column_op(rfloordiv)(left, right)
@@ -502,25 +486,8 @@ def floordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
             raise TypeError("Floor division can not be applied to given types.")
         left_dtype = left.dtype
 
-        def fallback_div(x: PySparkColumn, y: PySparkColumn) -> PySparkColumn:
-            return x.__div__(y)
-
-        safe_div: Callable[[PySparkColumn, PySparkColumn], PySparkColumn] = (
-            F.try_divide if is_ansi else fallback_div
-        )
-
         def floordiv(lc: PySparkColumn, rc: Any) -> PySparkColumn:
-            expr = F.when(F.lit(rc is np.nan), np.nan).otherwise(
-                F.when(
-                    F.lit(rc != 0) | F.lit(rc).isNull(),
-                    F.floor(lc.__div__(rc)),
-                ).otherwise(
-                    F.when(F.lit(lc == np.inf) | F.lit(lc == -np.inf), lc).otherwise(
-                        safe_div(F.lit(np.inf), lc)
-                    )
-                )
-            )
-            return _cast_back_float(expr, left_dtype, right)
+            return _cast_back_float(_floor_divide_func(lc, F.lit(rc)), left_dtype, right)
 
         new_right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type)
         return numpy_column_op(floordiv)(left, new_right)
@@ -552,9 +519,7 @@ def rfloordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
             raise TypeError("Floor division can not be applied to given types.")
 
         def rfloordiv(left: PySparkColumn, right: Any) -> PySparkColumn:
-            return F.when(F.lit(left == 0), F.lit(np.inf).__div__(right)).otherwise(
-                F.when(F.lit(left) == np.nan, np.nan).otherwise(F.floor(F.lit(right).__div__(left)))
-            )
+            return _floor_divide_func(F.lit(right), left)
 
         right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type)
         return numpy_column_op(rfloordiv)(left, right)

python/pyspark/pandas/numpy_compat.py

@@ -21,6 +21,7 @@
 from pyspark.loose_version import LooseVersion
 from pyspark.pandas._typing import SeriesOrIndex
 from pyspark.pandas.base import IndexOpsMixin
+from pyspark.pandas.utils import _floor_divide_func
 from pyspark.sql import Column
 from pyspark.sql import functions as F
 from pyspark.sql.pandas.functions import pandas_udf
@@ -187,102 +188,6 @@ 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 (
-        # 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) != (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(
-            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))
-    )
-
-
 # NumPy 2.3.0 changed how fmax/fmin break a signed-zero tie: for equal operands
 # (for example +0.0 and -0.0) it returns the first operand, while older versions
 # returned the second. Track the installed NumPy so the result keeps the matching

python/pyspark/pandas/tests/data_type_ops/test_num_mul_div.py

@@ -118,6 +118,51 @@ def test_floordiv(self):
             self.assertRaises(TypeError, lambda: psdf["decimal"] // 0.1)
             self.assertRaises(TypeError, lambda: 0.1 // psdf["decimal"])
 
+        # A divisor that is not exactly representable: 1.0 / 0.1 rounds up to exactly 10.0,
+        # so flooring the quotient gives 10 where pandas returns 9.
+        pser = pd.Series([1.0, 10.0, -1.0, 2.5])
+        psser = ps.from_pandas(pser)
+        self.assert_eq(pser // 0.1, psser // 0.1)
+
+        # An integral operand above 2**53 loses its low bits when divided as double. pandas
+        # returns int64 here, which Spark's division cannot, so compare the values.
+        pser = pd.Series([9007199254740993, -9007199254740993])
+        psser = ps.from_pandas(pser)
+        self.assert_eq((pser // 2).astype(float), psser // 2)
+        self.assert_eq((pser // 3).astype(float), psser // 3)
+
+        # An infinite divisor leaves a quotient of 0 or -1, and an infinite dividend has no
+        # finite floor, which pandas reports as nan.
+        pser = pd.Series([1.0, -1.0, np.inf, -np.inf])
+        psser = ps.from_pandas(pser)
+        self.assert_eq(pser // np.inf, psser // np.inf)
+        self.assert_eq(pser // -np.inf, psser // -np.inf)
+        pser = pd.Series([np.inf, -np.inf, np.nan, 1.0])
+        psser = ps.from_pandas(pser)
+        self.assert_eq(pser // 2.0, psser // 2.0)
+
+        # Finite operands whose quotient overflows to an infinity, which is its own floor.
+        edge_pdf = pd.DataFrame({"a": [1e300, -1e300], "b": [1e-300, 1e-300]})
+        edge_psdf = ps.from_pandas(edge_pdf)
+        self.assert_eq(edge_pdf.a // edge_pdf.b, edge_psdf.a // edge_psdf.b)
+
+        # A negative zero divisor negates the result, and a zero dividend keeps its own sign.
+        pser = pd.Series([1.0, -1.0, 2.5])
+        psser = ps.from_pandas(pser)
+        self.assert_eq(pser // -0.0, psser // -0.0)
+        pser = pd.Series([-0.0, 0.0])
+        psser = ps.from_pandas(pser)
+        self.assert_eq(pser // 3.0, psser // 3.0)
+        # An equality check cannot see the sign of a zero, so compare it directly.
+        self.assertEqual(
+            np.signbit(pser // 3.0).tolist(), np.signbit((psser // 3.0).to_pandas()).tolist()
+        )
+
+        # The only quotient that does not fit in a long, where pandas wraps around.
+        pser = pd.Series([-(2**63)])
+        psser = ps.from_pandas(pser)
+        self.assert_eq((pser // -1).astype(float), psser // -1)
+
     def test_mod(self):
         pdf, psdf = self.pdf, self.psdf
 

python/pyspark/pandas/tests/data_type_ops/test_num_reverse.py

@@ -103,6 +103,28 @@ def test_rfloordiv(self):
             self.assertRaises(TypeError, lambda: datetime.date(1994, 1, 1) // psser)
             self.assertRaises(TypeError, lambda: datetime.datetime(1994, 1, 1) // psser)
 
+        # A divisor that is not exactly representable: 1.0 / 0.1 rounds up to exactly 10.0,
+        # so flooring the quotient gives 10 where pandas returns 9.
+        pser = pd.Series([0.1, 0.3, 0.7])
+        psser = ps.from_pandas(pser)
+        self.assert_eq(1.0 // pser, 1.0 // psser)
+
+        # An integral operand above 2**53 loses its low bits when divided as double. pandas
+        # returns int64 here, which Spark's division cannot, so compare the values.
+        pser = pd.Series([2, -2])
+        psser = ps.from_pandas(pser)
+        self.assert_eq((9007199254740993 // pser).astype(float), 9007199254740993 // psser)
+
+        # A zero dividend and a zero divisor, which pandas reports as nan.
+        pser = pd.Series([0, 2])
+        psser = ps.from_pandas(pser)
+        self.assert_eq(0 // pser, 0 // psser)
+
+        # A quotient that overflows to an infinity, which is its own floor.
+        pser = pd.Series([1e-300, -1e-300])
+        psser = ps.from_pandas(pser)
+        self.assert_eq(1e300 // pser, 1e300 // psser)
+
     def test_rpow(self):
         pdf, psdf = self.pdf, self.psdf
         for col in self.numeric_df_cols:

python/pyspark/pandas/tests/test_numpy_compat.py

@@ -363,7 +363,7 @@ def test_np_modf(self):
         self.assert_eq(ps_integral, pd_integral, almost=True)
 
     def test_floor_divide_func(self):
-        from pyspark.pandas.numpy_compat import _floor_divide_func
+        from pyspark.pandas.utils import _floor_divide_func
 
         def floor_divided(pdf):
             psdf = ps.from_pandas(pdf)

python/pyspark/pandas/utils.py

@@ -1046,6 +1046,102 @@ def compare_allow_null(
     return left.isNull() | right.isNull() | comp(left, right)
 
 
+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 (
+        # 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) != (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(
+            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))
+    )
+
+
 def log_advice(message: str) -> None:
     """
     Display advisory logs for functions to be aware of when using pandas API on Spark