From 9b6f7c05badc3dcfc2445d34314184ac9d9df7dc Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 26 Aug 2026 05:16:32 +0000 Subject: [PATCH 1/5] fix: normalize noncanonical NaN literals in comparisons --- .../apache/comet/rules/CometExecRule.scala | 4 +- .../apache/comet/CometExpressionSuite.scala | 55 ++++++++++++++++++- 2 files changed, 56 insertions(+), 3 deletions(-) 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 c69602fc80..50907aa8eb 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/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 => diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index 80f92c4b7f..2adcf02791 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, 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,59 @@ 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)) + withParquetDataFrame(rows, 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])) + } + checkSparkAnswerAndOperator( + df.filter(value === literal), + Seq(classOf[CometFilterExec])) + } + } + } + } + } + test("parquet default values") { withTable("t1") { sql("create table t1(col1 boolean) using parquet") From 4a151c4c706b7753f15dc6715db0b1c7c826b338 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Thu, 27 Aug 2026 17:13:06 +0000 Subject: [PATCH 2/5] fix: normalize floating IN operands and expand native filter tests --- .../org/apache/comet/serde/predicates.scala | 24 +++++-- .../apache/comet/CometExpressionSuite.scala | 68 +++++++++++++++++-- 2 files changed, 84 insertions(+), 8 deletions(-) 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..dfa2094898 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,19 @@ 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, @@ -372,8 +386,10 @@ object ComparisonUtils { inputs: Seq[Attribute], binding: Boolean, negate: Boolean): Option[Expr] = { - val valueExpr = exprToProtoInternal(value, inputs, binding) - val listExprs = list.map(exprToProtoInternal(_, inputs, binding)) + // Spark treats all NaNs as equal and both signs of zero as equal in IN and InSet too. + // Normalize both sides, including the fused NOT IN path that calls this method directly. + val valueExpr = exprToProtoInternal(normalizeInOperand(value), inputs, binding) + val listExprs = list.map(e => exprToProtoInternal(normalizeInOperand(e), inputs, binding)) if (valueExpr.isDefined && listExprs.forall(_.isDefined)) { val builder = ExprOuterClass.In.newBuilder() builder.setInValue(valueExpr.get) diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index 2adcf02791..5d655134c3 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -197,7 +197,8 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { (Some(Float.NegativeInfinity), Some(Double.NegativeInfinity)), (Some(Float.PositiveInfinity), Some(Double.PositiveInfinity)), (None, None)) - withParquetDataFrame(rows, withDictionary = false) { df => + 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) @@ -215,14 +216,73 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { 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)) + 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)) + 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])) } - checkSparkAnswerAndOperator( - df.filter(value === literal), - Seq(classOf[CometFilterExec])) } } } } + } } test("parquet default values") { From 9dd5c044cf6e427c04a71ad18d6a8554eb14fced Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Fri, 28 Aug 2026 01:12:08 +0000 Subject: [PATCH 3/5] fix: retain fallback reasons for normalized IN operands --- .../apache/comet/serde/QueryPlanSerde.scala | 2 +- .../org/apache/comet/serde/predicates.scala | 10 +- .../comet/rules/CometExecRuleSuite.scala | 114 +++++++++++++++++- 3 files changed, 121 insertions(+), 5 deletions(-) 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 6802dfaa64..c0a69d71f7 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/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 ++= _) 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 dfa2094898..69411bb776 100644 --- a/spark/src/main/scala/org/apache/comet/serde/predicates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/predicates.scala @@ -388,8 +388,10 @@ object ComparisonUtils { negate: Boolean): Option[Expr] = { // Spark treats all NaNs as equal and both signs of zero as equal in IN and InSet too. // Normalize both sides, including the fused NOT IN path that calls this method directly. - val valueExpr = exprToProtoInternal(normalizeInOperand(value), inputs, binding) - val listExprs = list.map(e => exprToProtoInternal(normalizeInOperand(e), inputs, binding)) + val normalizedValue = normalizeInOperand(value) + val normalizedList = list.map(normalizeInOperand) + 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) @@ -401,6 +403,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 } } 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 5444a89fa3..551385293e 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -23,7 +23,7 @@ import scala.util.Random import org.apache.spark.sql._ import org.apache.spark.sql.catalyst.FunctionIdentifier -import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionInfo} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression, ExpressionInfo, In, Literal, Not} import org.apache.spark.sql.catalyst.expressions.aggregate.BloomFilterAggregate import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec @@ -31,10 +31,12 @@ import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.QueryStageExec import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec} import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec} -import org.apache.spark.sql.types.{DataTypes, StructField, StructType} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{DataTypes, DoubleType, FloatType, StructField, StructType} import org.apache.comet.{CometConf, CometExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark42Plus} +import org.apache.comet.serde.QueryPlanSerde import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator} /** @@ -110,6 +112,114 @@ class CometExecRuleSuite extends CometTestBase { } } + for (dataType <- Seq(FloatType, DoubleType)) { + 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(1.0f), Literal(3.0f)) + case DoubleType => Seq(Literal(1.0d), 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 { + s"$column IN (CAST(1 AS ${dataType.sql}), CAST(3 AS ${dataType.sql}))" + } + 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 " is not supported" message used to hide. No serde in the tree is in that From d48800c93b089088724cb43af7dc5584e5756439 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Fri, 28 Aug 2026 16:41:43 +0000 Subject: [PATCH 4/5] fix: preserve pruning for finite floating IN lists --- .../org/apache/comet/serde/predicates.scala | 15 ++-- .../comet/rules/CometExecRuleSuite.scala | 72 +++++++++++++++++-- 2 files changed, 79 insertions(+), 8 deletions(-) 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 69411bb776..93e0cf609a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/predicates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/predicates.scala @@ -386,10 +386,17 @@ object ComparisonUtils { inputs: Seq[Attribute], binding: Boolean, negate: Boolean): Option[Expr] = { - // Spark treats all NaNs as equal and both signs of zero as equal in IN and InSet too. - // Normalize both sides, including the fused NOT IN path that calls this method directly. - val normalizedValue = normalizeInOperand(value) - val normalizedList = list.map(normalizeInOperand) + // NaNs and either sign of zero cannot match a finite 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.isFinite(v) && v != 0.0f + case Literal(v: Double, DoubleType) => java.lang.Double.isFinite(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)) { 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 551385293e..c066b7b3d4 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -23,8 +23,9 @@ import scala.util.Random import org.apache.spark.sql._ import org.apache.spark.sql.catalyst.FunctionIdentifier -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression, ExpressionInfo, In, Literal, Not} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression, ExpressionInfo, In, InSet, KnownFloatingPointNormalized, Literal, Not} import org.apache.spark.sql.catalyst.expressions.aggregate.BloomFilterAggregate +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._ @@ -113,6 +114,69 @@ class CometExecRuleSuite extends CometTestBase { } for (dataType <- Seq(FloatType, DoubleType)) { + test(s"floating ${dataType.sql} IN serialization preserves prunable finite 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 nullLiteral = Literal.create(null, dataType) + val lists: Seq[(Seq[Expression], Boolean)] = Seq( + ordinary -> false, + (ordinary :+ nullLiteral) -> false, + Seq(nullLiteral) -> false, + Seq(value, other) -> true) ++ + Seq(Double.NaN, 0.0d, -0.0d, Double.PositiveInfinity, Double.NegativeInfinity) + .map(v => (ordinary :+ literal(v)) -> true) ++ + (if (isSpark35Plus) Seq(Seq.empty[Expression] -> false) else Nil) + for ((list, needsNormalization) <- 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) + 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) + } + } + } + } + } + } + test( s"floating ${dataType.sql} IN serialization retains normalized operand fallback reasons") { val expressionNames = Seq("Literal", "KnownFloatingPointNormalized") @@ -127,8 +191,8 @@ class CometExecRuleSuite extends CometTestBase { val value = AttributeReference("value", dataType)() val other = AttributeReference("other", dataType)() val literals = dataType match { - case FloatType => Seq(Literal(1.0f), Literal(3.0f)) - case DoubleType => Seq(Literal(1.0d), Literal(3.0d)) + case FloatType => Seq(Literal(0.0f), Literal(3.0f)) + case DoubleType => Seq(Literal(0.0d), Literal(3.0d)) } // Exercise temporary literals and normalizers in both the value and the list. val in = @@ -179,7 +243,7 @@ class CometExecRuleSuite extends CometTestBase { val predicate = if (literalValue) { s"CAST(1 AS ${dataType.sql}) IN ($column, -$column)" } else { - s"$column IN (CAST(1 AS ${dataType.sql}), CAST(3 AS ${dataType.sql}))" + s"$column IN (CAST(0 AS ${dataType.sql}), CAST(3 AS ${dataType.sql}))" } val expression = if (negate) s"NOT ($predicate)" else predicate val df = sql(s"SELECT $expression AS hit FROM range(0, 4, 1, 1)") From be281b7b770534397eda13d8fc6ad30f94f2a02b Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Fri, 28 Aug 2026 18:40:44 +0000 Subject: [PATCH 5/5] Preserve Parquet pruning for infinity IN literals --- .../scala/org/apache/comet/serde/predicates.scala | 6 +++--- .../scala/org/apache/comet/CometExpressionSuite.scala | 8 ++++++-- .../org/apache/comet/rules/CometExecRuleSuite.scala | 11 ++++++++--- 3 files changed, 17 insertions(+), 8 deletions(-) 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 93e0cf609a..1eae1f92e3 100644 --- a/spark/src/main/scala/org/apache/comet/serde/predicates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/predicates.scala @@ -386,12 +386,12 @@ object ComparisonUtils { inputs: Seq[Attribute], binding: Boolean, negate: Boolean): Option[Expr] = { - // NaNs and either sign of zero cannot match a finite nonzero literal. Leave such lists + // 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.isFinite(v) && v != 0.0f - case Literal(v: Double, DoubleType) => java.lang.Double.isFinite(v) && v != 0.0d + 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. diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index 5d655134c3..1c189e2567 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -241,7 +241,9 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { (2, Some(-0.0f), Some(-0.0d)), (3, Some(13.0f), Some(13.0d)), (4, Some(1.0f), Some(1.0d)), - (5, None, None)) + (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( ( @@ -251,7 +253,9 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { java.lang.Float.intBitsToFloat(0xffc00002), java.lang.Double.longBitsToDouble(0xfff8000000000002L)), (0.0f, 0.0d), - (-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] = 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 c066b7b3d4..4866a3585e 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -114,7 +114,7 @@ class CometExecRuleSuite extends CometTestBase { } for (dataType <- Seq(FloatType, DoubleType)) { - test(s"floating ${dataType.sql} IN serialization preserves prunable finite literal lists") { + 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)() @@ -123,14 +123,19 @@ class CometExecRuleSuite extends CometTestBase { 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)] = Seq( ordinary -> false, (ordinary :+ nullLiteral) -> false, + infinities -> false, + (infinities :+ nullLiteral) -> false, Seq(nullLiteral) -> false, Seq(value, other) -> true) ++ - Seq(Double.NaN, 0.0d, -0.0d, Double.PositiveInfinity, Double.NegativeInfinity) - .map(v => (ordinary :+ literal(v)) -> true) ++ + Seq(Double.PositiveInfinity, Double.NegativeInfinity) + .map(v => (ordinary :+ literal(v)) -> false) ++ + Seq(Double.NaN, 0.0d, -0.0d) + .flatMap(v => Seq(ordinary, infinities).map(list => (list :+ literal(v)) -> true)) ++ (if (isSpark35Plus) Seq(Seq.empty[Expression] -> false) else Nil) for ((list, needsNormalization) <- lists; asSet <- Seq(false, true) if !asSet || list.forall(_.isInstanceOf[Literal]);