diff --git a/python/pyspark/pandas/data_type_ops/num_ops.py b/python/pyspark/pandas/data_type_ops/num_ops.py index 8baeacf3b0187..c1926d4d68781 100644 --- a/python/pyspark/pandas/data_type_ops/num_ops.py +++ b/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) diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index aaa7d843cd75c..c365fa63a5d3e 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/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 diff --git a/python/pyspark/pandas/tests/data_type_ops/test_num_mul_div.py b/python/pyspark/pandas/tests/data_type_ops/test_num_mul_div.py index d272b22cc475c..2bb2d024f3c2f 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_num_mul_div.py +++ b/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 diff --git a/python/pyspark/pandas/tests/data_type_ops/test_num_reverse.py b/python/pyspark/pandas/tests/data_type_ops/test_num_reverse.py index f2b85085e7d1b..83e52e45f6ce0 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_num_reverse.py +++ b/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: diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py b/python/pyspark/pandas/tests/test_numpy_compat.py index a733155b7199d..9a2c0ca76e3d3 100644 --- a/python/pyspark/pandas/tests/test_numpy_compat.py +++ b/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) diff --git a/python/pyspark/pandas/utils.py b/python/pyspark/pandas/utils.py index 7d36f8431b813..6a3eb14daac92 100644 --- a/python/pyspark/pandas/utils.py +++ b/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