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]
