zhengruifeng commented on code in PR #58306:
URL: https://github.com/apache/spark/pull/58306#discussion_r3867874946
##########
python/pyspark/pandas/numpy_compat.py:
##########
@@ -177,10 +177,66 @@ 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"]
+ # Dispatched on type twice: floating operands need IEEE answers for
infinities and signed
Review Comment:
what about dispatching earlier by datatype?
when(is int, _floor_divide_integral)
.when(is float, _floor_divide_floating)
.otherwise
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
