From 41489eff3938d99827f384337a7153e2a4347acb Mon Sep 17 00:00:00 2001 From: David Mollitor Date: Fri, 18 Sep 2026 03:16:28 +0000 Subject: [PATCH] [SPARK-59696][SQL] Remove redundant second modulo in integral MathUtils.pmod The integral `MathUtils.pmod` overloads (Int, Long, Byte, Short) compute `if (r < 0) (r + n) % n else r`, where `r = a % n`. When the divisor `n > 0`, a negative `r` lies in `(-n, 0)`, so `r + n` is already in `[0, n)` and the trailing `% n` is a no-op; it is skipped there. When `n < 0`, `r + n` can fall below `n` and still needs reducing, so the `% n` is retained: if (r >= 0) r else if (n > 0) r + n else (r + n) % n This is behavior-preserving for every input, including negative divisors. An unconditional drop of the `% n` would change e.g. `pmod(-3, -5)` from `-3` to `-8` (out of range) and diverge from the (unchanged) decimal `pmod` path. The guarded form was verified equivalent to the original against an exhaustive oracle: all byte and short (a, n) pairs, and all 2^32 int dividends across positive and negative divisors. Added regression cases in `ArithmeticExpressionSuite` for a negative dividend with a negative divisor -- the `r < 0`, `n < 0` path that an unconditional drop would break and that existing tests did not cover. The float/double overloads intentionally keep the second modulo: floating-point rounding can make `r + n` round up to exactly `n`, which `% n` folds back to 0. `pmod` backs the `pmod` SQL function and `HashPartitioning`'s per-row partition-id expression; for positive divisors this removes an integer division from the hot path (~3.3x faster on the changed branch in a JMH microbenchmark; numbers in the PR description). Co-authored-by: Isaac --- .../spark/sql/catalyst/util/MathUtils.scala | 19 +++++++++++++++---- .../ArithmeticExpressionSuite.scala | 8 ++++++++ 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/MathUtils.scala b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/MathUtils.scala index e318a60ce99ab..271574079cf6a 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/MathUtils.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/MathUtils.scala @@ -92,25 +92,36 @@ object MathUtils { // Positive modulo (`pmod`): the remainder `a % n` adjusted to share the sign of `n`. // Unlike `floorMod`, this matches the `pmod` SQL function / `HashPartitioning` semantics. // Shared by `Pmod`'s eval and codegen paths so the two never diverge. + // + // The `r < 0` branch shifts the remainder by `n`. When `n > 0`, `r` lies in `(-n, 0)` so the + // shifted value `r + n` is already in `[0, n)` and the extra `% n` is a no-op -- it is skipped. + // When `n < 0`, `r + n` can fall below `n`, so the `% n` is retained to preserve the original + // result. The float/double overloads always keep it because `r + n` can round up to exactly `n`. def pmod(a: Int, n: Int): Int = { val r = a % n - if (r < 0) (r + n) % n else r + if (r >= 0) r + else if (n > 0) r + n + else (r + n) % n } def pmod(a: Long, n: Long): Long = { val r = a % n - if (r < 0) (r + n) % n else r + if (r >= 0) r + else if (n > 0) r + n + else (r + n) % n } def pmod(a: Byte, n: Byte): Byte = { val r = a % n - if (r < 0) ((r + n) % n).toByte else r.toByte + val result = if (r >= 0) r else if (n > 0) r + n else (r + n) % n + result.toByte } def pmod(a: Short, n: Short): Short = { val r = a % n - if (r < 0) ((r + n) % n).toShort else r.toShort + val result = if (r >= 0) r else if (n > 0) r + n else (r + n) % n + result.toShort } def pmod(a: Float, n: Float): Float = { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ArithmeticExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ArithmeticExpressionSuite.scala index f8cdf825bc4c8..60b09fa08f798 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ArithmeticExpressionSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ArithmeticExpressionSuite.scala @@ -582,6 +582,14 @@ class ArithmeticExpressionSuite extends SparkFunSuite with ExpressionEvalHelper checkEvaluation(Pmod(positiveShort, negativeShort), positiveShort.toShort) checkEvaluation(Pmod(positiveInt, negativeInt), positiveInt) checkEvaluation(Pmod(positiveLong, negativeLong), positiveLong) + // Negative divisor (n < 0): `pmod` is only positive for a positive divisor, so these expected + // values are intentionally <= 0 (released behavior). They guard the r < 0, n < 0 path where + // `r + n` still needs `% n` -- dropping it goes out of range (pmod(-3, -5) would be -8). + checkEvaluation(Pmod(Literal(-3), Literal(-5)), -3) + checkEvaluation(Pmod(Literal(-7), Literal(-3)), -1) + checkEvaluation(Pmod(Literal(-3L), Literal(-5L)), -3L) + checkEvaluation(Pmod(Literal((-3).toShort), Literal((-5).toShort)), (-3).toShort) + checkEvaluation(Pmod(Literal((-7).toByte), Literal((-3).toByte)), (-1).toByte) Seq("true", "false").foreach { failOnError => withSQLConf(SQLConf.ANSI_ENABLED.key -> failOnError) {