fix: normalize noncanonical NaN literals in comparisons

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -562,8 +562,8 @@ case class CometExecRule(session: SparkSession)
   private def normalizeNaNAndZero(expr: Expression): Expression = {
     expr match {
       case _: KnownFloatingPointNormalized => expr
-      case FloatLiteral(f) if !f.equals(-0.0f) => expr
-      case DoubleLiteral(d) if !d.equals(-0.0d) => expr
+      case FloatLiteral(f) if !f.isNaN && !f.equals(-0.0f) => expr
+      case DoubleLiteral(d) if !d.isNaN && !d.equals(-0.0d) => expr
       case _ =>
         expr.dataType match {
           case _: FloatType | _: DoubleType =>

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -876,7 +876,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
    * converted, so lifting one off a tree that converted fine would attribute a stale reason to an
    * operator that has no problem.
    */
-  private def liftFallbackReasons(from: Expression, to: Expression): Unit = {
+  private[serde] def liftFallbackReasons(from: Expression, to: Expression): Unit = {
     val reasons = mutable.Set.empty[String]
     from.foreach { e =>
       e.getTagValue(CometExplainInfo.FALLBACK_REASONS).foreach(reasons ++= _)

spark/src/main/scala/org/apache/comet/serde/predicates.scala

@@ -21,9 +21,10 @@ package org.apache.comet.serde
 
 import scala.jdk.CollectionConverters._
 
-import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BinaryExpression, EqualNullSafe, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, In, InSet, IsNaN, IsNotNull, IsNull, LessThan, LessThanOrEqual, Literal, Not, Or}
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, BinaryExpression, EqualNullSafe, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, In, InSet, IsNaN, IsNotNull, IsNull, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, Literal, Not, Or}
+import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.BooleanType
+import org.apache.spark.sql.types.{BooleanType, DoubleType, FloatType}
 
 import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus}
@@ -365,15 +366,39 @@ object ComparisonUtils {
   val inUnsupportedReasons: Seq[String] =
     Seq(nonDefaultCollationDocReason, legacyNullInEmptyListReason)
 
+  private def normalizeInOperand(expr: Expression): Expression = expr.dataType match {
+    case FloatType | DoubleType =>
+      expr match {
+        case _: KnownFloatingPointNormalized => expr
+        // DataFusion's static IN filter hashes raw floating-point bits. Fold literal
+        // normalization here so the list remains scalar and can still use that filter.
+        case literal: Literal =>
+          Literal(NormalizeNaNAndZero(literal).eval(), literal.dataType)
+        case _ => KnownFloatingPointNormalized(NormalizeNaNAndZero(expr))
+      }
+    case _ => expr
+  }
+
   def in(
       expr: Expression,
       value: Expression,
       list: Seq[Expression],
       inputs: Seq[Attribute],
       binding: Boolean,
       negate: Boolean): Option[Expr] = {
-    val valueExpr = exprToProtoInternal(value, inputs, binding)
-    val listExprs = list.map(exprToProtoInternal(_, inputs, binding))
+    // NaNs and either sign of zero cannot match a non-NaN nonzero literal. Leave such lists
+    // unwrapped so native Parquet scans can still prune using the column's statistics.
+    val needsNormalization = !list.forall {
+      case Literal(null, _) => true
+      case Literal(v: Float, FloatType) => !java.lang.Float.isNaN(v) && v != 0.0f
+      case Literal(v: Double, DoubleType) => !java.lang.Double.isNaN(v) && v != 0.0d
+      case _ => false
+    }
+    // Otherwise normalize both sides for Spark's NaN/zero equality, including fused NOT IN.
+    val normalizedValue = if (needsNormalization) normalizeInOperand(value) else value
+    val normalizedList = if (needsNormalization) list.map(normalizeInOperand) else list
+    val valueExpr = exprToProtoInternal(normalizedValue, inputs, binding)
+    val listExprs = normalizedList.map(exprToProtoInternal(_, inputs, binding))
     if (valueExpr.isDefined && listExprs.forall(_.isDefined)) {
       val builder = ExprOuterClass.In.newBuilder()
       builder.setInValue(valueExpr.get)
@@ -385,6 +410,10 @@ object ComparisonUtils {
           .setIn(builder)
           .build())
     } else {
+      // Normalization creates temporary wrappers and literals outside the original tree. Keep
+      // their failure reasons on the membership expression so the operator can explain fallback.
+      liftFallbackReasons(normalizedValue, expr)
+      normalizedList.foreach(liftFallbackReasons(_, expr))
       None
     }
   }

spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala

@@ -25,7 +25,7 @@ import org.apache.hadoop.fs.Path
 import org.apache.spark.sql.{Column, CometTestBase, DataFrame, Row}
 import org.apache.spark.sql.catalyst.expressions.{Alias, Cast, FromUnixTime, Literal, StructsToJson, TruncDate, TruncTimestamp}
 import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, OptimizeIn, SimplifyExtractValueOps}
-import org.apache.spark.sql.comet.CometProjectExec
+import org.apache.spark.sql.comet.{CometFilterExec, CometProjectExec}
 import org.apache.spark.sql.execution.{ProjectExec, SparkPlan}
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.functions._
@@ -172,6 +172,123 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
+  Seq(
+    (
+      "float",
+      "_1",
+      Seq[Any](
+        java.lang.Float.intBitsToFloat(0x7fc00001),
+        java.lang.Float.intBitsToFloat(0xffc00002))),
+    (
+      "double",
+      "_2",
+      Seq[Any](
+        java.lang.Double.longBitsToDouble(0x7ff8000000000001L),
+        java.lang.Double.longBitsToDouble(0xfff8000000000002L)))).foreach {
+    case (dataType, column, nanLiterals) =>
+      test(s"compare $dataType columns with noncanonical NaN literals") {
+        withSQLConf(SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "false") {
+          val rows = Seq(
+            (Some(Float.NaN), Some(Double.NaN)),
+            (Some(-0.0f), Some(-0.0d)),
+            (Some(0.0f), Some(0.0d)),
+            (Some(-1.0f), Some(-1.0d)),
+            (Some(1.0f), Some(1.0d)),
+            (Some(Float.NegativeInfinity), Some(Double.NegativeInfinity)),
+            (Some(Float.PositiveInfinity), Some(Double.PositiveInfinity)),
+            (None, None))
+          val identifiedRows = rows.zipWithIndex.map { case ((f, d), id) => (f, d, id) }
+          withParquetDataFrame(identifiedRows, withDictionary = false) { df =>
+            // Parquet canonicalizes stored NaNs, so construct the signed/payload literals here.
+            // Compare Boolean results so Spark's NaN-aware answer checker cannot hide a mismatch.
+            val value = df(column)
+            nanLiterals.foreach { nan =>
+              val literal = lit(nan)
+              Seq((value, literal), (literal, value)).foreach { case (left, right) =>
+                val comparisons = Seq(
+                  left === right,
+                  left =!= right,
+                  left.eqNullSafe(right),
+                  left < right,
+                  left <= right,
+                  left > right,
+                  left >= right)
+                checkSparkAnswerAndOperator(
+                  df.select(comparisons: _*),
+                  Seq(classOf[CometProjectExec]))
+                comparisons.foreach { comparison =>
+                  // Compare surviving identities, not just NaN-aware row values. Keep Parquet
+                  // pushdown disabled so every ordering predicate executes in CometFilterExec.
+                  checkSparkAnswerAndOperator(
+                    df.filter(comparison).select("_3"),
+                    Seq(classOf[CometFilterExec]))
+                }
+              }
+            }
+          }
+        }
+      }
+  }
+
+  for ((name, threshold) <- Seq(("In", 10), ("InSet", 1))) {
+    test(s"floating $name and NOT $name normalize NaNs and signed zeros") {
+      withSQLConf(
+        SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "false",
+        "spark.sql.optimizer.inSetConversionThreshold" -> threshold.toString) {
+        val rows = Seq(
+          (0, Some(Float.NaN), Some(Double.NaN)),
+          (1, Some(0.0f), Some(0.0d)),
+          (2, Some(-0.0f), Some(-0.0d)),
+          (3, Some(13.0f), Some(13.0d)),
+          (4, Some(1.0f), Some(1.0d)),
+          (5, None, None),
+          (6, Some(Float.PositiveInfinity), Some(Double.PositiveInfinity)),
+          (7, Some(Float.NegativeInfinity), Some(Double.NegativeInfinity)))
+        withParquetDataFrame(rows, withDictionary = false) { df =>
+          val cases = Seq(
+            (
+              java.lang.Float.intBitsToFloat(0x7fc00001),
+              java.lang.Double.longBitsToDouble(0x7ff8000000000001L)),
+            (
+              java.lang.Float.intBitsToFloat(0xffc00002),
+              java.lang.Double.longBitsToDouble(0xfff8000000000002L)),
+            (0.0f, 0.0d),
+            (-0.0f, -0.0d),
+            (Float.PositiveInfinity, Double.PositiveInfinity),
+            (Float.NegativeInfinity, Double.NegativeInfinity))
+          for ((f, d) <- cases; includeNull <- Seq(false, true)) {
+            val floatCandidates: Seq[Any] = Seq(f, 13.0f) ++ (if (includeNull) Seq(null) else Nil)
+            val doubleCandidates: Seq[Any] =
+              Seq(d, 13.0d) ++ (if (includeNull) Seq(null) else Nil)
+            // Negation creates negative NaNs after the Parquet scan, so the membership value
+            // must be normalized as well as the programmatically constructed list literals.
+            val predicates = Seq(df("_2"), -df("_2")).map(_.isin(floatCandidates: _*)) ++
+              Seq(df("_3"), -df("_3")).map(_.isin(doubleCandidates: _*))
+            val projected = df.select(df("_1") +: predicates.flatMap(p => Seq(p, !p)): _*)
+            val optimized = projected.queryExecution.optimizedPlan
+            val membership = optimized.expressions.flatMap(_.collect {
+              case _: org.apache.spark.sql.catalyst.expressions.In => "In"
+              case _: org.apache.spark.sql.catalyst.expressions.InSet => "InSet"
+            })
+            // A singleton IN is rewritten to equality and would not exercise the faulty kernel.
+            assert(membership.nonEmpty && membership.forall(_ == name), optimized.toString)
+            checkSparkAnswerAndOperator(projected, Seq(classOf[CometProjectExec]))
+            for (p <- predicates; negate <- Seq(false, true)) {
+              val filtered = df.filter(if (negate) !p else p).select("_1")
+              if (includeNull && negate) {
+                // NOT IN with a null candidate can never be true. Spark legitimately replaces
+                // this filter with an empty LocalRelation before physical planning.
+                checkSparkAnswer(filtered)
+              } else {
+                checkSparkAnswerAndOperator(filtered, Seq(classOf[CometFilterExec]))
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
   test("parquet default values") {
     withTable("t1") {
       sql("create table t1(col1 boolean) using parquet")