[spark] branch master updated: [SPARK-37388][SQL] Fix NPE in WidthBucket in WholeStageCodegenExec

[we...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

wenchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new 77f3974  [SPARK-37388][SQL] Fix NPE in WidthBucket in WholeStageCodegenExec
77f3974 is described below

commit 77f397464823a547c27d98ed306703cb9c73cec3
Author: Tom van Bussel <to...@databricks.com>
AuthorDate: Mon Nov 22 16:18:30 2021 +0800

    [SPARK-37388][SQL] Fix NPE in WidthBucket in WholeStageCodegenExec
    
    ### What changes were proposed in this pull request?
    
    This PR fixes a `NullPointerException` in `WholeStageCodegenExec` caused by the expression `WidthBucket`. The cause of this NPE is that `WidthBucket` calls `WidthBucket.computeBucketNumber`, which can return `null`, but the generated code cannot deal with `null`s.
    
    ### Why are the changes needed?
    
    This fixes a `NullPointerException` in Spark SQL.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No
    
    ### How was this patch tested?
    
    Added tests to `MathExpressionsSuite`. This suite already had tests for `WidthBucket` with interval inputs, but lacked tests with double inputs. I checked that the tests failed without the fix, and succeed with the fix.
    
    Closes #34670 from tomvanbussel/SPARK-37388.
    
    Authored-by: Tom van Bussel <to...@databricks.com>
    Signed-off-by: Wenchen Fan <we...@databricks.com>
---
 .../sql/catalyst/expressions/mathExpressions.scala | 34 +++++++++++++++++-----
 .../expressions/MathExpressionsSuite.scala         | 28 ++++++++++++++++++
 2 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala
index e3541dc..03f9da6 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala
@@ -1547,14 +1547,27 @@ case class BRound(child: Expression, scale: Expression)
 }
 
 object WidthBucket {
-
   def computeBucketNumber(value: Double, min: Double, max: Double, numBucket: Long): jl.Long = {
-    if (numBucket <= 0 || numBucket == Long.MaxValue || jl.Double.isNaN(value) || min == max ||
-        jl.Double.isNaN(min) || jl.Double.isInfinite(min) ||
-        jl.Double.isNaN(max) || jl.Double.isInfinite(max)) {
-      return null
+    if (isNull(value, min, max, numBucket)) {
+      null
+    } else {
+      computeBucketNumberNotNull(value, min, max, numBucket)
     }
+  }
 
+  /** This function is called by generated Java code, so it needs to be public. */
+  def isNull(value: Double, min: Double, max: Double, numBucket: Long): Boolean = {
+    numBucket <= 0 ||
+      numBucket == Long.MaxValue ||
+      jl.Double.isNaN(value) ||
+      min == max ||
+      jl.Double.isNaN(min) || jl.Double.isInfinite(min) ||
+      jl.Double.isNaN(max) || jl.Double.isInfinite(max)
+  }
+
+  /** This function is called by generated Java code, so it needs to be public. */
+  def computeBucketNumberNotNull(
+      value: Double, min: Double, max: Double, numBucket: Long): jl.Long = {
     val lower = Math.min(min, max)
     val upper = Math.max(min, max)
 
@@ -1666,9 +1679,14 @@ case class WidthBucket(
   }
 
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
-    defineCodeGen(ctx, ev, (input, min, max, numBucket) =>
-      "org.apache.spark.sql.catalyst.expressions.WidthBucket" +
-        s".computeBucketNumber($input, $min, $max, $numBucket)")
+    nullSafeCodeGen(ctx, ev, (input, min, max, numBucket) => {
+      s"""${ev.isNull} = org.apache.spark.sql.catalyst.expressions.WidthBucket
+         |  .isNull($input, $min, $max, $numBucket);
+         |if (!${ev.isNull}) {
+         |  ${ev.value} = org.apache.spark.sql.catalyst.expressions.WidthBucket
+         |    .computeBucketNumberNotNull($input, $min, $max, $numBucket);
+         |}""".stripMargin
+    })
   }
 
   override def first: Expression = value
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MathExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MathExpressionsSuite.scala
index bd133e75..ea0d619 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MathExpressionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MathExpressionsSuite.scala
@@ -760,4 +760,32 @@ class MathExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
       checkEvaluation(WidthBucket(Literal(v), Literal(s), Literal(e), Literal(n)), expected)
     }
   }
+
+  test("SPARK-37388: width_bucket") {
+    val nullDouble = Literal.create(null, DoubleType)
+    val nullLong = Literal.create(null, LongType)
+
+    checkEvaluation(WidthBucket(5.35, 0.024, 10.06, 5L), 3L)
+    checkEvaluation(WidthBucket(-2.1, 1.3, 3.4, 3L), 0L)
+    checkEvaluation(WidthBucket(8.1, 0.0, 5.7, 4L), 5L)
+    checkEvaluation(WidthBucket(-0.9, 5.2, 0.5, 2L), 3L)
+    checkEvaluation(WidthBucket(nullDouble, 0.024, 10.06, 5L), null)
+    checkEvaluation(WidthBucket(5.35, nullDouble, 10.06, 5L), null)
+    checkEvaluation(WidthBucket(5.35, 0.024, nullDouble, 5L), null)
+    checkEvaluation(WidthBucket(5.35, nullDouble, nullDouble, 5L), null)
+    checkEvaluation(WidthBucket(5.35, 0.024, 10.06, nullLong), null)
+    checkEvaluation(WidthBucket(nullDouble, nullDouble, nullDouble, nullLong), null)
+    checkEvaluation(WidthBucket(5.35, 0.024, 10.06, -5L), null)
+    checkEvaluation(WidthBucket(5.35, 0.024, 10.06, Long.MaxValue), null)
+    checkEvaluation(WidthBucket(Double.NaN, 0.024, 10.06, 5L), null)
+    checkEvaluation(WidthBucket(Double.NegativeInfinity, 0.024, 10.06, 5L), 0L)
+    checkEvaluation(WidthBucket(Double.PositiveInfinity, 0.024, 10.06, 5L), 6L)
+    checkEvaluation(WidthBucket(5.35, 0.024, 0.024, 5L), null)
+    checkEvaluation(WidthBucket(5.35, Double.NaN, 10.06, 5L), null)
+    checkEvaluation(WidthBucket(5.35, Double.NegativeInfinity, 10.06, 5L), null)
+    checkEvaluation(WidthBucket(5.35, Double.PositiveInfinity, 10.06, 5L), null)
+    checkEvaluation(WidthBucket(5.35, 0.024, Double.NaN, 5L), null)
+    checkEvaluation(WidthBucket(5.35, 0.024, Double.NegativeInfinity, 5L), null)
+    checkEvaluation(WidthBucket(5.35, 0.024, Double.PositiveInfinity, 5L), null)
+  }
 }

---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org