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 c69602fc80b..50907aa8ebd 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/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 6802dfaa646..c0a69d71f79 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 0e7bb02cd68..1eae1f92e30 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,19 @@ object ComparisonUtils { inputs: Seq[Attribute], binding: Boolean, negate: Boolean): Option[Expr] = { - val valueExpr = exprToProtoInternal(value, inputs, binding) - val listExprs = list.map(exprToProtoInternal(_, inputs, binding)) + // NaNs and either sign of zero cannot match a non-NaN nonzero literal. Leave such lists + // unwrapped so native Parquet scans can still prune using the column's statistics. + val needsNormalization = !list.forall { + case Literal(null, _) => true + case Literal(v: Float, FloatType) => !java.lang.Float.isNaN(v) && v != 0.0f + case Literal(v: Double, DoubleType) => !java.lang.Double.isNaN(v) && v != 0.0d + case _ => false + } + // Otherwise normalize both sides for Spark's NaN/zero equality, including fused NOT IN. + val normalizedValue = if (needsNormalization) normalizeInOperand(value) else value + val normalizedList = if (needsNormalization) list.map(normalizeInOperand) else list + val valueExpr = exprToProtoInternal(normalizedValue, inputs, binding) + val listExprs = normalizedList.map(exprToProtoInternal(_, inputs, binding)) if (valueExpr.isDefined && listExprs.forall(_.isDefined)) { val builder = ExprOuterClass.In.newBuilder() builder.setInValue(valueExpr.get) @@ -385,6 +410,10 @@ object ComparisonUtils { .setIn(builder) .build()) } else { + // Normalization creates temporary wrappers and literals outside the original tree. Keep + // their failure reasons on the membership expression so the operator can explain fallback. + liftFallbackReasons(normalizedValue, expr) + normalizedList.foreach(liftFallbackReasons(_, expr)) None } } diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index 80f92c4b7f5..1c189e25679 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,123 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + Seq( + ( + "float", + "_1", + Seq[Any]( + java.lang.Float.intBitsToFloat(0x7fc00001), + java.lang.Float.intBitsToFloat(0xffc00002))), + ( + "double", + "_2", + Seq[Any]( + java.lang.Double.longBitsToDouble(0x7ff8000000000001L), + java.lang.Double.longBitsToDouble(0xfff8000000000002L)))).foreach { + case (dataType, column, nanLiterals) => + test(s"compare $dataType columns with noncanonical NaN literals") { + withSQLConf(SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "false") { + val rows = Seq( + (Some(Float.NaN), Some(Double.NaN)), + (Some(-0.0f), Some(-0.0d)), + (Some(0.0f), Some(0.0d)), + (Some(-1.0f), Some(-1.0d)), + (Some(1.0f), Some(1.0d)), + (Some(Float.NegativeInfinity), Some(Double.NegativeInfinity)), + (Some(Float.PositiveInfinity), Some(Double.PositiveInfinity)), + (None, None)) + val identifiedRows = rows.zipWithIndex.map { case ((f, d), id) => (f, d, id) } + withParquetDataFrame(identifiedRows, withDictionary = false) { df => + // Parquet canonicalizes stored NaNs, so construct the signed/payload literals here. + // Compare Boolean results so Spark's NaN-aware answer checker cannot hide a mismatch. + val value = df(column) + nanLiterals.foreach { nan => + val literal = lit(nan) + Seq((value, literal), (literal, value)).foreach { case (left, right) => + val comparisons = Seq( + left === right, + left =!= right, + left.eqNullSafe(right), + left < right, + left <= right, + left > right, + left >= right) + checkSparkAnswerAndOperator( + df.select(comparisons: _*), + Seq(classOf[CometProjectExec])) + comparisons.foreach { comparison => + // Compare surviving identities, not just NaN-aware row values. Keep Parquet + // pushdown disabled so every ordering predicate executes in CometFilterExec. + checkSparkAnswerAndOperator( + df.filter(comparison).select("_3"), + Seq(classOf[CometFilterExec])) + } + } + } + } + } + } + } + + for ((name, threshold) <- Seq(("In", 10), ("InSet", 1))) { + test(s"floating $name and NOT $name normalize NaNs and signed zeros") { + withSQLConf( + SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "false", + "spark.sql.optimizer.inSetConversionThreshold" -> threshold.toString) { + val rows = Seq( + (0, Some(Float.NaN), Some(Double.NaN)), + (1, Some(0.0f), Some(0.0d)), + (2, Some(-0.0f), Some(-0.0d)), + (3, Some(13.0f), Some(13.0d)), + (4, Some(1.0f), Some(1.0d)), + (5, None, None), + (6, Some(Float.PositiveInfinity), Some(Double.PositiveInfinity)), + (7, Some(Float.NegativeInfinity), Some(Double.NegativeInfinity))) + withParquetDataFrame(rows, withDictionary = false) { df => + val cases = Seq( + ( + java.lang.Float.intBitsToFloat(0x7fc00001), + java.lang.Double.longBitsToDouble(0x7ff8000000000001L)), + ( + java.lang.Float.intBitsToFloat(0xffc00002), + java.lang.Double.longBitsToDouble(0xfff8000000000002L)), + (0.0f, 0.0d), + (-0.0f, -0.0d), + (Float.PositiveInfinity, Double.PositiveInfinity), + (Float.NegativeInfinity, Double.NegativeInfinity)) + for ((f, d) <- cases; includeNull <- Seq(false, true)) { + val floatCandidates: Seq[Any] = Seq(f, 13.0f) ++ (if (includeNull) Seq(null) else Nil) + val doubleCandidates: Seq[Any] = + Seq(d, 13.0d) ++ (if (includeNull) Seq(null) else Nil) + // Negation creates negative NaNs after the Parquet scan, so the membership value + // must be normalized as well as the programmatically constructed list literals. + val predicates = Seq(df("_2"), -df("_2")).map(_.isin(floatCandidates: _*)) ++ + Seq(df("_3"), -df("_3")).map(_.isin(doubleCandidates: _*)) + val projected = df.select(df("_1") +: predicates.flatMap(p => Seq(p, !p)): _*) + val optimized = projected.queryExecution.optimizedPlan + val membership = optimized.expressions.flatMap(_.collect { + case _: org.apache.spark.sql.catalyst.expressions.In => "In" + case _: org.apache.spark.sql.catalyst.expressions.InSet => "InSet" + }) + // A singleton IN is rewritten to equality and would not exercise the faulty kernel. + assert(membership.nonEmpty && membership.forall(_ == name), optimized.toString) + checkSparkAnswerAndOperator(projected, Seq(classOf[CometProjectExec])) + for (p <- predicates; negate <- Seq(false, true)) { + val filtered = df.filter(if (negate) !p else p).select("_1") + if (includeNull && negate) { + // NOT IN with a null candidate can never be true. Spark legitimately replaces + // this filter with an empty LocalRelation before physical planning. + checkSparkAnswer(filtered) + } else { + checkSparkAnswerAndOperator(filtered, Seq(classOf[CometFilterExec])) + } + } + } + } + } + } + } + test("parquet default values") { withTable("t1") { sql("create table t1(col1 boolean) using parquet") 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 5444a89fa36..4866a3585e7 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -23,18 +23,21 @@ 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, 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._ 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 +113,182 @@ 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)] = Seq( + ordinary -> false, + (ordinary :+ nullLiteral) -> false, + infinities -> false, + (infinities :+ nullLiteral) -> false, + Seq(nullLiteral) -> false, + Seq(value, other) -> 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]); + 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") + 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(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 = + 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(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)") + 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