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) {