[SPARK-36753][SQL] ArrayExcept handle duplicated Double.NaN and Float.NaN

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala

@@ -4073,14 +4073,20 @@ case class ArrayExcept(left: Expression, right: Expression) extends ArrayBinaryL
     if (TypeUtils.typeWithProperEquals(elementType)) {
       (array1, array2) =>
         val hs = new OpenHashSet[Any]
+        val isNaN = SQLOpenHashSet.isNaN(elementType)
         var notFoundNullElement = true
+        var notFoundNaNElement = true
         var i = 0
         while (i < array2.numElements()) {
           if (array2.isNullAt(i)) {
             notFoundNullElement = false
           } else {
             val elem = array2.get(i, elementType)
-            hs.add(elem)
+            if (isNaN(elem)) {
+              notFoundNaNElement = false
+            } else {
+              hs.add(elem)
+            }
           }
           i += 1
         }
@@ -4094,9 +4100,16 @@ case class ArrayExcept(left: Expression, right: Expression) extends ArrayBinaryL
             }
           } else {
             val elem = array1.get(i, elementType)
-            if (!hs.contains(elem)) {
-              arrayBuffer += elem
-              hs.add(elem)
+            if (isNaN(elem)) {
+              if (notFoundNaNElement) {
+                arrayBuffer += elem
+                notFoundNaNElement = false
+              }
+            } else {
+              if (!hs.contains(elem)) {
+                arrayBuffer += elem
+                hs.add(elem)
+              }
             }
           }
           i += 1
@@ -4168,13 +4181,20 @@ case class ArrayExcept(left: Expression, right: Expression) extends ArrayBinaryL
       nullSafeCodeGen(ctx, ev, (array1, array2) => {
         val notFoundNullElement = ctx.freshName("notFoundNullElement")
         val nullElementIndex = ctx.freshName("nullElementIndex")
+        val notFoundNaNElement = ctx.freshName("notFoundNaNElement")
         val builder = ctx.freshName("builder")
         val openHashSet = classOf[OpenHashSet[_]].getName
         val classTag = s"scala.reflect.ClassTag$$.MODULE$$.$hsTypeName()"
         val hashSet = ctx.freshName("hashSet")
         val arrayBuilder = classOf[mutable.ArrayBuilder[_]].getName
         val arrayBuilderClass = s"$arrayBuilder$$of$ptName"
 
+        val isNaNMethod = elementType match {
+          case DoubleType => Some(s"java.lang.Double.isNaN((double)$value)")
+          case FloatType => Some(s"java.lang.Float.isNaN((float)$value)")
+          case _ => None
+        }
+
         def withArray2NullCheck(body: String): String =
           if (right.dataType.asInstanceOf[ArrayType].containsNull) {
             if (left.dataType.asInstanceOf[ArrayType].containsNull) {
@@ -4197,11 +4217,21 @@ case class ArrayExcept(left: Expression, right: Expression) extends ArrayBinaryL
             body
           }
 
+        def withArray2NaNCheck(body: String): String = {
+          isNaNMethod.map { isNaN =>
+            s"""
+               |if ($isNaN) {
+               |  $notFoundNaNElement = false;
+               |} else {
+               |  $body
+               |}
+             """.stripMargin
+          }
+        }.getOrElse(body)
+
         val writeArray2ToHashSet = withArray2NullCheck(
-          s"""
-             |$jt $value = ${genGetValue(array2, i)};
-             |$hashSet.add$hsPostFix($hsValueCast$value);
-           """.stripMargin)
+          s"$jt $value = ${genGetValue(array2, i)};" +
+            withArray2NaNCheck(s"$hashSet.add$hsPostFix($hsValueCast$value);"))
 
         def withArray1NullAssignment(body: String) =
           if (left.dataType.asInstanceOf[ArrayType].containsNull) {
@@ -4221,17 +4251,35 @@ case class ArrayExcept(left: Expression, right: Expression) extends ArrayBinaryL
             body
           }
 
-        val processArray1 = withArray1NullAssignment(
+        def withArray1NaNCheck(body: String): String = {
+          isNaNMethod.map { isNaN =>
+            s"""
+               |if ($isNaN) {
+               |  if ($notFoundNaNElement) {
+               |    $notFoundNaNElement = false;
+               |    $size++;
+               |    $builder.$$plus$$eq($value);
+               |  }
+               |} else {
+               |  $body
+               |}
+             """.stripMargin
+          }
+        }.getOrElse(body)
+
+        val body =
           s"""
-             |$jt $value = ${genGetValue(array1, i)};
              |if (!$hashSet.contains($hsValueCast$value)) {
              |  if (++$size > ${ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH}) {
              |    break;
              |  }
              |  $hashSet.add$hsPostFix($hsValueCast$value);
              |  $builder.$$plus$$eq($value);
              |}
-           """.stripMargin)
+           """.stripMargin
+
+        val processArray1 = withArray1NullAssignment(
+          s"$jt $value = ${genGetValue(array1, i)};" + withArray1NaNCheck(body))
 
         // Only need to track null element index when array1's element is nullable.
         val declareNullTrackVariables = if (left.dataType.asInstanceOf[ArrayType].containsNull) {
@@ -4243,9 +4291,16 @@ case class ArrayExcept(left: Expression, right: Expression) extends ArrayBinaryL
           ""
         }
 
+        // Only need to track NaN element index when array1's element is DoubleType or FloatType.
+        val declareNaNTrackVariables = elementType match {
+          case DoubleType | FloatType => s"boolean $notFoundNaNElement = true;"
+          case _ => ""
+        }
+
         s"""
            |$openHashSet $hashSet = new $openHashSet$hsPostFix($classTag);
            |$declareNullTrackVariables
+           |$declareNaNTrackVariables
            |for (int $i = 0; $i < $array2.numElements(); $i++) {
            |  $writeArray2ToHashSet
            |}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala

@@ -2326,4 +2326,21 @@ class CollectionExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper
       Literal.create(Seq(Float.NaN, null, 1f), ArrayType(FloatType))),
       Seq(Float.NaN, null, 1f))
   }
+
+  test("SPARK-36753: ArrayExcept should handle duplicated Double.NaN and Float.Nan") {
+    checkEvaluation(ArrayExcept(
+      Literal.apply(Array(Double.NaN, 1d)), Literal.apply(Array(Double.NaN))),
+      Seq(1d))
+    checkEvaluation(ArrayExcept(
+      Literal.create(Seq(null, Double.NaN, null, 1d), ArrayType(DoubleType)),
+      Literal.create(Seq(Double.NaN, null), ArrayType(DoubleType))),
+      Seq(1d))
+    checkEvaluation(ArrayExcept(
+      Literal.apply(Array(Float.NaN, 1f)), Literal.apply(Array(Float.NaN))),
+      Seq(1f))
+    checkEvaluation(ArrayExcept(
+      Literal.create(Seq(null, Float.NaN, null, 1f), ArrayType(FloatType)),
+      Literal.create(Seq(Float.NaN, null), ArrayType(FloatType))),
+      Seq(1f))
+  }
 }