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

github-merge-queue[bot] pushed a commit to branch 
gh-readonly-queue/main/pr-5472-e6e92f1a1a10e1bca2bcf32858113ab61b2ba685
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git

commit de2e70291460ed0310428493d8c3aa5223293d08
Author: Chao Sun <[email protected]>
AuthorDate: Wed Sep 23 21:36:30 2026 +0000

    fix: normalize noncanonical NaN literals in comparisons (#5472)
    
    * fix: normalize noncanonical NaN literals in comparisons
    
    * fix: normalize floating IN operands and expand native filter tests
    
    * fix: retain fallback reasons for normalized IN operands
    
    * fix: preserve pruning for finite floating IN lists
    
    * Preserve Parquet pruning for infinity IN literals
    
    * fix: preserve pruning for literal zero IN lists
---
 .../latest/compatibility/floating-point.md         |  12 +-
 .../org/apache/comet/rules/CometExecRule.scala     |   4 +-
 .../org/apache/comet/serde/QueryPlanSerde.scala    |   2 +-
 .../scala/org/apache/comet/serde/predicates.scala  | 123 ++++++++++++-
 .../org/apache/comet/CometExpressionSuite.scala    | 151 +++++++++++++++-
 .../apache/comet/rules/CometExecRuleSuite.scala    | 192 ++++++++++++++++++++-
 6 files changed, 469 insertions(+), 15 deletions(-)

diff --git a/docs/source/user-guide/latest/compatibility/floating-point.md 
b/docs/source/user-guide/latest/compatibility/floating-point.md
index 1d3c4b2957..7d8817db6f 100644
--- a/docs/source/user-guide/latest/compatibility/floating-point.md
+++ b/docs/source/user-guide/latest/compatibility/floating-point.md
@@ -23,10 +23,14 @@ Spark normalizes NaN and zero for floating point numbers 
for several cases. See
 However, one exception is comparison. Spark does not normalize NaN and zero 
when comparing values
 because they are handled well in Spark (e.g., 
`SQLOrderingUtil.compareFloats`). But the comparison
 functions of arrow-rs used by DataFusion do not normalize NaN and zero (e.g., 
[arrow::compute::kernels::cmp::eq](https://docs.rs/arrow/latest/arrow/compute/kernels/cmp/fn.eq.html#)).
-So Comet adds additional normalization expression of NaN and zero for 
comparisons, and may still have differences
-to Spark in some cases, especially when the data contains both positive and 
negative zero. This is likely an edge
-case that is not of concern for many users. If it is a concern, setting 
`spark.comet.exec.strictFloatingPoint=true`
-will make relevant operations fall back to Spark.
+For top-level `FLOAT` and `DOUBLE` comparisons, Comet normalizes both operands 
before native
+execution, including noncanonical NaN literals. Top-level `IN`, `InSet`, and 
`NOT IN` membership
+also normalize dynamic candidates and lists containing NaN. When every 
candidate is a non-NaN
+literal, Comet keeps DataFusion's static filter and pruning path, enumerating 
both signed-zero
+forms when a list contains zero.
+
+This scalar membership handling does not yet recurse into floating-point 
leaves nested in arrays
+or structs; see 
[#6019](https://github.com/apache/datafusion-comet/issues/6019).
 
 ## Nested equality and membership
 
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala 
b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
index 54f8891624..286f5adea0 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -712,8 +712,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 =>
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala 
b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index aa4399c3b3..3944b13f49 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -924,7 +924,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 ++= _)
diff --git a/spark/src/main/scala/org/apache/comet/serde/predicates.scala 
b/spark/src/main/scala/org/apache/comet/serde/predicates.scala
index 0e7bb02cd6..129979cff7 100644
--- a/spark/src/main/scala/org/apache/comet/serde/predicates.scala
+++ b/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,6 +366,102 @@ object ComparisonUtils {
   val inUnsupportedReasons: Seq[String] =
     Seq(nonDefaultCollationDocReason, legacyNullInEmptyListReason)
 
+  /**
+   * Normalize one top-level floating-point membership operand to Spark's NaN 
and signed-zero
+   * equality. Literal normalization is folded immediately so it remains a 
scalar literal; a
+   * dynamic expression gets Spark's normalization wrapper; non-floating 
operands are unchanged.
+   *
+   * @param expr
+   *   The value or candidate operand being serialized.
+   * @return
+   *   The operand to serialize, possibly folded or wrapped for 
Spark-compatible equality.
+   */
+  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
+  }
+
+  /**
+   * Keep an all-literal floating-point membership list on DataFusion's 
static-filter path while
+   * matching Spark's signed-zero equality. DataFusion hashes the raw 
floating-point bits for a
+   * static list, so a list containing one zero sign also needs the other 
sign. A NaN literal or a
+   * dynamic candidate cannot be made equivalent by enumeration and must use 
normalization
+   * instead.
+   *
+   * @param list
+   *   The original Spark `IN` / `InSet` candidates.
+   * @return
+   *   `Some` with a static candidate list, including any missing 
opposite-signed zeros, or `None`
+   *   when the caller must normalize the membership operands.
+   */
+  private def staticFloatingInList(list: Seq[Expression]): 
Option[Seq[Expression]] = {
+    val canStayStatic = list.forall {
+      case Literal(null, _) => true
+      case Literal(v: Float, FloatType) => !java.lang.Float.isNaN(v)
+      case Literal(v: Double, DoubleType) => !java.lang.Double.isNaN(v)
+      case _ => false
+    }
+    if (!canStayStatic) {
+      None
+    } else {
+      val hasPositiveFloatZero = list.exists {
+        case Literal(v: Float, FloatType) =>
+          java.lang.Float.floatToRawIntBits(v) == 
java.lang.Float.floatToRawIntBits(0.0f)
+        case _ => false
+      }
+      val hasNegativeFloatZero = list.exists {
+        case Literal(v: Float, FloatType) =>
+          java.lang.Float.floatToRawIntBits(v) == 
java.lang.Float.floatToRawIntBits(-0.0f)
+        case _ => false
+      }
+      val hasPositiveDoubleZero = list.exists {
+        case Literal(v: Double, DoubleType) =>
+          java.lang.Double.doubleToRawLongBits(v) == 
java.lang.Double.doubleToRawLongBits(0.0d)
+        case _ => false
+      }
+      val hasNegativeDoubleZero = list.exists {
+        case Literal(v: Double, DoubleType) =>
+          java.lang.Double.doubleToRawLongBits(v) == 
java.lang.Double.doubleToRawLongBits(-0.0d)
+        case _ => false
+      }
+      val missingZeroSigns = Seq(
+        if (hasPositiveFloatZero && !hasNegativeFloatZero) 
Some(Literal(-0.0f)) else None,
+        if (hasNegativeFloatZero && !hasPositiveFloatZero) Some(Literal(0.0f)) 
else None,
+        if (hasPositiveDoubleZero && !hasNegativeDoubleZero) 
Some(Literal(-0.0d)) else None,
+        if (hasNegativeDoubleZero && !hasPositiveDoubleZero) 
Some(Literal(0.0d))
+        else None).flatten
+      Some(list ++ missingZeroSigns)
+    }
+  }
+
+  /**
+   * Serialize Spark membership while preserving Spark-compatible top-level 
floating-point
+   * equality and DataFusion's static-filter pruning whenever that is safe.
+   *
+   * @param expr
+   *   The original membership expression, used as the fallback-reason owner.
+   * @param value
+   *   The value being tested for membership.
+   * @param list
+   *   The membership candidates.
+   * @param inputs
+   *   The attributes available for bound-expression serialization.
+   * @param binding
+   *   Whether attributes should be bound to input ordinals.
+   * @param negate
+   *   Whether to serialize `NOT IN` rather than `IN`.
+   * @return
+   *   The native membership protobuf, or `None` after copying any 
synthesized-expression fallback
+   *   reasons back to `expr`.
+   */
   def in(
       expr: Expression,
       value: Expression,
@@ -372,8 +469,21 @@ object ComparisonUtils {
       inputs: Seq[Attribute],
       binding: Boolean,
       negate: Boolean): Option[Expr] = {
-    val valueExpr = exprToProtoInternal(value, inputs, binding)
-    val listExprs = list.map(exprToProtoInternal(_, inputs, binding))
+    val serializedOperands: (Expression, Seq[Expression]) = value.dataType 
match {
+      case FloatType | DoubleType =>
+        staticFloatingInList(list) match {
+          // All-literal non-NaN lists stay static so native Parquet scans can 
still prune using
+          // column statistics. Enumerating both zero signs makes raw-bit 
membership match Spark.
+          case Some(staticList) => value -> staticList
+          // NaN literals and dynamic candidates need both sides normalized, 
including fused NOT IN.
+          case None => normalizeInOperand(value) -> 
list.map(normalizeInOperand)
+        }
+      // Nested floating-point leaves are intentionally outside this scalar 
normalization path.
+      case _ => value -> list
+    }
+    val (serializedValue, serializedList) = serializedOperands
+    val valueExpr = exprToProtoInternal(serializedValue, inputs, binding)
+    val listExprs = serializedList.map(exprToProtoInternal(_, inputs, binding))
     if (valueExpr.isDefined && listExprs.forall(_.isDefined)) {
       val builder = ExprOuterClass.In.newBuilder()
       builder.setInValue(valueExpr.get)
@@ -385,6 +495,11 @@ object ComparisonUtils {
           .setIn(builder)
           .build())
     } else {
+      // Normalization and static-list expansion create temporary wrappers and 
literals outside the
+      // original tree. Keep their failure reasons on the membership 
expression so the operator can
+      // explain fallback.
+      liftFallbackReasons(serializedValue, expr)
+      serializedList.foreach(liftFallbackReasons(_, expr))
       None
     }
   }
diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala 
b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
index 56692797f2..f26ef595a6 100644
--- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
+++ b/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, 
In, InSet, Literal, StructsToJson, TruncDate, TruncTimestamp}
 import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, 
OptimizeIn, SimplifyExtractValueOps}
-import org.apache.spark.sql.comet.{CometProjectExec, CometSortExec, 
CometTakeOrderedAndProjectExec}
+import org.apache.spark.sql.comet.{CometFilterExec, CometProjectExec, 
CometSortExec, CometTakeOrderedAndProjectExec}
 import org.apache.spark.sql.execution.{LocalTableScanExec, ProjectExec, 
SparkPlan}
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.functions._
@@ -328,6 +328,155 @@ 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("floating IN with a column candidate normalizes signed zeros") {
+    withSQLConf(
+      SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "false",
+      "spark.sql.optimizer.inSetConversionThreshold" -> "10") {
+      val rows = Seq(
+        (0, Some(-0.0d), Some(0.0d)),
+        (1, Some(0.0d), Some(-0.0d)),
+        (2, Some(1.0d), Some(2.0d)),
+        (3, None, Some(0.0d)))
+      withParquetDataFrame(rows, withDictionary = false) { df =>
+        // The first candidate is another column, so DataFusion evaluates 
dynamic equality rather
+        // than a static literal filter. The second candidate keeps Spark from 
rewriting a
+        // singleton IN to ordinary equality before Comet serializes it.
+        val predicate = df("_2").isin(df("_3"), lit(13.0d))
+        val projected = df.select(df("_1"), predicate)
+        val optimized = projected.queryExecution.optimizedPlan
+        val membership = optimized.expressions.flatMap(_.collect {
+          case in: org.apache.spark.sql.catalyst.expressions.In => in
+        })
+        assert(
+          membership.size == 1 &&
+            membership.forall(_.list.exists(candidate => 
!candidate.isInstanceOf[Literal])),
+          optimized.toString)
+
+        checkSparkAnswerAndOperator(projected, Seq(classOf[CometProjectExec]))
+        checkSparkAnswerAndOperator(
+          df.filter(predicate).select("_1"),
+          Seq(classOf[CometFilterExec]))
+      }
+    }
+  }
+
   test("parquet default values") {
     withTable("t1") {
       sql("create table t1(col1 boolean) using parquet")
diff --git 
a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala 
b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
index 8bf6097754..d59a9d3207 100644
--- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
@@ -24,8 +24,9 @@ import scala.util.Random
 import org.apache.logging.log4j.Level
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.FunctionIdentifier
-import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, 
ExpressionInfo, Literal}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, 
AttributeReference, Expression, ExpressionInfo, In, InSet, 
KnownFloatingPointNormalized, Literal, Not}
 import 
org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, 
BloomFilterAggregate, Final, Min, Partial, PartialMerge}
+import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero
 import org.apache.spark.sql.comet._
 import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
 import org.apache.spark.sql.execution._
@@ -33,11 +34,11 @@ import 
org.apache.spark.sql.execution.adaptive.{QueryStageExec, ShuffleQueryStag
 import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, 
ObjectHashAggregateExec}
 import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, 
ShuffleExchangeExec}
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.{DataTypes, StructField, StructType}
+import org.apache.spark.sql.types.{DataTypes, DoubleType, FloatType, 
StructField, StructType}
 
 import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo}
 import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, 
isSpark40Plus, isSpark42Plus, withFallbackReason}
-import org.apache.comet.serde.{CometAggregateExpressionSerde, Compatible, 
ExprOuterClass, Unsupported}
+import org.apache.comet.serde.{CometAggregateExpressionSerde, Compatible, 
ExprOuterClass, QueryPlanSerde, Unsupported}
 import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator}
 
 /**
@@ -137,6 +138,191 @@ class CometExecRuleSuite extends CometTestBase {
     }
   }
 
+  for (dataType <- Seq(FloatType, DoubleType)) {
+    test(s"floating ${dataType.sql} IN serialization preserves prunable 
literal lists") {
+      withSQLConf("spark.sql.legacy.nullInEmptyListBehavior" -> "false") {
+        val value = AttributeReference("value", dataType)()
+        val other = AttributeReference("other", dataType)()
+        def literal(v: Double): Literal = dataType match {
+          case FloatType => Literal(v.toFloat)
+          case DoubleType => Literal(v)
+        }
+        val ordinary = Seq(1.0d, 3.0d).map(literal)
+        val infinities = Seq(Double.PositiveInfinity, 
Double.NegativeInfinity).map(literal)
+        val nullLiteral = Literal.create(null, dataType)
+        val lists: Seq[(Seq[Expression], Boolean, Int)] = Seq(
+          (ordinary, false, 0),
+          (ordinary :+ nullLiteral, false, 0),
+          (infinities, false, 0),
+          (infinities :+ nullLiteral, false, 0),
+          (Seq(nullLiteral), false, 0),
+          (Seq(value, other), true, 0)) ++
+          Seq(Double.PositiveInfinity, Double.NegativeInfinity)
+            .map(v => (ordinary :+ literal(v), false, 0)) ++
+          Seq(Double.NaN)
+            .flatMap(v => Seq(ordinary, infinities).map(list => (list :+ 
literal(v), true, 0))) ++
+          Seq(0.0d, -0.0d)
+            .flatMap(v =>
+              Seq(ordinary, infinities).map(list => (list :+ literal(v), 
false, 1))) ++
+          Seq((ordinary ++ Seq(literal(0.0d), literal(-0.0d)), false, 0)) ++
+          (if (isSpark35Plus) Seq((Seq.empty[Expression], false, 0)) else Nil)
+        for ((list, needsNormalization, extraStaticCandidates) <- lists;
+          asSet <- Seq(false, true) if !asSet || 
list.forall(_.isInstanceOf[Literal]);
+          negate <- Seq(false, true);
+          alreadyNormalized <- Seq(false, true)) {
+          withClue(s"list=$list, asSet=$asSet, negate=$negate, 
normalized=$alreadyNormalized: ") {
+            val needle = if (alreadyNormalized) {
+              KnownFloatingPointNormalized(NormalizeNaNAndZero(value))
+            } else {
+              value
+            }
+            val in = if (asSet) {
+              InSet(needle, list.collect { case l: Literal => l.value }.toSet)
+            } else {
+              In(needle, list)
+            }
+            val result = QueryPlanSerde
+              .exprToProto(if (negate) Not(in) else in, Seq(value, other))
+              .get
+            // NOT InSet uses a separate Not node; NOT In is fused into the 
membership node.
+            val serialized = if (result.hasNot) result.getNot.getChild else 
result
+            assert(serialized.hasIn)
+            assert(result.hasNot == (asSet && negate))
+            assert(serialized.getIn.getNegated == (negate && !asSet))
+            val serializedValue = serialized.getIn.getInValue
+            if (needsNormalization || alreadyNormalized) {
+              assert(serializedValue.hasNormalizeNanAndZero)
+              assert(serializedValue.getNormalizeNanAndZero.getChild.hasBound)
+            } else {
+              assert(serializedValue.hasBound)
+            }
+            assert(serialized.getIn.getListsCount == list.size + 
extraStaticCandidates)
+            for (i <- list.indices) {
+              val candidate = serialized.getIn.getLists(i)
+              if (list(i).isInstanceOf[Literal]) {
+                assert(candidate.hasLiteral)
+              } else {
+                assert(candidate.hasNormalizeNanAndZero)
+                assert(candidate.getNormalizeNanAndZero.getChild.hasBound)
+              }
+            }
+            for (i <- list.size until list.size + extraStaticCandidates) {
+              assert(serialized.getIn.getLists(i).hasLiteral)
+            }
+          }
+        }
+      }
+    }
+
+    test(
+      s"floating ${dataType.sql} IN serialization retains normalized operand 
fallback reasons") {
+      val expressionNames = Seq("Literal", "KnownFloatingPointNormalized")
+      for (disabled <- None +: expressionNames.map(Some(_));
+        literalValue <- Seq(false, true);
+        negate <- Seq(false, true)) {
+        val configs = expressionNames.map { name =>
+          CometConf.getExprEnabledConfigKey(name) -> 
(!disabled.contains(name)).toString
+        }
+        withSQLConf(configs: _*) {
+          withClue(s"disabled=$disabled, literalValue=$literalValue, 
negate=$negate: ") {
+            val value = AttributeReference("value", dataType)()
+            val other = AttributeReference("other", dataType)()
+            val literals = dataType match {
+              case FloatType => Seq(Literal(Float.NaN), Literal(3.0f))
+              case DoubleType => Seq(Literal(Double.NaN), Literal(3.0d))
+            }
+            // Exercise temporary literals and normalizers in both the value 
and the list.
+            val in =
+              if (literalValue) In(literals.head, Seq(value, other)) else 
In(value, literals)
+            val expr = if (negate) Not(in) else in
+            val result = QueryPlanSerde.exprToProto(expr, Seq(value, other))
+            val reasons = in
+              .getTagValue(CometExplainInfo.FALLBACK_REASONS)
+              .getOrElse(Set.empty[String])
+            disabled match {
+              case Some(name) =>
+                assert(result.isEmpty)
+                val key = CometConf.getExprEnabledConfigKey(name)
+                assert(
+                  reasons == Set(s"Expression support is disabled. Set 
$key=true to enable it."))
+              case None =>
+                assert(result.exists(_.hasIn))
+                assert(result.get.getIn.getNegated == negate)
+                assert(reasons.isEmpty)
+            }
+          }
+        }
+      }
+    }
+
+    test(s"floating ${dataType.sql} IN planning retains normalized operand 
fallback reasons") {
+      val expressionNames = Seq("Literal", "KnownFloatingPointNormalized")
+      for (disabled <- None +: expressionNames.map(Some(_));
+        strict <- Seq(false, true);
+        literalValue <- Seq(false, true);
+        negate <- Seq(false, true)) {
+        val configs = Seq(
+          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+          SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
+          "spark.sql.optimizer.inSetConversionThreshold" -> "100",
+          CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true",
+          CometConf.COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST.key -> 
"Range",
+          CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false",
+          CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false",
+          CometConf.COMET_STRICT_FALLBACK_REASONS.key -> strict.toString) ++
+          expressionNames.map { name =>
+            CometConf.getExprEnabledConfigKey(name) -> 
(!disabled.contains(name)).toString
+          }
+        withSQLConf(configs: _*) {
+          withClue(
+            s"disabled=$disabled, strict=$strict, literalValue=$literalValue, 
negate=$negate: ") {
+            val column = s"CAST(id AS ${dataType.sql})"
+            val predicate = if (literalValue) {
+              s"CAST(1 AS ${dataType.sql}) IN ($column, -$column)"
+            } else {
+              // Keep a dynamic candidate here: all-literal non-NaN zero lists 
now stay static and
+              // intentionally avoid the normalization wrappers this fallback 
test exercises.
+              s"$column IN (CAST(0 AS ${dataType.sql}), -$column)"
+            }
+            val expression = if (negate) s"NOT ($predicate)" else predicate
+            val df = sql(s"SELECT $expression AS hit FROM range(0, 4, 1, 1)")
+            val optimized = df.queryExecution.optimizedPlan
+            val expressions = optimized.flatMap(_.expressions)
+            val membership = expressions.flatMap(_.collect { case in: In => in 
})
+            // A folded predicate, singleton equality, or InSet would miss 
this serializer.
+            assert(membership.size == 1 && membership.head.list.size == 2, 
optimized.toString)
+            assert(membership.head.value.isInstanceOf[Literal] == literalValue)
+            assert(expressions.exists(_.exists {
+              case Not(_: In) => true
+              case _ => false
+            }) == negate)
+
+            // Planning itself used to throw in strict mode, before a native 
task could run.
+            val plan = df.queryExecution.executedPlan
+            val projects = plan.collect { case p: ProjectExec => p }
+            disabled match {
+              case Some(name) =>
+                assert(projects.size == 1, plan.toString)
+                assert(plan.find(_.isInstanceOf[CometProjectExec]).isEmpty, 
plan.toString)
+                val key = CometConf.getExprEnabledConfigKey(name)
+                val reasons = projects.head
+                  .getTagValue(CometExplainInfo.FALLBACK_REASONS)
+                  .getOrElse(Set.empty[String])
+                assert(
+                  reasons == Set(s"Expression support is disabled. Set 
$key=true to enable it."))
+              case None =>
+                assert(projects.isEmpty, plan.toString)
+                assert(plan.find(_.isInstanceOf[CometProjectExec]).isDefined, 
plan.toString)
+                assert(
+                  !plan.exists(
+                    
_.getTagValue(CometExplainInfo.FALLBACK_REASONS).exists(_.nonEmpty)))
+            }
+          }
+        }
+      }
+    }
+  }
+
   test("strict mode fails an operator that Comet declined without recording a 
reason") {
     // The bug this guards against is a serde returning None and forgetting to 
say why, which the
     // generic "<operator> is not supported" message used to hide. No serde in 
the tree is in that


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to