From c0577e39b9ce42d202bf0d18e27ae96e84480b33 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 4 Oct 2026 17:26:45 +0000 Subject: [PATCH 1/9] fix: preserve skipped unbase64 rows --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../expression-audits/string_funcs.md | 4 + .../user-guide/latest/compatibility/index.md | 14 + docs/source/user-guide/latest/expressions.md | 2 +- .../scala/org/apache/comet/CometConf.scala | 11 + .../apache/comet/rules/CometExecRule.scala | 162 ++++- .../comet/serde/CometExpressionSerde.scala | 14 + .../apache/comet/serde/QueryPlanSerde.scala | 18 + .../org/apache/comet/serde/strings.scala | 20 +- .../misc/unbase64_operator_masks.sql | 47 ++ .../serde/CometEvaluationMaskSuite.scala | 648 ++++++++++++++++++ 12 files changed, 934 insertions(+), 8 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/misc/unbase64_operator_masks.sql create mode 100644 spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 44cd5c56a9f..528bdf7b9ca 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -621,6 +621,7 @@ jobs: org.apache.comet.CometStringDecodeSuite org.apache.comet.CometWidthBucketSuite org.apache.comet.CometUuidExpressionSuite + org.apache.comet.serde.CometEvaluationMaskSuite org.apache.comet.serde.CometScalarFunctionSuite org.apache.comet.serde.CometLiteralSuite org.apache.comet.serde.QueryPlanSerdeSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index dbc76c05d46..d4fff1a3d82 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -325,6 +325,7 @@ jobs: org.apache.comet.CometStringDecodeSuite org.apache.comet.CometWidthBucketSuite org.apache.comet.CometUuidExpressionSuite + org.apache.comet.serde.CometEvaluationMaskSuite org.apache.comet.serde.CometScalarFunctionSuite org.apache.comet.serde.CometLiteralSuite org.apache.comet.serde.QueryPlanSerdeSuite diff --git a/docs/source/contributor-guide/expression-audits/string_funcs.md b/docs/source/contributor-guide/expression-audits/string_funcs.md index 12721470b4a..96ce449acc9 100644 --- a/docs/source/contributor-guide/expression-audits/string_funcs.md +++ b/docs/source/contributor-guide/expression-audits/string_funcs.md @@ -253,6 +253,10 @@ - Spark 4.0.1 (audited 2026-08-24): `NullIntolerant` becomes `override def nullIntolerant: Boolean = true`; `inputTypes` widens to `StringTypeWithCollation(supportsTrimCollation = true)`. Behaviour is byte-level and collation-independent, so no divergence for `UTF8_BINARY` and no shim is needed. - Spark 4.1.1 (audited 2026-08-24): adds `contextIndependentFoldable`; no behavioural change on the decode path. - Comet native implementation (`spark_unbase64`, `native/spark-expr/src/string_funcs/unbase64.rs`) ports the JDK MIME decoder rules: 256-entry decode LUT, a reused per-batch scratch `Vec` copied into a preallocated `BinaryBuilder`, all four error messages reproduced verbatim. `CometUnBase64` handles `failOnError = false` natively; `failOnError = true` (reachable from `to_binary('base64')` / `try_to_binary`) requires strict RFC 4648 validation and is not yet implemented natively, so those cases stay on the JVM codegen dispatcher via `CodegenDispatchFallback`. +- Native decoding is limited to column and literal children. Compound children use JVM codegen dispatch when enabled, otherwise they fall back to Spark. Both native decoding and JVM dispatch evaluate batches, so `CometExecRule` also preserves Spark's row pipeline where a limit, preordered top-K or unpartitioned `WindowGroupLimit`, or a first-match semi/anti join condition can skip malformed input. Partitioned window limits drain their groups and do not independently require this fallback, but an outer limit can still protect them. Ordinary decoding projections below a blocking sort or exchange remain eligible unless another limiting consumer intervenes. +- With AQE and an outer limit, a final aggregate that decodes its results and uses incompatible intermediate buffers stays in Spark together with its partial aggregate before stages materialize. This also applies when an intervening sort currently consumes all input, because AQE can remove that sort later. A final aggregate only merging already-decoded inputs does not need this result-evaluation protection. +- `spark.comet.exec.preserveEvaluationMasks.enabled` defaults to `true`. Users who know their Base64 input is valid can set it to `false` to retain native execution across these operator boundaries, including the AQE aggregate protection above. With this opt-out, eager batch evaluation can raise errors on malformed values in rows Spark would skip. Expression support restrictions, including strict decoding and compound-child dispatch, still apply. +- Remaining evaluation-mask gaps, including other throwing expressions and enclosing filter/conditional masks, are tracked in [#6006](https://github.com/apache/datafusion-comet/issues/6006). ## upper diff --git a/docs/source/user-guide/latest/compatibility/index.md b/docs/source/user-guide/latest/compatibility/index.md index 017140b14dd..6e23c71752a 100644 --- a/docs/source/user-guide/latest/compatibility/index.md +++ b/docs/source/user-guide/latest/compatibility/index.md @@ -119,6 +119,20 @@ divergence: - Spark 4.2 introduced additional ANSI arithmetic overflow behavior differences that Comet does not yet track ([#4967](https://github.com/apache/datafusion-comet/issues/4967)). +## Errors from rows Spark skips + +Batch evaluation can raise a data-dependent error on a row that Spark's row pipeline would +never evaluate. Comet preserves operator-level evaluation masks for `unbase64` under limits, +first-match semi/anti join conditions, and ordered top-K or unpartitioned window limits. +This protection is enabled by default. Setting +`spark.comet.exec.preserveEvaluationMasks.enabled=false` retains native execution for known-valid +input, but can raise errors from malformed Base64 in rows Spark would skip. + +This policy currently enrolls only `UnBase64`; it does not establish error-evaluation parity for +ANSI arithmetic, casts, decimal division, or `element_at`. Per-row conditional evaluation +(`AND`, `CASE`) and aggregate `FILTER` require separate checks. The broader audit and remaining +work are tracked in [#6006](https://github.com/apache/datafusion-comet/issues/6006). + ## Known result-value divergences The following native paths silently return values that differ from Spark for edge-case inputs. diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 74232aac652..49c79e25617 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -625,7 +625,7 @@ to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enab | `try_to_binary` | ✅ | — | Rewrites to `try_eval(to_binary(...))`, which routes through the JVM codegen dispatcher | | `try_to_number` | ✅ | Codegen dispatch | Routed through the JVM codegen dispatcher | | `ucase` | ✅ | Hybrid | | -| `unbase64` | ✅ | Codegen dispatch | Column or literal input runs natively; other child expressions, and the strict form used by `to_binary(str, 'base64')`, route through the JVM codegen dispatcher | +| `unbase64` | ✅ | Codegen dispatch | Column or literal input runs natively; other child expressions, and the strict form used by `to_binary(str, 'base64')`, route through the JVM codegen dispatcher. Where limits, preordered top-K or unpartitioned window limits, or first-match semi/anti join conditions can skip malformed input, the affected pipeline stays in Spark's row execution. For known-valid input, set `spark.comet.exec.preserveEvaluationMasks.enabled=false` to retain native execution. Disabling this protection can raise errors on malformed values in rows Spark would skip. | | `upper` | ✅ | Hybrid | | --- diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index f2f19ac5ce9..905ba4bd221 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -260,6 +260,17 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(true) + val COMET_EXEC_PRESERVE_EVALUATION_MASKS_ENABLED: ConfigEntry[Boolean] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.preserveEvaluationMasks.enabled") + .category(CATEGORY_EXEC) + .doc( + "Preserve Spark's row evaluation when operators such as LIMIT or semi/anti joins " + + "can skip inputs to throwing expressions. Currently this applies to unbase64. " + + "Set to false to retain native execution when the input is known to be valid. " + + "When disabled, Comet may fail on malformed Base64 in rows Spark would skip.") + .booleanConf + .createWithDefault(true) + val COMET_EXEC_PROJECT_ENABLED: ConfigEntry[Boolean] = createExecEnabledConfig("project", defaultValue = true) val COMET_EXEC_FILTER_ENABLED: ConfigEntry[Boolean] = 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 09f6f726a39..206147a24db 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -25,9 +25,10 @@ import scala.collection.mutable.ListBuffer import scala.jdk.CollectionConverters._ import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder} +import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder, SortOrder} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero +import org.apache.spark.sql.catalyst.plans.{JoinType, LeftAnti, LeftSemi} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreeNodeTag import org.apache.spark.sql.catalyst.util.sideBySide @@ -50,7 +51,7 @@ import org.apache.spark.sql.execution.datasources.v2.csv.CSVScan import org.apache.spark.sql.execution.datasources.v2.json.JsonScan import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike} -import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, HashJoin, ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -93,6 +94,9 @@ object CometExecRule { val COMET_UNSAFE_PARTIAL: TreeNodeTag[String] = TreeNodeTag[String]("comet.unsafePartialAgg") + private[rules] val UNSAFE_EXPRESSION_EVALUATION: TreeNodeTag[String] = + TreeNodeTag[String]("comet.unsafeExpressionEvaluation") + /** * Info message for a native operator that keeps a `Scan` over an input that now reads its * shuffle directly, because its native plan's leaves do not line up with its inputs. @@ -391,6 +395,9 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) // spotless:on private def transform(plan: SparkPlan): SparkPlan = { def convertNode(op: SparkPlan): SparkPlan = op match { + case op if op.getTagValue(CometExecRule.UNSAFE_EXPRESSION_EVALUATION).isDefined => + withFallbackReason(op, op.getTagValue(CometExecRule.UNSAFE_EXPRESSION_EVALUATION).get) + // Scan marker produced by an optional, out-of-tree scan contrib (e.g. contrib/delta). // Matched by trait (no compile-time dependency on the contrib) and present only when that // contrib is on the classpath. The marker carries its own serde handler and typically wraps @@ -971,9 +978,10 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) // corresponding Final or PartialMerge cannot be converted and the intermediate buffer // formats are incompatible. This runs before transform() so the tags are checked // during the bottom-up conversion. Tags persist through AQE stage creation. - tagUnsafePartialAggregates(planWithJoinRewritten) + val planWithEvaluationMasks = preserveEvaluationMasks(planWithJoinRewritten) + tagUnsafePartialAggregates(planWithEvaluationMasks) - var newPlan = revertUnsafePartialAggregates(transform(planWithJoinRewritten)) + var newPlan = revertUnsafePartialAggregates(transform(planWithEvaluationMasks)) // if the plan cannot be run fully natively then explain why (when appropriate // config is enabled) @@ -1086,6 +1094,150 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } } + /** Keep opted-in expressions in Spark's row pipeline where an operator can skip inputs. */ + private def preserveEvaluationMasks(plan: SparkPlan): SparkPlan = { + if (!CometConf.COMET_EXEC_PRESERVE_EVALUATION_MASKS_ENABLED.get()) return plan + + def findEvaluationMaskName(expr: Expression): Option[String] = { + var name: Option[String] = None + expr.exists { child => + name = QueryPlanSerde.evaluationMaskName(child) + name.isDefined + } + name + } + + def originalPlan(node: SparkPlan): SparkPlan = node match { + case scan: CometScanExec => + scan.wrapped + .copy(partitionFilters = scan.partitionFilters, dataFilters = scan.dataFilters) + case comet: CometExec => comet.originalPlan + case shuffle: CometShuffleExchangeExec => shuffle.originalPlan + case broadcast: CometBroadcastExchangeExec => broadcast.originalPlan + case _ => node + } + + // Most plans contain no opted-in expression. Include native originals and sticky tags + // so repeated AQE passes still restore an already-protected subtree when necessary. + if (!plan.exists(node => + node.getTagValue(CometExecRule.UNSAFE_EXPRESSION_EVALUATION).isDefined || + originalPlan(node).expressions.exists(findEvaluationMaskName(_).isDefined))) { + return plan + } + + def firstMatch(joinType: JoinType): Boolean = joinType match { + case LeftSemi | LeftAnti => true + case _ => false + } + + def protect( + node: SparkPlan, + belowLimit: Boolean, + hasLimitAncestor: Boolean): (SparkPlan, Option[String]) = { + val original = originalPlan(node) + val startsLimit = original match { + // Offset-only collection does not stop its input early. + case collect: CollectLimitExec => collect.limit >= 0 + case _: LocalLimitExec | _: GlobalLimitExec => true + case topK: TakeOrderedAndProjectExec => + SortOrder.orderingSatisfies(node.children.head.outputOrdering, topK.sortOrder) + case windowLimit + if ShimCometWindowGroupLimit.windowGroupLimitClass.exists( + _.isInstance(windowLimit)) => + // Partitioned limits drain each group, but an outer LIMIT can still stop them. + // Keep the conservative behavior if a future rank function cannot be extracted. + ShimCometWindowGroupLimit.extract(windowLimit).forall(_.partitionSpec.isEmpty) && + SortOrder.orderingSatisfies( + node.children.head.outputOrdering, + windowLimit.requiredChildOrdering.head) + case _ => false + } + // These operators consume their input before yielding rows. Still visit their children: + // an inner LocalLimit below an exchange must establish its own evaluation boundary. + val materializesInput = original match { + case _: SortExec | _: HashAggregateExec | _: ObjectHashAggregateExec | + _: ShuffleExchangeLike | _: BroadcastExchangeLike | _: QueryStageExec | + _: ReusedExchangeExec => + true + case _ => false + } + // A top-K is still a LIMIT even when its current input requires sorting. + val limitAncestor = hasLimitAncestor || startsLimit || + original.isInstanceOf[TakeOrderedAndProjectExec] + val protectedChildren = node.children.map { child => + protect(child, startsLimit || (belowLimit && !materializesInput), limitAncestor) + } + val childReason = protectedChildren.flatMap(_._2).headOption + val condition = original match { + case join: HashJoin if firstMatch(join.joinType) => join.condition + case join: SortMergeJoinExec if firstMatch(join.joinType) => join.condition + case join: BroadcastNestedLoopJoinExec if firstMatch(join.joinType) => join.condition + case _ => None + } + val finalAggregate = original match { + case agg: BaseAggregateExec + if (agg.isInstanceOf[HashAggregateExec] || + agg.isInstanceOf[ObjectHashAggregateExec]) && + agg.aggregateExpressions.map(_.mode).distinct == Seq(Final) => + Some(agg) + case _ => None + } + val limitName = if (belowLimit) { + // Final merges buffers; it does not reevaluate the aggregate's original inputs. + val expressions = finalAggregate.map(_.resultExpressions).getOrElse(original.expressions) + expressions.iterator.flatMap(findEvaluationMaskName).take(1).toSeq.headOption + } else { + None + } + // AQE can remove an intervening sort after a native Partial has materialized. Choose + // compatible buffers before that happens, even if a current operator drains its input. + val aggregateBufferName = if (hasLimitAncestor && conf.adaptiveExecutionEnabled) { + finalAggregate + .filterNot(agg => + QueryPlanSerde.allAggsSupportNativePartialToSparkFinal(agg.aggregateExpressions)) + .flatMap( + _.resultExpressions.iterator.flatMap(findEvaluationMaskName).take(1).toSeq.headOption) + } else { + None + } + val ownReason = limitName + .map(name => s"$name requires Spark evaluation below LIMIT") + .orElse(aggregateBufferName.map(name => + s"$name requires Spark aggregate buffers below LIMIT with AQE")) + .orElse( + condition + .flatMap(findEvaluationMaskName) + .map(name => s"$name requires Spark evaluation in first-match join conditions")) + .orElse(node.getTagValue(CometExecRule.UNSAFE_EXPRESSION_EVALUATION)) + // Exchanges can restart native execution after consuming a Spark row pipeline. + val restartsNative = original.isInstanceOf[ShuffleExchangeLike] || + original.isInstanceOf[BroadcastExchangeLike] + val reason = ownReason.orElse(if (restartsNative) None else childReason) + // Early tags keep fresh aggregate buffers safe before AQE materialization. Reused native + // buffer chains are restored by revertUnsafePartialAggregates after conversion. + val children = protectedChildren.map(_._1) + val prepared = node match { + // Do not refill a batch between the row decoder and its short-circuiting consumer. + case _: RowToColumnarExec | _: CometSparkToColumnarExec if childReason.isDefined => + children.head + case _: ColumnarToRowExec | _: CometColumnarToRowExec | _: CometNativeColumnarToRowExec + if childReason.isDefined && !children.head.supportsColumnar => + children.head + case _ if (original ne node) && (reason.isDefined || children != node.children) => + // AQE can reuse an existing native subtree. Rebuild affected ancestors as well so + // their serialized native plans do not retain the decoder that just fell back. + val restored = original.withNewChildren(children) + node.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).foreach(restored.setLogicalLink) + restored + case _ => node.withNewChildren(children) + } + reason.foreach(prepared.setTagValue(CometExecRule.UNSAFE_EXPRESSION_EVALUATION, _)) + (prepared, reason) + } + + protect(plan, belowLimit = false, hasLimitAncestor = false)._1 + } + /** Convert a Spark plan to a Comet plan using the specified serde handler */ private def convertToComet(op: SparkPlan, handler: CometOperatorSerde[_]): Option[SparkPlan] = { val converted = tryConvertToComet(op, handler) @@ -1583,6 +1735,8 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) private def canAggregateBeConverted( agg: BaseAggregateExec, expectedMode: AggregateMode): Boolean = { + if (agg.getTagValue(CometExecRule.UNSAFE_EXPRESSION_EVALUATION).isDefined) return false + val handler = allExecs.get(agg.getClass) if (handler.isEmpty) return false val serde = handler.get.asInstanceOf[CometOperatorSerde[SparkPlan]] diff --git a/spark/src/main/scala/org/apache/comet/serde/CometExpressionSerde.scala b/spark/src/main/scala/org/apache/comet/serde/CometExpressionSerde.scala index 206bd7c8207..9ae1a3d4585 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometExpressionSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometExpressionSerde.scala @@ -124,6 +124,20 @@ trait CometExpressionSerde[T <: Expression] { */ trait CodegenDispatchFallback extends NativeOptInAvailable { self: CometExpressionSerde[_] => } +/** + * Opts an expression instance into the planner's existing operator-level evaluation protection. + * This is not an exhaustive classification of expressions that can throw, nor a guarantee that + * all evaluation contexts are protected. The policy is independent of native or JVM-dispatch + * support and must inspect the expression without evaluating it or serializing it. + * + * Return a stable function name for fallback explanations, or None when the instance is not + * enrolled. None does not prove safety. Future enrollment requires separate compatibility and + * performance tests, including evaluation mode, input types, and enclosing expressions. + */ +trait RequiresSparkEvaluationMask[T <: Expression] { self: CometExpressionSerde[T] => + def evaluationMaskName(expr: T): Option[String] +} + /** * Marker for serdes that have a native implementation the user can opt into. Normally these * serdes run a Spark-compatible path by default. A serde can set `hasConditionalNativeDefault` 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 430d36289c1..38ed0355b7c 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -409,6 +409,24 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { val exprSerdeMap: Map[Class[_ <: Expression], CometExpressionSerde[_]] = serdeGroups.map(_._2).reduce(_ ++ _) + // Cache enrollment, not the per-instance policy result. Include subclasses because a whole-tree + // JVM dispatcher can evaluate them without serializing each child through the exact-class map. + // Unmarked subclasses cannot opt out of an ancestor's protection. If several marked classes + // match, any Some result requires protection; class-name order only chooses the explanation. + private val evaluationMaskSerdes = exprSerdeMap.toSeq + .collect { case (exprClass, serde: RequiresSparkEvaluationMask[_]) => + exprClass -> serde.asInstanceOf[RequiresSparkEvaluationMask[Expression]] + } + .sortBy(_._1.getName) + + private[comet] def evaluationMaskName(expr: Expression): Option[String] = + evaluationMaskSerdes.iterator + .filter { case (exprClass, _) => exprClass.isInstance(expr) } + .flatMap { case (_, serde) => serde.evaluationMaskName(expr) } + .take(1) + .toSeq + .headOption + /** * Mapping of Spark aggregate expression class to Comet expression handler. */ diff --git a/spark/src/main/scala/org/apache/comet/serde/strings.scala b/spark/src/main/scala/org/apache/comet/serde/strings.scala index 68d0e252df8..cd4166abdab 100644 --- a/spark/src/main/scala/org/apache/comet/serde/strings.scala +++ b/spark/src/main/scala/org/apache/comet/serde/strings.scala @@ -846,12 +846,20 @@ object CometBase64 extends CometExpressionSerde[Base64] { // `unbase64(...)`) only when the child is a column reference or literal, since native // ScalarFunctionExpr evaluates its arguments eagerly and would bypass Spark's short-circuit // semantics for compound children (see apache/datafusion-comet#5451). More complex children -// stay on the JVM codegen dispatcher via CodegenDispatchFallback. When failOnError = true +// stay on the JVM codegen dispatcher via CodegenDispatchFallback. CometExecRule preserves +// row evaluation below LIMIT and in first-match join conditions for both execution paths, +// unless spark.comet.exec.preserveEvaluationMasks.enabled is disabled for known-valid input. +// When failOnError = true // (from `to_binary('base64')` / `try_to_binary`), Spark uses a stricter RFC 4648 validator, so // those cases also stay on the dispatcher. Error messages match Spark byte-for-byte (pinned in // the Rust unit tests), but the wrapping exception class does not; kept as Compatible() because // Spark surfaces these as bare IllegalArgumentException without a SQL error class. -object CometUnBase64 extends CometExpressionSerde[UnBase64] with CodegenDispatchFallback { +object CometUnBase64 + extends CometExpressionSerde[UnBase64] + with CodegenDispatchFallback + with RequiresSparkEvaluationMask[UnBase64] { + + override def evaluationMaskName(expr: UnBase64): Option[String] = Some("unbase64") private val failOnErrorReason = "unbase64 with failOnError = true uses stricter RFC 4648 validation that is not yet" + @@ -861,8 +869,14 @@ object CometUnBase64 extends CometExpressionSerde[UnBase64] with CodegenDispatch "unbase64 with a non-trivial child expression uses the JVM codegen dispatcher to preserve" + " Spark's short-circuit evaluation (native path is limited to column and literal children)" + private val maskedEvaluationReason = + "unbase64 below LIMIT or in first-match join conditions falls back to Spark so malformed" + + " input is not decoded on rows Spark skips. For known-valid input, set " + + "spark.comet.exec.preserveEvaluationMasks.enabled=false to retain native execution; " + + "malformed input may then fail even on rows Spark would skip" + override def getUnsupportedReasons(): Seq[String] = - Seq(failOnErrorReason, nonTrivialChildReason) + Seq(failOnErrorReason, nonTrivialChildReason, maskedEvaluationReason) override def getSupportLevel(expr: UnBase64): SupportLevel = { if (expr.failOnError) { diff --git a/spark/src/test/resources/sql-tests/expressions/misc/unbase64_operator_masks.sql b/spark/src/test/resources/sql-tests/expressions/misc/unbase64_operator_masks.sql new file mode 100644 index 00000000000..4d06ba54eaa --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/misc/unbase64_operator_masks.sql @@ -0,0 +1,47 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- Regression for https://github.com/apache/datafusion-comet/issues/5532. +-- Both native decoding and JVM dispatch must preserve rows skipped by LIMIT. +-- ConfigMatrix: spark.comet.exec.scalaUDF.codegen.enabled=false,true +-- Config: spark.sql.shuffle.partitions=1 + +-- One writer preserves the VALUES order in a single Parquet file: the valid value precedes +-- the malformed suffix. This pins physical scan order, not SQL ordering without ORDER BY. +statement +CREATE TABLE test_unbase64_operator_limit USING parquet AS +SELECT /*+ COALESCE(1) */ bad FROM VALUES ('YWJj'), ('A') AS t(bad) + +-- Spark stops after decoding 'abc' and never evaluates the malformed second row. +query expect_fallback(unbase64 requires Spark evaluation below LIMIT) +SELECT hex(unbase64(bad)) FROM test_unbase64_operator_limit LIMIT 1 + +query expect_fallback(unbase64 requires Spark evaluation below LIMIT) +SELECT bad FROM test_unbase64_operator_limit +WHERE unbase64(bad) <=> X'616263' LIMIT 1 + +-- A compound argument reaches the dispatcher, which also evaluates whole batches. +query expect_fallback(unbase64 requires Spark evaluation below LIMIT) +SELECT hex(unbase64(concat(bad, ''))) FROM test_unbase64_operator_limit LIMIT 1 + +-- Consuming the malformed row must still raise the decoder error. +query expect_error(Last unit does not have enough valid bits) +SELECT hex(unbase64(bad)) FROM test_unbase64_operator_limit WHERE bad = 'A' LIMIT 1 + +-- Without an early-stop operator, decoding valid input retains the native kernel. +query expect_native(unbase64) +SELECT hex(unbase64(bad)) FROM test_unbase64_operator_limit WHERE bad = 'YWJj' diff --git a/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala b/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala new file mode 100644 index 00000000000..e206f0e2f07 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala @@ -0,0 +1,648 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.serde + +import scala.reflect.ClassTag + +import org.apache.spark.sql.{CometTestBase, DataFrame, Row} +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.optimizer.BuildRight +import org.apache.spark.sql.catalyst.plans.logical.LocalRelation +import org.apache.spark.sql.catalyst.plans.physical.SinglePartition +import org.apache.spark.sql.comet._ +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, LogicalQueryStage} +import org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ShuffleExchangeExec, ShuffleExchangeLike} +import org.apache.spark.sql.execution.joins._ +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{IntegerType, StringType} + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus +import org.apache.comet.rules.CometExecRule +import org.apache.comet.shims.ShimCometWindowGroupLimit + +class CometEvaluationMaskSuite extends CometTestBase { + private val limitReason = "unbase64 requires Spark evaluation below LIMIT" + private val joinReason = "unbase64 requires Spark evaluation in first-match join conditions" + private val decodeMessage = "Last unit does not have enough valid bits" + + private def withModes(f: (Boolean, Boolean) => Unit): Unit = { + for (aqe <- Seq(false, true); dispatch <- Seq(false, true)) { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString, + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> dispatch.toString) { + withClue(s"AQE=$aqe, dispatcher=$dispatch: ")(f(aqe, dispatch)) + } + } + } + + private def withInputs(inputs: (String, String)*)(f: => Unit): Unit = { + withTempPath { dir => + withTempView(inputs.map(_._1): _*) { + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + inputs.foreach { case (name, query) => + val path = s"${dir.getCanonicalPath}/$name" + sql(query).coalesce(1).write.parquet(path) + val input = spark.read.parquet(path) + assert(input.inputFiles.length == 1) + input.createOrReplaceTempView(name) + } + } + f + } + } + } + + private def sparkPlan(query: String): SparkPlan = { + var plan: SparkPlan = null + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + plan = sql(query).queryExecution.executedPlan + } + plan + } + + private def applyRule(plan: SparkPlan): SparkPlan = + CometExecRule(spark).apply(stripAQEPlan(plan)) + private def count[T <: SparkPlan: ClassTag](plan: SparkPlan): Int = + collect(plan) { case node: T => node }.size + private def aggregateCounts(plan: SparkPlan): (Int, Int) = + (count[CometHashAggregateExec](plan), count[ObjectHashAggregateExec](plan)) + private def nativeScans(plan: SparkPlan): Int = collect(plan) { + case _: CometScanExec | _: CometNativeScanExec => true + }.size + private def original(plan: SparkPlan): SparkPlan = plan match { + case comet: CometExec => comet.originalPlan + case other => other + } + private def hasDecoder(expr: Expression): Boolean = expr.exists(_.isInstanceOf[UnBase64]) + private def malformed(error: Throwable): Boolean = + causeChain(error).exists(e => Option(e.getMessage).exists(_.contains(decodeMessage))) + private def decodeError(body: => Any): Unit = assert(malformed(intercept[Exception](body))) + private def decodeErrors(query: String): Unit = { + val (sparkError, cometError) = checkSparkAnswerMaybeThrows(sql(query)) + Seq(sparkError, cometError).foreach(error => assert(error.exists(malformed), query)) + } + + test("evaluation-mask protection can be disabled for known-valid input") { + val preserve = CometConf.COMET_EXEC_PRESERVE_EVALUATION_MASKS_ENABLED + assert(preserve.defaultValue.contains(true)) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + withInputs( + "opt_out_valid" -> "SELECT * FROM VALUES (1, 'YWJj'), (1, 'YWFh') AS t(k, bad)", + "opt_out_left" -> "SELECT * FROM VALUES (1, X'616262'), (2, X'616262') AS t(k, expected)") { + for (enabled <- Seq(true, false)) { + withSQLConf(preserve.key -> enabled.toString) { + val projection = "SELECT hex(unbase64(bad)) FROM opt_out_valid LIMIT 1" + val compound = "SELECT hex(unbase64(concat(bad, ''))) FROM opt_out_valid LIMIT 1" + if (enabled) { + checkSparkAnswerAndFallbackReason(projection, limitReason) + checkSparkAnswerAndFallbackReason(compound, limitReason) + } else { + checkSparkAnswerAndImpl(sql(projection), native = Seq("unbase64")) + checkSparkAnswerAndImpl(sql(compound), dispatched = Seq("unbase64")) + } + for (kind <- Seq("SEMI", "ANTI")) { + val query = s"""SELECT /*+ BROADCAST(r) */ l.k + |FROM opt_out_left l LEFT $kind JOIN opt_out_valid r + |ON l.k = r.k AND unbase64(r.bad) > l.expected""".stripMargin + if (enabled) { + checkSparkAnswerAndFallbackReason(query, joinReason) + } else { + val (_, plan) = checkSparkAnswerAndImpl(sql(query), native = Seq("unbase64")) + assert(count[CometBroadcastHashJoinExec](plan) == 1, plan.toString) + } + } + // AQE can remove the sort after Partial materializes. The opt-out keeps both + // aggregate halves native instead of preemptively restoring Spark's buffers. + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val aggregate = sparkPlan( + "SELECT unbase64(bad) AS decoded, collect_list(k) FROM " + + "VALUES (1, 'YWJj'), (1, 'YWFh') AS t(k, bad) GROUP BY bad") + assert(aggregateCounts(aggregate) == (0, 2)) + val sorted = SortExec( + Seq(SortOrder(aggregate.output.head, Ascending)), + global = false, + aggregate) + val plan = applyRule(CollectLimitExec(1, sorted)) + assert(aggregateCounts(plan) == (if (enabled) (0, 2) else (2, 0)), plan.toString) + } + } + } + } + } + } + + test("LIMIT masks follow the decoder policy with ANSI, strict and compound inputs") { + assert(QueryPlanSerde.exprSerdeMap.collect { case (cls, _: RequiresSparkEvaluationMask[_]) => + cls + }.toSet == Set(classOf[UnBase64])) + val input = AttributeReference("encoded", StringType)() + withInputs( + "masked" -> "SELECT * FROM VALUES (1, 'YWJj'), (2, 'A') AS t(k, bad)", + "valid" -> "SELECT * FROM VALUES (1, 'YWJj'), (1, 'YWFh'), (1, NULL) AS t(k, bad)", + "inner_left" -> "SELECT 1 AS k, X'616262' AS expected", + "one_group" -> "SELECT 1 AS k, 'YWJj' AS bad") { + withModes { (_, _) => + for (ansi <- Seq(false, true)) { + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansi.toString) { + for (child <- Seq(input, Literal("YWJj"), Concat(Seq(input, Literal("")))); + strict <- Seq(false, true)) { + val expressions = Seq(UnBase64(child, strict), new LimitUnBase64(child, strict)) + assert( + expressions.forall(QueryPlanSerde.evaluationMaskName(_).contains("unbase64"))) + } + // Existing unbase64/to_binary SQL tests exercise native and dispatched serialization. + // Other throwing expressions: https://github.com/apache/datafusion-comet/issues/6006. + Seq( + Add(Literal(Int.MaxValue), Literal(1)), + Cast(input, IntegerType), + Concat(Seq(input, Literal(""))), + FormatString(Literal("%s"), input), + Like(input, Literal("%"), escapeChar = '\\')).foreach { expr => + assert(QueryPlanSerde.evaluationMaskName(expr).isEmpty) + } + Seq( + "SELECT hex(unbase64(bad)) FROM masked LIMIT 1", + "SELECT bad FROM masked WHERE unbase64(bad) <=> X'616263' LIMIT 1", + "SELECT hex(unbase64(concat(bad, ''))) FROM masked LIMIT 1", + "SELECT hex(to_binary(bad, 'base64')) FROM masked LIMIT 1").foreach { query => + checkSparkAnswerAndFallbackReason(query, limitReason) + } + decodeErrors("SELECT hex(unbase64(bad)) FROM masked") + decodeErrors("SELECT hex(unbase64(bad)) FROM masked WHERE bad = 'A' LIMIT 1") + checkSparkAnswerAndOperator(sql("SELECT hex(unbase64(bad)) FROM valid")) + checkSparkAnswerAndOperator( + sql("SELECT /*+ BROADCAST(r) */ l.* FROM inner_left l " + + "INNER JOIN valid r ON l.k = r.k AND unbase64(r.bad) > l.expected")) + // AQE can remove this sort after Partial materializes; offset-only collection stays native. + checkSparkAnswerAndOperator( + sql("SELECT unbase64(bad) AS decoded, collect_list(k) FROM one_group " + + "GROUP BY bad ORDER BY decoded OFFSET 1")) + } + } + } + } + } + + test("whole-tree JVM dispatch preserves an unregistered decoder subclass below LIMIT") { + val name = "comet_test_unbase64" + spark.sessionState.functionRegistry.createOrReplaceTempFunction( + name, + children => new LimitUnBase64(children.head), + "scala_udf") + try { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + withInputs("subclass_input" -> "SELECT * FROM VALUES ('YQ=='), ('A') AS t(encoded)") { + val query = s"SELECT format_string('%d', length($name(encoded))) FROM subclass_input" + val native = sql(query).queryExecution.executedPlan.collect { + case p: CometProjectExec => p + } + assert( + native.size == 1 && native.head.nativeOp.getProjection + .getProjectList(0) + .hasJvmScalarUdf) + val df = sql(s"$query LIMIT 1") + val (_, plan) = checkSparkAnswerAndFallbackReason(df, limitReason) + assert(count[ProjectExec](plan) == 1 && count[CometProjectExec](plan) == 0) + assert(nativeScans(plan) == 1) + checkAnswer(df, Seq(Row("1"))) + decodeErrors(s"$query LIMIT 2") + } + } + } finally { + spark.sessionState.catalog.dropTempFunction(name, ignoreIfNotExists = true) + } + } + + test("unbase64 LIMIT masks survive physical boundaries and native reuse") { + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + withModes { (_, _) => + val input = sparkPlan("SELECT * FROM VALUES (1, 'YWJj'), (2, 'A') AS t(id, bad)") + def project = ProjectExec( + Seq(input.output.head, Alias(UnBase64(input.output(1)), "decoded")()), + input) + val staged = applyRule( + GlobalLimitExec(1, ShuffleExchangeExec(SinglePartition, LocalLimitExec(1, project)))) + assert(count[CometGlobalLimitExec](staged) == 1) + assert(count[CometLocalTableScanExec](staged) == 1) + val exchange = collect(staged) { case s: ShuffleExchangeLike => s }.head + for (plan <- Seq(applyRule(staged), applyRule(exchange))) { + assert(count[LocalLimitExec](plan) == 1) + assert(count[ProjectExec](plan) == 1) + assert(count[CometProjectExec](plan) == 0) + } + val barriers: Seq[SparkPlan => SparkPlan] = Seq( + child => SortExec(Seq(SortOrder(child.output.head, Ascending)), global = false, child), + child => ShuffleExchangeExec(SinglePartition, child)) + barriers.foreach { barrier => + val blocking = barrier(project) + // Materialization is safe for both fresh plans and native subtrees reused by AQE. + for (child <- Seq(blocking, applyRule(blocking))) { + val plan = applyRule(LocalLimitExec(1, child)) + withClue(blocking.nodeName) { + assert(count[CometProjectExec](plan) == 1) + assert(count[ProjectExec](plan) == 0) + } + } + } + val bridges: Seq[SparkPlan => SparkPlan] = Seq( + child => child, + child => RowToColumnarExec(ColumnarToRowExec(child)), + child => CometSparkToColumnarExec(CometColumnarToRowExec(child)), + child => CometSparkToColumnarExec(CometNativeColumnarToRowExec(child))) + bridges.foreach { bridge => + val decoded = project + val native = applyRule(ProjectExec(decoded.output, decoded)) + assert(count[CometProjectExec](native) == 2) + val logicalStage = LogicalQueryStage(LocalRelation(native.output), native) + native.setLogicalLink(logicalStage) + val transformed = applyRule(GlobalLimitExec(1, bridge(native))) + for (plan <- Seq(transformed, applyRule(transformed))) { + assert(count[ProjectExec](plan) == 2) + assert(count[CometProjectExec](plan) == 0) + assert(collect(plan) { case r: RowToColumnarTransition => r }.isEmpty) + val link = plan.children.head.getTagValue(SparkPlan.LOGICAL_PLAN_TAG) + assert(link.exists(_ eq logicalStage)) + } + } + } + } + } + + test("unbase64 chooses compatible aggregate buffers before AQE materialization") { + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + withModes { (aqe, _) => + for { + decoderInResult <- Seq(false, true) + boundary <- Seq("direct", "sort", "shuffle", "top-k", "offset") + } { + val result = if (decoderInResult) { + "unbase64(bad) AS decoded, collect_list(k) AS collected" + } else { + "bad, collect_list(unbase64(bad)) AS collected" + } + val aggregate = sparkPlan( + s"SELECT $result FROM " + + "VALUES (1, 'YWJj'), (2, 'YWJj') AS t(k, bad) GROUP BY bad") + assert(count[ObjectHashAggregateExec](aggregate) == 2) + val native = applyRule(aggregate) + assert(count[CometHashAggregateExec](native) == 2) + for (child <- Seq(aggregate, native)) { + val order = Seq(SortOrder(child.output.head, Ascending)) + val plan = boundary match { + case "direct" => CollectLimitExec(1, child) + case "sort" => CollectLimitExec(1, SortExec(order, global = false, child)) + case "shuffle" => CollectLimitExec(1, ShuffleExchangeExec(SinglePartition, child)) + case "top-k" => TakeOrderedAndProjectExec(1, order, child.output, child) + case "offset" => + CollectLimitExec(-1, SortExec(order, global = false, child), offset = 1) + } + // Final does not reevaluate decoder inputs, and offset-only never stops early. + val sparkBuffers = decoderInResult && + (boundary == "direct" || (aqe && boundary != "offset")) + val transformed = applyRule(plan) + withClue(s"decoderInResult=$decoderInResult, boundary=$boundary: ") { + for (replanned <- Seq(transformed, applyRule(transformed))) { + assert(aggregateCounts(replanned) == (if (sparkBuffers) (0, 2) else (2, 0))) + } + } + } + } + } + } + } + + test("unbase64 keeps aggregate buffers compatible through an actual AQE join change") { + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + "spark.sql.adaptive.autoBroadcastJoinThreshold" -> "10485760", + "spark.sql.join.preferSortMergeJoin" -> "true", + "spark.sql.adaptive.maxShuffledHashJoinLocalMapThreshold" -> "0", + CometConf.COMET_FORCE_SHJ.key -> "false") { + // One-file inputs keep the malformed suffix last; two bad groups retain LIMIT. + withInputs( + "aqe_valid" -> """SELECT 1 AS k, base64(cast(cast(id AS STRING) AS BINARY)) AS bad + |FROM range(0, 12, 1, 1)""".stripMargin, + "aqe_bad_last" -> + "SELECT k, CASE WHEN bad='MTE=' THEN 'A' ELSE bad END AS bad FROM aqe_valid", + "aqe_bad_only" -> """SELECT 1 AS k, CASE WHEN id=0 THEN 'A' ELSE 'B' END AS bad + |FROM range(0, 2, 1, 1)""".stripMargin, + "aqe_right" -> "SELECT 'MA==' AS bad") { + withModes { (aqe, _) => + for (join <- Seq("INNER", "LEFT SEMI")) { + withClue(s"join=$join: ") { + def query(input: String, decoder: String = "bad", agg: String = "collect_list(k)") + : String = + s"""SELECT a.decoded, a.collected FROM ( + | SELECT bad, unbase64($decoder) AS decoded, $agg AS collected + | FROM $input GROUP BY bad + |) a $join JOIN aqe_right b ON a.bad = b.bad LIMIT 1""".stripMargin + def checkSuccess(input: String, agg: String = "collect_list(k)"): SparkPlan = { + val text = query(input, agg = agg) + var expected = Seq.empty[Row] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + expected = sql(text).collect().toSeq + } + val df = sql(text) + val initial = df.queryExecution.executedPlan + val initialBuffers = aggregateCounts(initial) + assert( + collect(initial) { + case p if original(p).isInstanceOf[SortMergeJoinExec] => p + }.size == 1, + initial.toString) + checkAnswer(df, expected) + if (aqe && agg == "collect_list(k)") { + assert(initialBuffers == (0, 2), initial.toString) + } + val plan = df.queryExecution.executedPlan + if (aqe) { + assert(plan.asInstanceOf[AdaptiveSparkPlanExec].isFinalPlan) + assert( + collect(plan) { + case j: BroadcastHashJoinExec => j.buildSide + case j: CometBroadcastHashJoinExec => j.buildSide + } == Seq(BuildRight), + plan.toString) + } + plan + } + val valid = checkSuccess("aqe_valid") + assert(aggregateCounts(valid) == (if (aqe) (0, 2) else (2, 0))) + // SEMI streams the left side; INNER may materialize a left broadcast candidate. + if (aqe && join == "LEFT SEMI") checkSuccess("aqe_bad_last") + else if (!aqe) decodeErrors(query("aqe_bad_last")) + decodeErrors(query("aqe_valid", "CASE WHEN bad='MA==' THEN 'A' ELSE bad END")) + // ObjectHash evaluates results even before a nonmatching join discards them. + decodeErrors(query("aqe_bad_only")) + val compatible = checkSuccess("aqe_valid", "max(k)") + assert(count[CometHashAggregateExec](compatible) > 0) + if (aqe && join == "LEFT SEMI") { + // HashAggregate can leave this nonmatching output projection lazy. + checkSuccess("aqe_bad_only", "max(k)") + } + } + } + } + } + } + } + + test("ordered top-K masks child decoders but keeps its output projection native") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_FORCE_APPLY.key -> "true", + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> + ("org.apache.spark.sql.catalyst.optimizer.EliminateSorts," + + "org.apache.spark.sql.catalyst.optimizer.CollapseProject")) { + withModes { (aqe, _) => + for { + location <- Seq("child project", "child filter", "output projection") + comet <- Seq(false, true) + asRDD <- Seq(false, true) + } { + withSQLConf(CometConf.COMET_ENABLED.key -> comet.toString) { + withClue(s"location=$location, Comet=$comet, RDD=$asRDD: ") { + // Each ordered partition must stop before its fourth, malformed row. + def query(limit: Int): DataFrame = { + val input = spark + .range(0, 8, 1, 2) + .selectExpr("id", "if (id % 4 < 3, 'YWJj', 'A') AS encoded") + val ordered = location match { + case "child project" => + input.selectExpr("id", "hex(unbase64(encoded)) AS decoded").orderBy("id") + case "child filter" => + input.filter("unbase64(encoded) = X'616263'").orderBy("id") + case _ => + input.orderBy("id").selectExpr("id", "hex(unbase64(encoded)) AS decoded") + } + ordered.limit(limit) + } + def run(df: DataFrame): Seq[Row] = + (if (asRDD) df.rdd.collect() else df.collect()).toSeq + + val df = query(1) + val inOutput = location == "output projection" + assert( + run(df) == Seq(Row(0L, if (location == "child filter") "YWJj" else "616263"))) + val plan = df.queryExecution.executedPlan + assert(plan.isInstanceOf[AdaptiveSparkPlanExec] == aqe) + val topKs = collect(plan) { + case node if original(node).isInstanceOf[TakeOrderedAndProjectExec] => node + } + assert(topKs.size == 1, plan.toString) + val topK = original(topKs.head).asInstanceOf[TakeOrderedAndProjectExec] + val child = topKs.head.children.head + assert(SortOrder.orderingSatisfies(child.outputOrdering, topK.sortOrder)) + assert(child.outputPartitioning.numPartitions == 2) + assert(collect(child) { + case _: SortExec | _: CometSortExec | _: ShuffleExchangeLike => true + }.isEmpty) + val decoders = collect(child) { + case p: ProjectExec if p.projectList.exists(hasDecoder) => p + case f: FilterExec if hasDecoder(f.condition) => f + } + assert(decoders.size == (if (inOutput) 0 else 1), plan.toString) + assert(topK.projectList.exists(hasDecoder) == inOutput) + assert( + topKs.head.isInstanceOf[CometTakeOrderedAndProjectExec] == (comet && inOutput)) + decodeError(run(query(4))) + } + } + } + } + } + } + + test("first-match joins preserve candidate order and leave decoder-free joins native") { + val strategies = Seq( + ("BROADCAST", "=", classOf[BroadcastHashJoinExec], classOf[CometBroadcastHashJoinExec]), + ("SHUFFLE_HASH", "=", classOf[ShuffledHashJoinExec], classOf[CometHashJoinExec]), + ("MERGE", "=", classOf[SortMergeJoinExec], classOf[CometSortMergeJoinExec]), + ( + "BROADCAST", + "<=", + classOf[BroadcastNestedLoopJoinExec], + classOf[CometBroadcastNestedLoopJoinExec])) + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_EXEC_SORT_MERGE_JOIN_ENABLED.key -> "true", + CometConf.COMET_EXEC_SORT_MERGE_JOIN_WITH_JOIN_FILTER_ENABLED.key -> "true") { + withInputs( + "mask_left" -> "SELECT * FROM VALUES (1, X'616262'), (2, X'616262') AS t(k, expected)", + "mask_forward" -> "SELECT * FROM VALUES (1, 'YWJj'), (1, 'A') AS t(k, bad)", + "mask_reverse" -> "SELECT * FROM VALUES (1, 'A'), (1, 'YWJj') AS t(k, bad)") { + withModes { (aqe, _) => + for ((hint, equality, sparkJoin, nativeJoin) <- strategies; + kind <- Seq("SEMI", "ANTI")) { + withClue(s"$hint, predicate=$equality, LEFT $kind: ") { + // Hash joins visit duplicate keys backwards; merge and nested-loop visit forwards. + val table = + if (hint == "MERGE" || equality == "<=") "mask_forward" else "mask_reverse" + val reversed = if (table == "mask_forward") "mask_reverse" else "mask_forward" + def query(residual: String, input: String = table): String = + s"""SELECT /*+ $hint(r) */ l.k FROM mask_left l LEFT $kind JOIN $input r + |ON l.k $equality r.k $residual""".stripMargin + val (sparkPlan, cometPlan) = checkSparkAnswerAndFallbackReason( + query("AND unbase64(r.bad) > l.expected"), + joinReason) + Seq(sparkPlan, cometPlan).foreach { plan => + assert(plan.isInstanceOf[AdaptiveSparkPlanExec] == aqe) + if (aqe) assert(plan.asInstanceOf[AdaptiveSparkPlanExec].isFinalPlan) + val joins = collect(plan) { case node if sparkJoin.isInstance(node) => node } + assert( + joins.size == 1 && joins.head.expressions.exists(hasDecoder), + plan.toString) + assert(count[CometBroadcastExchangeExec](plan) == 0, plan.toString) + assert( + count[BroadcastExchangeLike](plan) == + (if (hint == "BROADCAST") 1 else 0), + plan.toString) + } + assert(collect(cometPlan) { case _: CometScanExec | _: CometNativeScanExec => + true + }.size == 2) + decodeErrors(query("AND unbase64(r.bad) < l.expected")) + decodeErrors(query("AND unbase64(r.bad) > l.expected", reversed)) + val (_, nativePlan) = checkSparkAnswerAndOperator(query("")) + assert(collect(nativePlan) { case p if nativeJoin.isInstance(p) => p }.size == 1) + } + } + } + } + } + } + + test("window group limits drain partitions unless an outer LIMIT stops consumption") { + assume(isSpark35Plus, "WindowGroupLimit requires Spark 3.5+") + withSQLConf( + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> + ("org.apache.spark.sql.catalyst.optimizer.EliminateSorts," + + "org.apache.spark.sql.catalyst.optimizer.CollapseProject"), + "spark.sql.optimizer.windowGroupLimitThreshold" -> "1000", + "spark.sql.execution.topKSortFallbackThreshold" -> "0", + CometConf.COMET_EXEC_WINDOW_GROUP_LIMIT_ENABLED.key -> "true") { + val decoder = "unbase64(IF(id = 0, 'YWJj', 'A'))" + def unpartitioned(project: Boolean, limit: Int = 1, order: String = "id"): String = { + val decoded = + if (project) + s", first_value($decoder) OVER (ORDER BY $order " + + "ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS decoded" + else "" + val filter = if (project) "" else s"WHERE hex($decoder) = '616263'" + s"""SELECT * FROM ( + |SELECT id, row_number() OVER (ORDER BY $order) AS rn $decoded + |FROM range(0, 2, 1, 1) $filter) WHERE rn <= $limit""".stripMargin + } + def partitioned(malformed: Boolean, limit: String = ""): String = { + val encoded = + if (malformed) "IF(id = 15, 'A', 'YWJj')" + else "IF(id % 2 = 0, 'YWJj', 'YWFh')" + // Sort below the decoder; put the bad row beyond WindowExec's group lookahead. + s"""SELECT k, id, hex(decoded) FROM ( + |SELECT k, id, decoded, row_number() OVER (PARTITION BY k ORDER BY id) AS rn + |FROM (SELECT k, id, unbase64(bad) AS decoded FROM ( + |SELECT id, id DIV 4 AS k, $encoded AS bad FROM range(0, 16, 1, 1) + |ORDER BY k, id))) WHERE rn <= 1 $limit""".stripMargin + } + def checkInput(plan: SparkPlan, partitioned: Boolean, native: Boolean): Unit = { + val windows = collect(plan) { + case node if ShimCometWindowGroupLimit.extract(original(node)).isDefined => node + } + assert(windows.size == 1, plan.toString) + val window = windows.head + val fields = ShimCometWindowGroupLimit.extract(original(window)).get + assert(fields.partitionSpec.nonEmpty == partitioned) + assert(window.isInstanceOf[CometExec] == native, plan.toString) + val child = window.children.head match { + case w: WholeStageCodegenExec => w.child + case other => other + } + val required = original(window).requiredChildOrdering.head + assert(SortOrder.orderingSatisfies(child.outputOrdering, required)) + assert(child.isInstanceOf[CometExec] == native, plan.toString) + assert(original(child).expressions.exists(hasDecoder), plan.toString) + } + withModes { (_, dispatch) => + for (project <- Seq(false, true)) { + val (sparkPlan, cometPlan) = + checkSparkAnswerAndFallbackReason(unpartitioned(project), limitReason) + Seq(sparkPlan, cometPlan).foreach(checkInput(_, partitioned = false, native = false)) + } + decodeErrors(unpartitioned(project = false, limit = 2)) + val sorted = unpartitioned(project = false, order = "id DESC") + if (dispatch) { + withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false") { + val unordered = collect(sparkPlan(unpartitioned(project = false))) { + case w if ShimCometWindowGroupLimit.extract(w).isDefined => w + }.head.transformExpressions { case order: SortOrder => + order.copy(direction = Descending) + } + assert( + !SortOrder.orderingSatisfies( + unordered.children.head.outputOrdering, + unordered.requiredChildOrdering.head)) + assert(collect(applyRule(unordered)) { + case f: CometFilterExec if hasDecoder(f.condition) => f + }.size == 1) + } + } + decodeErrors(sorted) + val (sparkValid, cometValid) = checkSparkAnswerAndOperator(partitioned(malformed = false)) + checkInput(sparkValid, partitioned = true, native = false) + checkInput(cometValid, partitioned = true, native = true) + decodeErrors(partitioned(malformed = true)) + val (sparkLimited, cometLimited) = checkSparkAnswerAndFallbackReason( + partitioned(malformed = true, limit = "LIMIT 1"), + limitReason) + Seq(sparkLimited, cometLimited).foreach(checkInput(_, partitioned = true, native = false)) + checkInput( + applyRule(LocalLimitExec(1, stripAQEPlan(cometValid))), + partitioned = true, + native = false) + } + } + } +} + +// Preserve the subclass through binding without capturing the suite in the dispatched expression. +private[serde] class LimitUnBase64(input: Expression, strict: Boolean = false) + extends UnBase64(input, strict) { + override protected def withNewChildInternal(newChild: Expression): UnBase64 = + new LimitUnBase64(newChild, failOnError) +} From 6da8a5cbf14da5377cdeed129ed4d0143ee1a62c Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 4 Oct 2026 17:29:56 +0000 Subject: [PATCH 2/9] fix: retain pruning-only scans and stabilize AQE mask tests --- .../apache/comet/rules/CometExecRule.scala | 13 ++++++++-- .../serde/CometEvaluationMaskSuite.scala | 26 ++++++++++++++----- 2 files changed, 30 insertions(+), 9 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 206147a24db..86e874f555e 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -1117,11 +1117,19 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) case _ => node } + def evaluatedExpressions(node: SparkPlan): Seq[Expression] = node match { + // V1 dataFilters include predicates Spark cannot push into Parquet. Without native row + // filtering they only inform metadata pruning; the Spark Filter above evaluates rows. + case _: FileSourceScanExec if !CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.get() => + Seq.empty + case _ => node.expressions + } + // Most plans contain no opted-in expression. Include native originals and sticky tags // so repeated AQE passes still restore an already-protected subtree when necessary. if (!plan.exists(node => node.getTagValue(CometExecRule.UNSAFE_EXPRESSION_EVALUATION).isDefined || - originalPlan(node).expressions.exists(findEvaluationMaskName(_).isDefined))) { + evaluatedExpressions(originalPlan(node)).exists(findEvaluationMaskName(_).isDefined))) { return plan } @@ -1184,7 +1192,8 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } val limitName = if (belowLimit) { // Final merges buffers; it does not reevaluate the aggregate's original inputs. - val expressions = finalAggregate.map(_.resultExpressions).getOrElse(original.expressions) + val expressions = + finalAggregate.map(_.resultExpressions).getOrElse(evaluatedExpressions(original)) expressions.iterator.flatMap(findEvaluationMaskName).take(1).toSeq.headOption } else { None diff --git a/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala b/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala index e206f0e2f07..fbf64a75422 100644 --- a/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala +++ b/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala @@ -192,7 +192,14 @@ class CometEvaluationMaskSuite extends CometTestBase { "SELECT bad FROM masked WHERE unbase64(bad) <=> X'616263' LIMIT 1", "SELECT hex(unbase64(concat(bad, ''))) FROM masked LIMIT 1", "SELECT hex(to_binary(bad, 'base64')) FROM masked LIMIT 1").foreach { query => - checkSparkAnswerAndFallbackReason(query, limitReason) + val (_, plan) = checkSparkAnswerAndFallbackReason(query, limitReason) + assert(nativeScans(plan) == 1, plan.toString) + } + withSQLConf(CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.key -> "true") { + val (_, plan) = checkSparkAnswerAndFallbackReason( + "SELECT bad FROM masked WHERE unbase64(bad) <=> X'616263' LIMIT 1", + limitReason) + assert(nativeScans(plan) == 0, plan.toString) } decodeErrors("SELECT hex(unbase64(bad)) FROM masked") decodeErrors("SELECT hex(unbase64(bad)) FROM masked WHERE bad = 'A' LIMIT 1") @@ -389,15 +396,20 @@ class CometEvaluationMaskSuite extends CometTestBase { if (aqe && agg == "collect_list(k)") { assert(initialBuffers == (0, 2), initial.toString) } + if (aqe && agg == "max(k)") { + assert(initialBuffers == (2, 0), initial.toString) + } val plan = df.queryExecution.executedPlan if (aqe) { assert(plan.asInstanceOf[AdaptiveSparkPlanExec].isFinalPlan) - assert( - collect(plan) { - case j: BroadcastHashJoinExec => j.buildSide - case j: CometBroadcastHashJoinExec => j.buildSide - } == Seq(BuildRight), - plan.toString) + val broadcastBuildSides = collect(plan) { + case j: BroadcastHashJoinExec => j.buildSide + case j: CometBroadcastHashJoinExec => j.buildSide + } + assert(broadcastBuildSides.size == 1, plan.toString) + if (join == "LEFT SEMI") { + assert(broadcastBuildSides == Seq(BuildRight), plan.toString) + } } plan } From 1211d7932af11c3d1d9454632003770f22f5158a Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 4 Oct 2026 17:55:11 +0000 Subject: [PATCH 3/9] test: construct aggregate count tuples explicitly for Scala 2.12 --- .../org/apache/comet/serde/CometEvaluationMaskSuite.scala | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala b/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala index fbf64a75422..c5b743504b2 100644 --- a/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala +++ b/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala @@ -144,7 +144,7 @@ class CometEvaluationMaskSuite extends CometTestBase { val aggregate = sparkPlan( "SELECT unbase64(bad) AS decoded, collect_list(k) FROM " + "VALUES (1, 'YWJj'), (1, 'YWFh') AS t(k, bad) GROUP BY bad") - assert(aggregateCounts(aggregate) == (0, 2)) + assert(aggregateCounts(aggregate) == Tuple2(0, 2)) val sorted = SortExec( Seq(SortOrder(aggregate.output.head, Ascending)), global = false, @@ -394,10 +394,10 @@ class CometEvaluationMaskSuite extends CometTestBase { initial.toString) checkAnswer(df, expected) if (aqe && agg == "collect_list(k)") { - assert(initialBuffers == (0, 2), initial.toString) + assert(initialBuffers == Tuple2(0, 2), initial.toString) } if (aqe && agg == "max(k)") { - assert(initialBuffers == (2, 0), initial.toString) + assert(initialBuffers == Tuple2(2, 0), initial.toString) } val plan = df.queryExecution.executedPlan if (aqe) { From 28e8ca1176d95a9609577a7848ff8e88b7e1f80c Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 4 Oct 2026 19:34:54 +0000 Subject: [PATCH 4/9] fix: preserve fused top-K evaluation across AQE passes --- .../apache/comet/rules/CometExecRule.scala | 14 ++++-- .../serde/CometEvaluationMaskSuite.scala | 50 +++++++++++++++++++ 2 files changed, 61 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 86e874f555e..caf278be68b 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -1111,6 +1111,9 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) case scan: CometScanExec => scan.wrapped .copy(partitionFilters = scan.partitionFilters, dataFilters = scan.dataFilters) + // The local selector saves its parent's original plan, but does not evaluate that + // parent's projection or offset. Only the enclosing top-K owns those operations. + case local: CometLocalTopKExec => local case comet: CometExec => comet.originalPlan case shuffle: CometShuffleExchangeExec => shuffle.originalPlan case broadcast: CometBroadcastExchangeExec => broadcast.originalPlan @@ -1118,6 +1121,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } def evaluatedExpressions(node: SparkPlan): Seq[Expression] = node match { + case local: CometLocalTopKExec => local.sortOrder // V1 dataFilters include predicates Spark cannot push into Parquet. Without native row // filtering they only inform metadata pruning; the Spark Filter above evaluates rows. case _: FileSourceScanExec if !CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.get() => @@ -1163,9 +1167,9 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) // These operators consume their input before yielding rows. Still visit their children: // an inner LocalLimit below an exchange must establish its own evaluation boundary. val materializesInput = original match { - case _: SortExec | _: HashAggregateExec | _: ObjectHashAggregateExec | - _: ShuffleExchangeLike | _: BroadcastExchangeLike | _: QueryStageExec | - _: ReusedExchangeExec => + case _: SortExec | _: CometLocalTopKExec | _: HashAggregateExec | + _: ObjectHashAggregateExec | _: ShuffleExchangeLike | _: BroadcastExchangeLike | + _: QueryStageExec | _: ReusedExchangeExec => true case _ => false } @@ -1232,6 +1236,10 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) case _: ColumnarToRowExec | _: CometColumnarToRowExec | _: CometNativeColumnarToRowExec if childReason.isDefined && !children.head.supportsColumnar => children.head + // This inserted selector has no standalone Spark original. Its parent restores + // the full top-K when a child changes, applying the projection and offset once. + case _: CometLocalTopKExec if reason.isDefined || children != node.children => + children.head case _ if (original ne node) && (reason.isDefined || children != node.children) => // AQE can reuse an existing native subtree. Rebuild affected ancestors as well so // their serialized native plans do not retain the decoder that just fell back. diff --git a/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala b/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala index c5b743504b2..0d2ae170df1 100644 --- a/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala +++ b/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala @@ -499,6 +499,56 @@ class CometEvaluationMaskSuite extends CometTestBase { } } + test("fused top-K evaluates its final decoder and offset once across AQE passes") { + withInputs( + "fused_mask" -> + "SELECT * FROM VALUES (1, 'YWJj'), (2, 'YWJj'), (3, 'YWJj'), (4, 'YWJj') AS t(k, bad)") { + val query = "SELECT k, hex(unbase64(bad)) AS decoded FROM " + + "(SELECT * FROM fused_mask ORDER BY k LIMIT 2 OFFSET 1) t" + val expected = Seq(Row(2, "616263"), Row(3, "616263")) + for (aqe <- Seq(true, false); fusion <- Seq(false, true)) { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString, + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_EXEC_TOPK_FUSION_ENABLED.key -> fusion.toString) { + val joined = sql(s"""SELECT /*+ BROADCAST(t) */ t.k, t.decoded + |FROM ($query) t JOIN range(1, 5) r ON t.k = r.id""".stripMargin) + checkAnswer(joined, expected) + val executed = joined.queryExecution.executedPlan + assert(count[CometLocalTopKExec](executed) == (if (fusion) 1 else 0), executed.toString) + assert(count[CometTakeOrderedAndProjectExec](executed) == 1, executed.toString) + assert(count[TakeOrderedAndProjectExec](executed) == 0, executed.toString) + assert(count[BroadcastExchangeLike](executed) == 1, executed.toString) + if (aqe) { + assert(executed.asInstanceOf[AdaptiveSparkPlanExec].isFinalPlan) + } + + // Reapplying the rule must not restore the local candidate selector as another + // full TakeOrderedAndProject: its saved original also contains the outer projection. + val once = collect(executed) { case topK: CometTakeOrderedAndProjectExec => topK }.head + val twice = applyRule(once) + assert(count[CometLocalTopKExec](twice) == (if (fusion) 1 else 0), twice.toString) + assert(count[CometTakeOrderedAndProjectExec](twice) == 1, twice.toString) + assert(count[TakeOrderedAndProjectExec](twice) == 0, twice.toString) + val replannedRows = twice.executeCollect().toSeq.map { row => + Row(row.getInt(0), row.getUTF8String(1).toString) + } + assert(replannedRows == expected) + + // An outer row limit really does require the final decoder to fall back. The + // candidate selector may stay native, but must never duplicate the final projection. + val limited = applyRule(LocalLimitExec(1, twice)) + assert(count[TakeOrderedAndProjectExec](limited) == 1, limited.toString) + assert(count[CometTakeOrderedAndProjectExec](limited) == 0, limited.toString) + val limitedRows = limited.executeCollect().toSeq.map { row => + Row(row.getInt(0), row.getUTF8String(1).toString) + } + assert(limitedRows == expected.take(1)) + } + } + } + } + test("first-match joins preserve candidate order and leave decoder-free joins native") { val strategies = Seq( ("BROADCAST", "=", classOf[BroadcastHashJoinExec], classOf[CometBroadcastHashJoinExec]), From f20168950e3ea7e71034fdaece92fdbc67735d34 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 15 Sep 2026 21:00:12 +0000 Subject: [PATCH 5/9] fix: preserve Spark evaluation for next_day and levenshtein --- .../apache/comet/rules/CometExecRule.scala | 74 +++++- .../apache/comet/serde/QueryPlanSerde.scala | 144 ++++++++++++ .../scala/org/apache/comet/serde/arrays.scala | 4 +- .../org/apache/comet/serde/datetime.scala | 21 +- .../datetime/next_day_ansi_collation.sql | 22 ++ .../next_day_collation_interpreted.sql | 65 ++++++ .../datetime/next_day_deferred_projection.sql | 113 +++++++++ .../levenshtein_collation_interpreted.sql | 40 ++++ .../levenshtein_interpreted_arguments.sql | 49 ++++ .../comet/rules/CometExecRuleSuite.scala | 221 +++++++++++++++++- 10 files changed, 745 insertions(+), 8 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/datetime/next_day_collation_interpreted.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/datetime/next_day_deferred_projection.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/string/levenshtein_collation_interpreted.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/string/levenshtein_interpreted_arguments.sql 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 caf278be68b..435eb6d0ce9 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -25,7 +25,8 @@ import scala.collection.mutable.ListBuffer import scala.jdk.CollectionConverters._ import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder, SortOrder} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LeafExpression, LessThan, LessThanOrEqual, NamedExpression, Remainder, SortOrder} +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero import org.apache.spark.sql.catalyst.plans.{JoinType, LeftAnti, LeftSemi} @@ -1142,10 +1143,32 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) case _ => false } + // Match Spark's whole-stage eligibility before relying on deferred generated expressions. + // In interpreted execution, or across an input adapter, a Project materializes every output. + def supportsWholeStage(node: SparkPlan): Boolean = node match { + case codegen: CodegenSupport + if conf.wholeStageEnabled && + !QueryPlanSerde.usesInterpretedProjection(conf) && + codegen.supportCodegen => + !node.expressions.exists(_.exists { + case _: LeafExpression => false + case _: CodegenFallback => true + case _ => false + }) && !WholeStageCodegenExec.isTooManyFields(conf, node.schema) && + !node.children.exists(child => WholeStageCodegenExec.isTooManyFields(conf, child.schema)) + case _ => false + } + + def eagerReferences(expr: Expression): AttributeSet = expr match { + case attribute: AttributeReference => AttributeSet(Seq(attribute)) + case _ => AttributeSet(QueryPlanSerde.eagerlyEvaluatedChildren(expr).flatMap(eagerReferences)) + } + def protect( node: SparkPlan, belowLimit: Boolean, - hasLimitAncestor: Boolean): (SparkPlan, Option[String]) = { + hasLimitAncestor: Boolean, + deferredInputs: AttributeSet): (SparkPlan, Option[String]) = { val original = originalPlan(node) val startsLimit = original match { // Offset-only collection does not stop its input early. @@ -1176,8 +1199,32 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) // A top-K is still a LIMIT even when its current input requires sorting. val limitAncestor = hasLimitAncestor || startsLimit || original.isInstanceOf[TakeOrderedAndProjectExec] + val childDeferredInputs = original match { + // AQE can present a previously converted subtree. Batch bridges must not hide + // the original Project from its Spark consumer's deferred-input contract. + case _: RowToColumnarExec | _: ColumnarToRowExec | _: CometColumnarToRowExec | + _: CometNativeColumnarToRowExec | _: CometSparkToColumnarExec => + deferredInputs + // Spark gives both children of these joins separate generated stages. Stage roots + // and input adapters likewise materialize rows before handing them to a consumer. + case _: SortMergeJoinExec | _: ShuffledHashJoinExec | _: WholeStageCodegenExec | + _: InputAdapter => + AttributeSet.empty + case project: ProjectExec if supportsWholeStage(project) => + val eagerOutputs = project.projectList.filterNot(expr => + expr.deterministic && deferredInputs.contains(expr.toAttribute)) + project.inputSet -- project.usedInputs -- + AttributeSet(eagerOutputs.flatMap(eagerReferences)) + case codegen: CodegenSupport if !materializesInput && supportsWholeStage(original) => + original.inputSet -- codegen.usedInputs + case _ => AttributeSet.empty + } val protectedChildren = node.children.map { child => - protect(child, startsLimit || (belowLimit && !materializesInput), limitAncestor) + protect( + child, + startsLimit || (belowLimit && !materializesInput), + limitAncestor, + childDeferredInputs) } val childReason = protectedChildren.flatMap(_._2).headOption val condition = original match { @@ -1213,10 +1260,29 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } else { None } + val deferredOutputExpressions = original match { + case project: ProjectExec => project.projectList + // Grouped hash aggregation materializes its input but defers deterministic result + // expressions to its consumer. Ungrouped aggregation evaluates its results eagerly. + case aggregate: HashAggregateExec if aggregate.groupingExpressions.nonEmpty => + aggregate.resultExpressions + case _ => Seq.empty + } + val deferredProjection = + if (deferredOutputExpressions.nonEmpty && supportsWholeStage(original)) { + deferredOutputExpressions + .filter(expr => expr.deterministic && deferredInputs.contains(expr.toAttribute)) + .flatMap(findEvaluationMaskName) + .headOption + .map(name => s"$name requires Spark evaluation in a deferred projection") + } else { + None + } val ownReason = limitName .map(name => s"$name requires Spark evaluation below LIMIT") .orElse(aggregateBufferName.map(name => s"$name requires Spark aggregate buffers below LIMIT with AQE")) + .orElse(deferredProjection) .orElse( condition .flatMap(findEvaluationMaskName) @@ -1252,7 +1318,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) (prepared, reason) } - protect(plan, belowLimit = false, hasLimitAncestor = false)._1 + protect(plan, belowLimit = false, hasLimitAncestor = false, AttributeSet.empty)._1 } /** Convert a Spark plan to a Comet plan using the specified serde handler */ 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 38ed0355b7c..af7921c9619 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -24,6 +24,7 @@ import java.util.concurrent.atomic.AtomicLong import scala.collection.mutable import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.expressions._ @@ -836,6 +837,34 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { } val fn = aggExpr.aggregateFunction + if (aggExpr.mode == Partial || aggExpr.mode == Complete) { + // Raw-input aggregates can skip arguments (FILTER, null-skipping, or state-dependent + // updates). Final/PartialMerge only read buffers, so their original children are irrelevant. + var pending = fn.children.toList + while (pending.nonEmpty) { + val child = pending.head + pending = pending.tail + child match { + case nextDay: NextDay + if CometNextDay.evaluationMaskName(nextDay).isDefined && + (aggExpr.filter.isDefined || fn.isInstanceOf[First] || fn.isInstanceOf[Last]) => + withFallbackReason( + aggExpr, + "next_day requires Spark evaluation in aggregate arguments: " + + "native aggregation may evaluate a masked child.") + return None + case _ => + } + // Imperative aggregates interpret their arguments even when codegen is enabled. + if (fn.isInstanceOf[ImperativeAggregate] && requiresSparkInterpretedEvaluation(child)) { + withFallbackReason( + aggExpr, + s"${child.prettyName} requires Spark interpreted evaluation in aggregate arguments") + return None + } + pending = child.children.toList ::: pending + } + } val cometExpr = aggrSerdeMap.get(fn.getClass) val protoAggExprOpt = cometExpr match { case Some(handler) => @@ -966,6 +995,24 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { result } + /** Known eval/codegen differences; callers inspect the enclosing expression tree. */ + private[serde] def requiresSparkInterpretedEvaluation(expr: Expression): Boolean = expr match { + // Interpreted evaluation unboxes a null threshold to zero; generated evaluation returns null. + case levenshtein: Levenshtein => levenshtein.children.lift(2).exists(_.nullable) + case nextDay: NextDay + if nextDay.dayOfWeek.foldable && !nextDay.dayOfWeek.isInstanceOf[Literal] => + // Match NextDay.doGenCode's constant evaluation. Spark can recover from a generation failure + // by interpreting the projection, but the dispatcher cannot. Successful constants stay + // eligible for native execution or dispatch. + try { + nextDay.dayOfWeek.eval() + false + } catch { + case NonFatal(_) => true + } + case _ => false + } + private def liftCoverageTags(from: Expression, to: Expression): Unit = { val exprs = from.collect { case e: Expression => e } Seq(CometExplainInfo.NATIVE_EXPRS, CometExplainInfo.CODEGEN_DISPATCH_EXPRS).foreach { tag => @@ -1023,6 +1070,34 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { + // Inspect the full tree before an enclosing expression can dispatch it as a whole. + var remaining = expr :: Nil + while (remaining.nonEmpty) { + val node = remaining.head + remaining = remaining.tail + node match { + case nextDay: NextDay if requiresSparkInterpretedEvaluation(nextDay) => + withFallbackReason( + expr, + "next_day requires Spark evaluation because its weekday fails during code generation") + return None + case _: Levenshtein if requiresSparkInterpretedEvaluation(node) => + withFallbackReason( + expr, + "levenshtein with a nullable threshold requires Spark evaluation: " + + "Spark interpreted and generated evaluation differ") + return None + case _ => + } + node.children.reverseIterator.foreach(child => remaining = child :: remaining) + } + unsupportedNextDayEvaluation(expr) match { + case Some(reason) => + withFallbackReason(expr, reason) + return None + case None => + } + def convert[T <: Expression](expr: T, handler: CometExpressionSerde[T]): Option[Expr] = { val exprConfName = handler.getExprConfigName(expr) if (!CometConf.isExprEnabled(exprConfName)) { @@ -1126,6 +1201,75 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { } } + /** Spark 4.2 changed this setting from String to an Enumeration.Value. */ + private[comet] def usesInterpretedProjection(conf: SQLConf): Boolean = + conf.getConf(SQLConf.CODEGEN_FACTORY_MODE).toString == "NO_CODEGEN" + + /** Inputs Spark evaluates whenever it consumes this expression's output. */ + private[comet] def eagerlyEvaluatedChildren(expr: Expression): Seq[Expression] = expr match { + case conditional: ConditionalExpression => conditional.alwaysEvaluatedInputs + // Generated TryEval can catch failures from a deferred child's code. + case _: TryEval => Seq.empty + case order: SortOrder => Seq(order.child) + case unary: UnaryExpression => Seq(unary.child) + case _: DateAdd | _: DateDiff | _: BinaryComparison => + expr.children.take(if (expr.children.head.nullable) 1 else 2) + case _: Concat | _: Greatest | _: Least => expr.children + case _ => Seq.empty + } + + /** Reject ANSI NextDay below native parents that can evaluate children Spark skips. */ + private def unsupportedNextDayEvaluation(expr: Expression): Option[String] = { + var stack: List[(Expression, Option[String])] = (expr, None) :: Nil + while (stack.nonEmpty) { + val (current, unsafeAncestor) = stack.head + stack = stack.tail + current match { + case nextDay: NextDay + if CometNextDay.evaluationMaskName(nextDay).isDefined && unsafeAncestor.isDefined => + return Some( + s"next_day requires Spark evaluation under ${unsafeAncestor.get}: " + + "native evaluation may evaluate a masked child.") + case _ => + } + + def pushChildren(context: Option[String]): Unit = { + current.children.reverseIterator.foreach { child => + stack = (child, context) :: stack + } + } + + val wholeExpressionDispatch = exprSerdeMap.get(current.getClass).exists { handler => + (handler.isInstanceOf[CodegenDispatchFallback] && + handler.asInstanceOf[CometExpressionSerde[Expression]] + .getSupportLevel(current).isInstanceOf[Unsupported]) || + (handler.isInstanceOf[CometCodegenDispatch[_]] && + !handler.isInstanceOf[NativeOptInAvailable]) + } + if (unsafeAncestor.isDefined) { + // An outer native parent can evaluate even a dispatched subtree too early. + pushChildren(unsafeAncestor) + } else if (!wholeExpressionDispatch) { + current match { + case _: If | _: CaseWhen | _: EqualNullSafe => + pushChildren(None) + case coalesce: Coalesce if coalesce.deterministic => + pushChildren(None) + case order: SortOrder => + // sameOrderExpressions are equivalences, not evaluated sort keys. + stack = (order.child, None) :: stack + case _ => + val eager = eagerlyEvaluatedChildren(current) + current.children.reverseIterator.foreach { child => + val context = if (eager.exists(_ eq child)) None else Some(current.nodeName) + stack = (child, context) :: stack + } + } + } + } + None + } + /** * Nodes that carry no computation of their own. They are excluded from the expression coverage * stats in extended explain because they appear in nearly every expression tree and would swamp diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index b2a75f84782..f5e0eef09f3 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -832,10 +832,12 @@ object CometArrayFilter extends CometExpressionSerde[ArrayFilter] { binding: Boolean): Option[ExprOuterClass.Expr] = { expr.function match { case LambdaFunction(IsNotNull(v: NamedLambdaVariable), Seq(lambdaVar), _) - if v.exprId == lambdaVar.exprId => + if v.exprId == lambdaVar.exprId && + !expr.argument.exists(requiresSparkInterpretedEvaluation) => // Fast path: Catalyst desugars `array_compact` to `filter(arr, x -> x IS NOT NULL)`, so // restore the native serde here (avoids per-batch JNI). Guard requires the IsNotNull // operand to be the lambda variable itself, not a captured column (#4830). + // ArrayFilter interprets its argument, so keep sensitive children in the full dispatcher. CometArrayCompact.convert(expr, inputs, binding) case _ => // General lambda: run Spark's own evaluation through the codegen dispatcher so the result diff --git a/spark/src/main/scala/org/apache/comet/serde/datetime.scala b/spark/src/main/scala/org/apache/comet/serde/datetime.scala index bf27ce2c552..6de0aa96e22 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -404,7 +404,26 @@ object CometConvertTimezone * 70ms dispatched against 89ms for Spark. See * https://github.com/apache/datafusion-comet/issues/5591. */ -object CometNextDay extends CometExpressionSerde[NextDay] with CodegenDispatchFallback { +object CometNextDay + extends CometExpressionSerde[NextDay] + with CodegenDispatchFallback + with RequiresSparkEvaluationMask[NextDay] { + + override def evaluationMaskName(expr: NextDay): Option[String] = { + if (!expr.failOnError) return None + expr.dayOfWeek match { + case Literal(null, _) => None + case Literal(value: org.apache.spark.unsafe.types.UTF8String, _) => + try { + org.apache.spark.sql.catalyst.util.DateTimeUtils.getDayOfWeekFromString(value) + None + } catch { + case scala.util.control.NonFatal(_) => Some("next_day") + } + case _ => Some("next_day") + } + } + private val collationReason = DatetimeCollation.reason("next_day") diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/next_day_ansi_collation.sql b/spark/src/test/resources/sql-tests/expressions/datetime/next_day_ansi_collation.sql index 7e9b93d400d..96bc19a7b61 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/next_day_ansi_collation.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/next_day_ansi_collation.sql @@ -17,6 +17,8 @@ -- MinSparkVersion: 4.0 -- Config: spark.sql.ansi.enabled=true +-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true +-- Config: spark.comet.batchSize=1024 -- ANSI mode with a collated dayOfWeek. A collated next_day runs through the JVM codegen -- dispatcher (see next_day_collation.sql for why), which evaluates Spark's own generated code, so @@ -39,3 +41,23 @@ SELECT next_day(date('2024-01-01'), 'NOT_A_DAY' COLLATE UNICODE_CI) -- ANSI rather than resolving to MONDAY. query expect_error(Illegal input for day of week) SELECT next_day(date('2024-01-01'), 'MON ' COLLATE UTF8_LCASE_RTRIM) + +-- One file keeps the valid first row and invalid second row in the same batch. +statement +CREATE TABLE test_next_day_ansi_masks USING parquet AS +SELECT /*+ COALESCE(1) */ expected, d, dow FROM VALUES + (date('2024-01-08'), date('2024-01-01'), 'Monday'), + (CAST(NULL AS DATE), date('2024-01-01'), 'NOT_A_DAY') AS t(expected, d, dow) + +-- Spark skips next_day when the comparison's left operand is NULL. +query expect_fallback(next_day requires Spark evaluation under EqualTo) +SELECT expected = next_day(d, dow COLLATE UTF8_LCASE) FROM test_next_day_ansi_masks + +-- LIMIT consumes only the valid row; the rest of its batch must remain unevaluated. +query expect_fallback(next_day requires Spark evaluation below LIMIT) +SELECT next_day(d, dow COLLATE UTF8_LCASE) FROM test_next_day_ansi_masks LIMIT 1 + +-- An invalid row that is actually consumed must still throw. +query expect_error(Illegal input for day of week) +SELECT next_day(d, dow COLLATE UTF8_LCASE) FROM test_next_day_ansi_masks +WHERE dow = 'NOT_A_DAY' LIMIT 1 diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/next_day_collation_interpreted.sql b/spark/src/test/resources/sql-tests/expressions/datetime/next_day_collation_interpreted.sql new file mode 100644 index 00000000000..174698d50a7 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/next_day_collation_interpreted.sql @@ -0,0 +1,65 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- MinSparkVersion: 4.0 +-- Config: spark.sql.ansi.enabled=true +-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true +-- ConfigMatrix: spark.sql.codegen.factoryMode=NO_CODEGEN,FALLBACK +-- Config: spark.sql.codegen.wholeStage=false + +-- The harness excludes ConstantFolding. Restore it so these cases exercise the default +-- optimizer: a failing constant is retained inside IF rather than evaluated during planning. +statement +SET spark.sql.optimizer.excludedRules= + +statement +CREATE TABLE test_next_day_collated_interpreted USING parquet AS +SELECT CAST(NULL AS DATE) AS null_d, date('2024-01-01') AS valid_d, + true AS flag, 'Monday' AS dow, 'NOT_A_DAY' AS bad_dow + +-- Safe weekday forms still dispatch. ConstantFolding reduces the collated literal to a Literal, +-- while the column weekday remains nonfoldable. +query +SELECT next_day(valid_d, dow COLLATE UTF8_LCASE), + next_day(valid_d, 'Monday' COLLATE UTF8_LCASE) FROM test_next_day_collated_interpreted + +-- Interpreted NextDay skips its weekday when the date is NULL. Generated evaluation fails while +-- compiling this foldable weekday. FALLBACK must let Spark recover with an interpreted projection. +query expect_fallback(weekday fails during code generation) +SELECT IF(flag, + next_day(null_d, CAST(1 / 0 AS STRING) COLLATE UTF8_LCASE), date('2024-01-01')) +FROM test_next_day_collated_interpreted + +-- AddMonths dispatches its whole subtree, so guarding only the NextDay serde is insufficient. +query expect_fallback(weekday fails during code generation) +SELECT add_months(IF(flag, + next_day(null_d, CAST(1 / 0 AS STRING) COLLATE UTF8_LCASE), date('2024-01-01')), 0) +FROM test_next_day_collated_interpreted + +-- A non-null date consumes the invalid weekday expression and must still raise its error. +query expect_error(DIVIDE_BY_ZERO) +SELECT IF(flag, + next_day(valid_d, CAST(1 / 0 AS STRING) COLLATE UTF8_LCASE), date('2024-01-01')) +FROM test_next_day_collated_interpreted + +-- Without whole-stage fusion, Spark eagerly evaluates the lower Project before the Filter. +-- Returning an empty result here would suppress an error Spark raises in this mode. +query expect_error(Illegal input for day of week) +SELECT x FROM ( + SELECT next_day(valid_d, bad_dow COLLATE UTF8_LCASE) AS x, rand(123L) AS r + FROM test_next_day_collated_interpreted) p +WHERE r < 0.0 diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/next_day_deferred_projection.sql b/spark/src/test/resources/sql-tests/expressions/datetime/next_day_deferred_projection.sql new file mode 100644 index 00000000000..a5de92a7150 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/next_day_deferred_projection.sql @@ -0,0 +1,113 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- MinSparkVersion: 4.0 +-- Config: spark.sql.ansi.enabled=true +-- Config: spark.sql.codegen.wholeStage=true +-- Config: spark.sql.codegen.factoryMode=FALLBACK +-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true +-- Config: spark.sql.shuffle.partitions=1 +-- ConfigMatrix: spark.sql.adaptive.enabled=false,true + +statement +CREATE TABLE test_next_day_deferred USING parquet AS +SELECT /*+ COALESCE(1) */ id, d, dow, CAST(NULL AS DATE) AS null_date FROM VALUES + (1, date('2024-01-01'), 'Monday'), + (2, date('2024-01-01'), 'NOT_A_DAY') AS t(id, d, dow) + +-- A strict parent consumes next_day on every row; deferral alone must not disable Comet. +query +SELECT date_add(x, 1), r FROM ( + SELECT next_day(d, 'Monday' COLLATE UTF8_LCASE) AS x, rand(123L) AS r + FROM test_next_day_deferred) p + +-- rand keeps the lower Project in place. Its BIGINT seed is required for Comet admission. +-- Spark evaluates the filter before evaluating the deferred next_day output. +query expect_fallback(deferred projection) +SELECT x FROM ( + SELECT next_day(d, dow COLLATE UTF8_LCASE) AS x, rand(123L) AS r + FROM test_next_day_deferred) p +WHERE r < 0.0 + +-- Null-safe join keys avoid an inferred is-not-null Filter. Only the matching row's payload +-- is evaluated. Retaining r in the output prevents pruning the nondeterministic expression. +query expect_fallback(deferred projection) +SELECT /*+ BROADCAST(k) */ p.x, p.r FROM ( + SELECT id, next_day(d, dow COLLATE UTF8_LCASE) AS x, rand(123L) AS r + FROM test_next_day_deferred) p +INNER JOIN (SELECT * FROM VALUES (1) AS keys(id)) k ON p.id <=> k.id + +-- A parent Project can also skip a lower Project's output, without any Filter or join. +query expect_fallback(deferred projection) +SELECT IF(r >= 0.0, date('2024-01-01'), x) FROM ( + SELECT next_day(d, dow COLLATE UTF8_LCASE) AS x, rand(123L) AS r + FROM test_next_day_deferred) p + +-- A comparison still skips its right operand when its left column is null. +query expect_fallback(deferred projection) +SELECT null_date < x, r FROM ( + SELECT null_date, next_day(d, dow COLLATE UTF8_LCASE) AS x, rand(123L) AS r + FROM test_next_day_deferred) p + +-- try_add catches errors from the deferred next_day inside Spark's generated try block. +query expect_fallback(deferred projection) +SELECT try_add(x, 1), r FROM ( + SELECT next_day(d, dow COLLATE UTF8_LCASE) AS x, rand(123L) AS r + FROM test_next_day_deferred) p + +-- When the filter consumes the invalid row, its deferred expression must still throw. +query expect_error(Illegal input for day of week) +SELECT x FROM ( + SELECT next_day(d, dow COLLATE UTF8_LCASE) AS x, rand(123L) AS r + FROM test_next_day_deferred) p +WHERE r >= 0.0 + +-- Spark keeps next_day inside the grouped aggregate's result expressions and evaluates it +-- only after the filter admits a group. The aggregate arguments themselves do not throw. +query expect_fallback(deferred projection) +SELECT d FROM ( + SELECT id, + next_day(DATE '2024-01-01', CAST(max(id) AS STRING) COLLATE UTF8_LCASE) AS d, + rand(123L) AS r + FROM test_next_day_deferred GROUP BY id) a +WHERE r < 0.0 + +-- An admitted group still evaluates the invalid aggregate result. +query expect_error(Illegal input for day of week) +SELECT d FROM ( + SELECT id, + next_day(DATE '2024-01-01', CAST(max(id) AS STRING) COLLATE UTF8_LCASE) AS d, + rand(123L) AS r + FROM test_next_day_deferred GROUP BY id) a +WHERE r >= 0.0 + +-- Without grouping keys, Spark evaluates aggregate results before the filter. The invalid +-- result must therefore throw even though the filter would discard it. +query expect_error(Illegal input for day of week) +SELECT d FROM ( + SELECT next_day(DATE '2024-01-01', CAST(max(id) AS STRING) COLLATE UTF8_LCASE) AS d, + rand(123L) AS r + FROM test_next_day_deferred) a +WHERE r < 0.0 + +-- Ungrouped result expressions remain accelerated when their evaluation is safe. +query +SELECT d FROM ( + SELECT next_day(date_add(DATE '2024-01-01', max(id)), 'Monday' COLLATE UTF8_LCASE) AS d, + rand(123L) AS r + FROM test_next_day_deferred) a +WHERE r < 0.0 diff --git a/spark/src/test/resources/sql-tests/expressions/string/levenshtein_collation_interpreted.sql b/spark/src/test/resources/sql-tests/expressions/string/levenshtein_collation_interpreted.sql new file mode 100644 index 00000000000..2d2d1384157 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/string/levenshtein_collation_interpreted.sql @@ -0,0 +1,40 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- MinSparkVersion: 4.0 +-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true +-- Config: spark.sql.codegen.factoryMode=NO_CODEGEN +-- Config: spark.sql.codegen.wholeStage=false + +-- Spark's interpreted Levenshtein unboxes a NULL threshold as zero, while its generated +-- code returns NULL. Keep nullable thresholds on Spark when interpreted execution is requested. +statement +CREATE TABLE test_levenshtein_collated_interpreted(s1 string, s2 string, threshold int) USING parquet + +statement +INSERT INTO test_levenshtein_collated_interpreted VALUES ('', '', NULL), ('a', 'b', NULL), ('kitten', 'sitting', 3) + +query expect_fallback(NO_CODEGEN) +SELECT levenshtein(s1 COLLATE UTF8_LCASE, s2, threshold) FROM test_levenshtein_collated_interpreted + +-- Hypot dispatches its whole subtree, so a check only on the Levenshtein serde is insufficient. +query expect_fallback(NO_CODEGEN) +SELECT hypot(levenshtein(s1 COLLATE UTF8_LCASE, s2, threshold), 1.0) FROM test_levenshtein_collated_interpreted + +-- Both safe forms still use the dispatcher in this mode. +query +SELECT levenshtein(s1 COLLATE UTF8_LCASE, s2), levenshtein(s1 COLLATE UTF8_LCASE, s2, coalesce(threshold, 0)) FROM test_levenshtein_collated_interpreted diff --git a/spark/src/test/resources/sql-tests/expressions/string/levenshtein_interpreted_arguments.sql b/spark/src/test/resources/sql-tests/expressions/string/levenshtein_interpreted_arguments.sql new file mode 100644 index 00000000000..bb87c6b6380 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/string/levenshtein_interpreted_arguments.sql @@ -0,0 +1,49 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- MinSparkVersion: 4.0 +-- Config: spark.sql.ansi.enabled=true +-- Config: spark.sql.codegen.factoryMode=FALLBACK +-- Config: spark.sql.codegen.wholeStage=true +-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true + +statement +CREATE TABLE test_levenshtein_arguments(s1 string, s2 string, threshold int) USING parquet + +statement +INSERT INTO test_levenshtein_arguments VALUES ('', '', NULL), ('a', 'b', NULL), ('a', 'b', 1) + +-- ArrayFilter interprets its argument even when codegen is enabled. NULL thresholds become +-- zero, so array_compact must retain [0] and [-1]. Ordinary inputs still use the native fast path. +query +SELECT array_compact(array(levenshtein(s1 COLLATE UTF8_LCASE, s2, threshold))), + array_compact(array(threshold, CAST(NULL AS INT))) FROM test_levenshtein_arguments + +-- Imperative aggregates call eval on their argument: NULL thresholds become zero, yielding +-- distances 0, -1 and 1. Generating the argument instead would discard the first two values. +query expect_fallback(interpreted evaluation in aggregate arguments) +SELECT sort_array(collect_list(levenshtein(s1 COLLATE UTF8_LCASE, s2, threshold))) FROM test_levenshtein_arguments + +-- Cover the ordinary imperative aggregate path as well as collect_list's object buffer. +query expect_fallback(interpreted evaluation in aggregate arguments) +SELECT approx_count_distinct(levenshtein(s1 COLLATE UTF8_LCASE, s2, threshold)) FROM test_levenshtein_arguments + +-- A nonnullable threshold is safe for collect_list. SUM uses generated arguments in this mode, +-- so its nullable-threshold argument also stays accelerated. +query +SELECT sort_array(collect_list(levenshtein(s1 COLLATE UTF8_LCASE, s2, coalesce(threshold, 0)))), + sum(levenshtein(s1 COLLATE UTF8_LCASE, s2, threshold)) FROM test_levenshtein_arguments 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 f0202c25b54..aa845416ad5 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -28,10 +28,12 @@ import org.scalatest.PrivateMethodTester._ import org.apache.logging.log4j.Level import org.apache.spark.sql._ +import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.catalyst.optimizer.{BuildRight, NormalizeNaNAndZero} +import org.apache.spark.sql.catalyst.plans.{Inner, LeftSemi} +import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, Attribute, AttributeReference, DateAdd, Expression, ExpressionInfo, GreaterThan, In, InSet, KnownFloatingPointNormalized, LessThan, Literal, NextDay, Not, SortOrder} import org.apache.spark.sql.catalyst.FunctionIdentifier -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.catalyst.plans.logical.{LocalRelation, LogicalPlan} import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RoundRobinPartitioning} import org.apache.spark.sql.catalyst.rules.Rule @@ -2344,4 +2346,219 @@ class CometExecRuleSuite extends CometTestBase { assert(warnings.head.contains(s"${CometConf.COMET_CONVERT_FROM_RANGE_ENABLED.key}=true")) } + test("ANSI next_day stays in Spark below LIMIT with or without dispatch") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val input = + createSparkPlan(spark, "SELECT * FROM VALUES ('Monday'), ('NOT_A_DAY') AS inputs(dow)") + val date = Literal.create(java.sql.Date.valueOf("2024-01-01"), DataTypes.DateType) + for ((dispatch, ansi) <- Seq(("true", true), ("false", true), ("true", false))) { + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> dispatch) { + val project = ProjectExec( + Seq(Alias(NextDay(date, input.output.head, failOnError = ansi), "next")()), + input) + val transformed = applyCometExecRule(CollectLimitExec(1, project)) + withClue(s"dispatch=$dispatch, ANSI=$ansi") { + if (ansi) { + assert(transformed.isInstanceOf[CollectLimitExec]) + assert(countOperators(transformed, classOf[ProjectExec]) == 1) + assert(countOperators(transformed, classOf[CometProjectExec]) == 0) + assert( + new ExtendedExplainInfo() + .getFallbackReasons(transformed) + .exists(_.contains("next_day requires Spark evaluation below LIMIT"))) + } else { + assert(countOperators(transformed, classOf[CometProjectExec]) == 1) + assert(countOperators(transformed, classOf[ProjectExec]) == 0) + } + assert(countOperators(transformed, classOf[CometLocalTableScanExec]) == 1) + } + } + } + } + } + + test("ANSI next_day first-match join conditions stay in Spark without dispatch") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { + val left = createSparkPlan(spark, "SELECT * FROM VALUES (1) AS inputs(id)") + val right = createSparkPlan( + spark, + "SELECT * FROM VALUES (1, 'Monday'), (1, 'NOT_A_DAY') AS inputs(id, dow)") + val date = Literal.create(java.sql.Date.valueOf("2024-01-01"), DataTypes.DateType) + val condition = Some(GreaterThan(NextDay(date, right.output(1), failOnError = true), date)) + val join = ShuffledHashJoinExec( + Seq(left.output.head), + Seq(right.output.head), + LeftSemi, + BuildRight, + condition, + left, + right) + val transformed = applyCometExecRule(join) + assert(transformed.isInstanceOf[ShuffledHashJoinExec]) + assert(new ExtendedExplainInfo() + .getFallbackReasons(transformed) + .exists(_.contains("next_day requires Spark evaluation in first-match join conditions"))) + assert(countOperators(transformed, classOf[CometLocalTableScanExec]) == 2) + + // A protected Spark join must also survive the optional SMJ-to-SHJ rewrite. + withSQLConf(CometConf.COMET_FORCE_SHJ.key -> "true") { + val sortMergeJoin = SortMergeJoinExec( + Seq(left.output.head), + Seq(right.output.head), + LeftSemi, + condition, + left, + right) + val preserved = applyCometExecRule(sortMergeJoin) + assert(preserved.isInstanceOf[SortMergeJoinExec]) + assert( + new ExtendedExplainInfo() + .getFallbackReasons(preserved) + .exists( + _.contains("next_day requires Spark evaluation in first-match join conditions"))) + } + } + } + + private def withAnsiNextDayProject(f: ProjectExec => Unit): Unit = { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val input = createSparkPlan( + spark, + "SELECT * FROM VALUES (1, 'Monday'), (2, 'NOT_A_DAY') AS inputs(id, dow)") + val date = Literal.create(java.sql.Date.valueOf("2024-01-01"), DataTypes.DateType) + f( + ProjectExec( + Seq( + input.output.head, + Alias(NextDay(date, input.output(1), failOnError = true), "next")()), + input)) + } + } + + test("ANSI next_day stays native across eager codegen input boundaries") { + withAnsiNextDayProject { project => + val sorted = + SortExec(Seq(SortOrder(project.output.head, Ascending)), global = false, project) + val filtered = FilterExec(LessThan(project.output.head, Literal(2)), sorted) + val right = createSparkPlan(spark, "SELECT * FROM VALUES (1) AS inputs(id)") + // Spark generates a separate stage for each shuffled-hash-join input, even without an + // explicit exchange in this physical plan. Its projected values are already materialized. + val joined = ShuffledHashJoinExec( + Seq(project.output.head), + Seq(right.output.head), + Inner, + BuildRight, + None, + ProjectExec(project.projectList, project.child), + right) + for (plan <- Seq(filtered, joined)) { + val transformed = applyCometExecRule(plan) + withClue(plan.nodeName) { + assert(countOperators(transformed, classOf[CometProjectExec]) == 1) + assert(countOperators(transformed, classOf[ProjectExec]) == 0) + assert( + !new ExtendedExplainInfo() + .getFallbackReasons(transformed) + .exists(_.contains("deferred projection"))) + } + } + } + } + + test("strict Project inputs stay native unless their output is deferred") { + withAnsiNextDayProject { project => + def strictProject: ProjectExec = ProjectExec( + Seq(project.output.head, Alias(DateAdd(project.output(1), Literal(1)), "next")()), + ProjectExec(project.projectList, project.child)) + + val eager = applyCometExecRule(strictProject) + assert(countOperators(eager, classOf[CometProjectExec]) == 2) + assert(countOperators(eager, classOf[ProjectExec]) == 0) + + // A filter can defer the strict parent's result, and therefore its input as well. + val deferred = strictProject + val filtered = + applyCometExecRule(FilterExec(LessThan(deferred.output.head, Literal(2)), deferred)) + assert(countOperators(filtered, classOf[CometProjectExec]) == 0) + assert(countOperators(filtered, classOf[ProjectExec]) == 2) + } + } + + test("deferred next_day fallback restores a reused native projection through batch bridges") { + withAnsiNextDayProject { project => + val bridges: Seq[SparkPlan => SparkPlan] = Seq( + child => RowToColumnarExec(ColumnarToRowExec(child)), + child => CometSparkToColumnarExec(CometNativeColumnarToRowExec(child))) + bridges.foreach { bridge => + val native = applyCometExecRule(ProjectExec(project.projectList, project.child)) + assert(countOperators(native, classOf[CometProjectExec]) == 1) + val logicalStage = LogicalQueryStage(LocalRelation(native.output), native) + native.setLogicalLink(logicalStage) + val filtered = FilterExec(LessThan(native.output.head, Literal(2)), bridge(native)) + val transformed = applyCometExecRule(filtered) + for (replanned <- Seq(transformed, applyCometExecRule(transformed))) { + assert(replanned.isInstanceOf[FilterExec]) + assert(countOperators(replanned, classOf[ProjectExec]) == 1) + assert(countOperators(replanned, classOf[CometProjectExec]) == 0) + assert(countOperators(replanned, classOf[CometLocalTableScanExec]) == 1) + assert(replanned.collect { case r: RowToColumnarTransition => r }.isEmpty) + assert( + new ExtendedExplainInfo() + .getFallbackReasons(replanned) + .exists(_.contains("next_day requires Spark evaluation in a deferred projection"))) + assert( + replanned.children.head + .getTagValue(SparkPlan.LOGICAL_PLAN_TAG) + .exists(_ eq logicalStage)) + } + } + } + } + + test("ANSI next_day preserves incompatible aggregate buffers before AQE removes a sort") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.USE_OBJECT_HASH_AGG.key -> "true", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val aggregate = createSparkPlan( + spark, + "SELECT next_day(DATE '2024-01-01', dow) AS next, collect_list(k) AS collected " + + "FROM VALUES (1, 'Monday'), (2, 'Monday') AS t(k, dow) GROUP BY dow") + assert(countOperators(aggregate, classOf[ObjectHashAggregateExec]) == 2) + val native = applyCometExecRule(aggregate) + assert(countOperators(native, classOf[CometHashAggregateExec]) == 2) + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + // The sort currently consumes every result, but AQE can remove it after the Partial + // materializes. Both a fresh plan and a reused native chain need Spark's binary buffers. + for (child <- Seq(aggregate, native)) { + val sorted = + SortExec(Seq(SortOrder(child.output.head, Ascending)), global = false, child) + val offset = applyCometExecRule(CollectLimitExec(-1, sorted, offset = 1)) + assert(countOperators(offset, classOf[CometHashAggregateExec]) == 2) + val transformed = applyCometExecRule(CollectLimitExec(1, sorted)) + for (replanned <- Seq(transformed, applyCometExecRule(transformed))) { + assert(countOperators(replanned, classOf[CometHashAggregateExec]) == 0) + assert(countOperators(replanned, classOf[ObjectHashAggregateExec]) == 2) + assert( + new ExtendedExplainInfo() + .getFallbackReasons(replanned) + .exists( + _.contains("next_day requires Spark aggregate buffers below LIMIT with AQE"))) + } + } + } + } + } } From 5c291b3065f0027ad9a80dd3df2b47698451641a Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 4 Oct 2026 17:58:31 +0000 Subject: [PATCH 6/9] fix: preserve safe expression admission and nullable threshold semantics --- .../user-guide/latest/compatibility/index.md | 16 ++++++++ .../apache/comet/rules/CometExecRule.scala | 5 ++- .../apache/comet/serde/QueryPlanSerde.scala | 6 ++- .../org/apache/comet/serde/datetime.scala | 1 - .../datetime/next_day_safe_evaluation.sql | 40 +++++++++++++++++++ .../levenshtein_collation_interpreted.sql | 6 +-- .../levenshtein_interpreted_arguments.sql | 13 +++--- .../string/levenshtein_nullable_threshold.sql | 34 ++++++++++++++++ .../string/levenshtein_threshold.sql | 4 +- .../comet/rules/CometExecRuleSuite.scala | 8 ++-- .../serde/CometEvaluationMaskSuite.scala | 2 +- 11 files changed, 115 insertions(+), 20 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/datetime/next_day_safe_evaluation.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/string/levenshtein_nullable_threshold.sql diff --git a/docs/source/user-guide/latest/compatibility/index.md b/docs/source/user-guide/latest/compatibility/index.md index 6e23c71752a..02d58a162de 100644 --- a/docs/source/user-guide/latest/compatibility/index.md +++ b/docs/source/user-guide/latest/compatibility/index.md @@ -153,3 +153,19 @@ so users hunting an unexpected value have a single place to check: `spark.sql.codegen.fallback` or when the generated code exceeds `spark.sql.codegen.hugeMethodLimit`, where an intermediate overflow that later cancels out returns `NULL` (or raises under ANSI) in Spark but the recovered value in Comet. + +### `next_day` and nullable `levenshtein` thresholds + +ANSI `next_day` participates in the evaluation-mask policy above. A valid literal weekday +stays native, including below a limit. Dynamic weekdays stay in Spark when a limit, first-match +join, or deferred projection can skip an invalid weekday; eager operands and ordinary sort keys +remain eligible for native execution. The same `spark.comet.exec.preserveEvaluationMasks.enabled` +setting controls the operator-level policy. Per-expression guards still preserve conditional +and filtered/state-dependent aggregate evaluation. The broader family is tracked in +[#6006](https://github.com/apache/datafusion-comet/issues/6006), and the shared policy is introduced +by [#5533](https://github.com/apache/datafusion-comet/pull/5533). + +Three-argument `levenshtein` with a nullable threshold stays in Spark because Spark's interpreted +and generated evaluators disagree for NULL thresholds. This applies to every code-generation +mode, including `FALLBACK`, whose evaluator can change after a compilation failure. Two-argument +calls and non-nullable thresholds remain eligible for native execution. 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 435eb6d0ce9..0c63b5806aa 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -26,8 +26,8 @@ import scala.jdk.CollectionConverters._ import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LeafExpression, LessThan, LessThanOrEqual, NamedExpression, Remainder, SortOrder} -import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge} +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero import org.apache.spark.sql.catalyst.plans.{JoinType, LeftAnti, LeftSemi} import org.apache.spark.sql.catalyst.rules.Rule @@ -1161,7 +1161,8 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) def eagerReferences(expr: Expression): AttributeSet = expr match { case attribute: AttributeReference => AttributeSet(Seq(attribute)) - case _ => AttributeSet(QueryPlanSerde.eagerlyEvaluatedChildren(expr).flatMap(eagerReferences)) + case _ => + AttributeSet(QueryPlanSerde.eagerlyEvaluatedChildren(expr).flatMap(eagerReferences)) } def protect( 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 af7921c9619..a69215a921d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -1241,8 +1241,10 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { val wholeExpressionDispatch = exprSerdeMap.get(current.getClass).exists { handler => (handler.isInstanceOf[CodegenDispatchFallback] && - handler.asInstanceOf[CometExpressionSerde[Expression]] - .getSupportLevel(current).isInstanceOf[Unsupported]) || + handler + .asInstanceOf[CometExpressionSerde[Expression]] + .getSupportLevel(current) + .isInstanceOf[Unsupported]) || (handler.isInstanceOf[CometCodegenDispatch[_]] && !handler.isInstanceOf[NativeOptInAvailable]) } diff --git a/spark/src/main/scala/org/apache/comet/serde/datetime.scala b/spark/src/main/scala/org/apache/comet/serde/datetime.scala index 6de0aa96e22..fa78fd9a4b5 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -424,7 +424,6 @@ object CometNextDay } } - private val collationReason = DatetimeCollation.reason("next_day") override def getIncompatibleReasons(): Seq[String] = diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/next_day_safe_evaluation.sql b/spark/src/test/resources/sql-tests/expressions/datetime/next_day_safe_evaluation.sql new file mode 100644 index 00000000000..292711e173a --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/next_day_safe_evaluation.sql @@ -0,0 +1,40 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- Config: spark.sql.ansi.enabled=true + +statement +CREATE TABLE next_day_safe(d DATE, dow STRING) USING parquet + +statement +INSERT INTO next_day_safe VALUES (DATE '2026-09-01', 'Monday'), (NULL, 'Tuesday') + +-- A valid literal cannot raise an invalid-weekday error, including below LIMIT. +query expect_native(next_day) +SELECT next_day(d, 'Monday') FROM next_day_safe LIMIT 1 + +-- Spark evaluates the date operand and these eager inputs for every consumed row. +query expect_native(next_day) +SELECT date_add(next_day(d, dow), 1), datediff(next_day(d, dow), d), + concat(CAST(next_day(d, dow) AS STRING), 'x'), greatest(next_day(d, dow), d) +FROM next_day_safe + +query expect_native(next_day) +SELECT d FROM next_day_safe ORDER BY next_day(d, dow) + +query expect_native(next_day) +SELECT max(next_day(d, dow)) FROM next_day_safe diff --git a/spark/src/test/resources/sql-tests/expressions/string/levenshtein_collation_interpreted.sql b/spark/src/test/resources/sql-tests/expressions/string/levenshtein_collation_interpreted.sql index 2d2d1384157..26500799fac 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/levenshtein_collation_interpreted.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/levenshtein_collation_interpreted.sql @@ -21,18 +21,18 @@ -- Config: spark.sql.codegen.wholeStage=false -- Spark's interpreted Levenshtein unboxes a NULL threshold as zero, while its generated --- code returns NULL. Keep nullable thresholds on Spark when interpreted execution is requested. +-- code returns NULL. Keep nullable thresholds on Spark in every factory mode. statement CREATE TABLE test_levenshtein_collated_interpreted(s1 string, s2 string, threshold int) USING parquet statement INSERT INTO test_levenshtein_collated_interpreted VALUES ('', '', NULL), ('a', 'b', NULL), ('kitten', 'sitting', 3) -query expect_fallback(NO_CODEGEN) +query expect_fallback(levenshtein with a nullable threshold requires Spark evaluation) SELECT levenshtein(s1 COLLATE UTF8_LCASE, s2, threshold) FROM test_levenshtein_collated_interpreted -- Hypot dispatches its whole subtree, so a check only on the Levenshtein serde is insufficient. -query expect_fallback(NO_CODEGEN) +query expect_fallback(levenshtein with a nullable threshold requires Spark evaluation) SELECT hypot(levenshtein(s1 COLLATE UTF8_LCASE, s2, threshold), 1.0) FROM test_levenshtein_collated_interpreted -- Both safe forms still use the dispatcher in this mode. diff --git a/spark/src/test/resources/sql-tests/expressions/string/levenshtein_interpreted_arguments.sql b/spark/src/test/resources/sql-tests/expressions/string/levenshtein_interpreted_arguments.sql index bb87c6b6380..6138b64e7c2 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/levenshtein_interpreted_arguments.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/levenshtein_interpreted_arguments.sql @@ -29,7 +29,7 @@ INSERT INTO test_levenshtein_arguments VALUES ('', '', NULL), ('a', 'b', NULL), -- ArrayFilter interprets its argument even when codegen is enabled. NULL thresholds become -- zero, so array_compact must retain [0] and [-1]. Ordinary inputs still use the native fast path. -query +query expect_fallback(levenshtein with a nullable threshold requires Spark evaluation) SELECT array_compact(array(levenshtein(s1 COLLATE UTF8_LCASE, s2, threshold))), array_compact(array(threshold, CAST(NULL AS INT))) FROM test_levenshtein_arguments @@ -42,8 +42,11 @@ SELECT sort_array(collect_list(levenshtein(s1 COLLATE UTF8_LCASE, s2, threshold) query expect_fallback(interpreted evaluation in aggregate arguments) SELECT approx_count_distinct(levenshtein(s1 COLLATE UTF8_LCASE, s2, threshold)) FROM test_levenshtein_arguments --- A nonnullable threshold is safe for collect_list. SUM uses generated arguments in this mode, --- so its nullable-threshold argument also stays accelerated. +-- A nonnullable threshold stays accelerated, including an imperative aggregate argument. query -SELECT sort_array(collect_list(levenshtein(s1 COLLATE UTF8_LCASE, s2, coalesce(threshold, 0)))), - sum(levenshtein(s1 COLLATE UTF8_LCASE, s2, threshold)) FROM test_levenshtein_arguments +SELECT sort_array(collect_list(levenshtein(s1 COLLATE UTF8_LCASE, s2, coalesce(threshold, 0)))) +FROM test_levenshtein_arguments + +-- A declarative aggregate can also fall back from generated to interpreted evaluation. +query expect_fallback(levenshtein with a nullable threshold requires Spark evaluation) +SELECT sum(levenshtein(s1 COLLATE UTF8_LCASE, s2, threshold)) FROM test_levenshtein_arguments diff --git a/spark/src/test/resources/sql-tests/expressions/string/levenshtein_nullable_threshold.sql b/spark/src/test/resources/sql-tests/expressions/string/levenshtein_nullable_threshold.sql new file mode 100644 index 00000000000..6d5f800f396 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/string/levenshtein_nullable_threshold.sql @@ -0,0 +1,34 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- MinSparkVersion: 3.5 +-- ConfigMatrix: spark.sql.codegen.factoryMode=FALLBACK,CODEGEN_ONLY,NO_CODEGEN + +statement +CREATE TABLE levenshtein_nullable(s STRING, t STRING, threshold INT) USING parquet + +statement +INSERT INTO levenshtein_nullable VALUES ('same', 'same', NULL), ('a', 'b', 2) + +-- A nullable threshold stays with Spark, including the default FALLBACK mode where a +-- runtime code-generation failure can switch Spark from generated to interpreted results. +query expect_fallback(levenshtein with a nullable threshold requires Spark evaluation) +SELECT levenshtein(s, t, threshold) FROM levenshtein_nullable + +-- The two-argument and known non-null threshold forms retain native execution. +query expect_native(levenshtein) +SELECT levenshtein(s, t), levenshtein(s, t, 2) FROM levenshtein_nullable diff --git a/spark/src/test/resources/sql-tests/expressions/string/levenshtein_threshold.sql b/spark/src/test/resources/sql-tests/expressions/string/levenshtein_threshold.sql index a0756a34033..8a9d9bf12b8 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/levenshtein_threshold.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/levenshtein_threshold.sql @@ -46,7 +46,7 @@ CREATE TABLE test_levenshtein_col(s1 string, s2 string, threshold int) USING par statement INSERT INTO test_levenshtein_col VALUES ('kitten', 'sitting', 2), ('frog', 'fog', 5), ('abc', 'abc', 0), ('hello', 'world', 3) -query +query expect_fallback(levenshtein with a nullable threshold requires Spark evaluation) SELECT levenshtein(s1, s2, threshold) FROM test_levenshtein_col -- threshold as column with NULLs @@ -56,5 +56,5 @@ CREATE TABLE test_levenshtein_col_nulls(s1 string, s2 string, threshold int) USI statement INSERT INTO test_levenshtein_col_nulls VALUES ('abc', 'adc', 2), ('hello', 'world', NULL), (NULL, 'test', 3), ('frog', 'fog', -1) -query +query expect_fallback(levenshtein with a nullable threshold requires Spark evaluation) SELECT levenshtein(s1, s2, threshold) FROM test_levenshtein_col_nulls 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 aa845416ad5..d38f43f983a 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -28,12 +28,11 @@ import org.scalatest.PrivateMethodTester._ import org.apache.logging.log4j.Level import org.apache.spark.sql._ -import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} -import org.apache.spark.sql.catalyst.optimizer.{BuildRight, NormalizeNaNAndZero} -import org.apache.spark.sql.catalyst.plans.{Inner, LeftSemi} -import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, Attribute, AttributeReference, DateAdd, Expression, ExpressionInfo, GreaterThan, In, InSet, KnownFloatingPointNormalized, LessThan, Literal, NextDay, Not, SortOrder} import org.apache.spark.sql.catalyst.FunctionIdentifier +import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, Attribute, AttributeReference, DateAdd, Expression, ExpressionInfo, GreaterThan, In, InSet, KnownFloatingPointNormalized, LessThan, Literal, NextDay, Not, SortOrder} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, BloomFilterAggregate, Final, Min, Partial, PartialMerge} +import org.apache.spark.sql.catalyst.optimizer.{BuildRight, NormalizeNaNAndZero} +import org.apache.spark.sql.catalyst.plans.{Inner, LeftSemi} import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan} import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RoundRobinPartitioning} import org.apache.spark.sql.catalyst.rules.Rule @@ -45,6 +44,7 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffl import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec} import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec} +import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataTypes, DoubleType, FloatType, StructField, StructType} diff --git a/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala b/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala index 0d2ae170df1..dfd062add3a 100644 --- a/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala +++ b/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala @@ -161,7 +161,7 @@ class CometEvaluationMaskSuite extends CometTestBase { test("LIMIT masks follow the decoder policy with ANSI, strict and compound inputs") { assert(QueryPlanSerde.exprSerdeMap.collect { case (cls, _: RequiresSparkEvaluationMask[_]) => cls - }.toSet == Set(classOf[UnBase64])) + }.toSet == Set(classOf[UnBase64], classOf[NextDay])) val input = AttributeReference("encoded", StringType)() withInputs( "masked" -> "SELECT * FROM VALUES (1, 'YWJj'), (2, 'A') AS t(k, bad)", From 3e19377ffd74ae8016ea767c93012ce55b7cb373 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 4 Oct 2026 19:40:20 +0000 Subject: [PATCH 7/9] test: exercise fused top-K projection on Spark 3.5 --- .../org/apache/comet/serde/CometEvaluationMaskSuite.scala | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala b/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala index dfd062add3a..8350a6c4dc2 100644 --- a/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala +++ b/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala @@ -503,8 +503,8 @@ class CometEvaluationMaskSuite extends CometTestBase { withInputs( "fused_mask" -> "SELECT * FROM VALUES (1, 'YWJj'), (2, 'YWJj'), (3, 'YWJj'), (4, 'YWJj') AS t(k, bad)") { - val query = "SELECT k, hex(unbase64(bad)) AS decoded FROM " + - "(SELECT * FROM fused_mask ORDER BY k LIMIT 2 OFFSET 1) t" + val query = "SELECT k, hex(unbase64(bad)) AS decoded " + + "FROM (SELECT * FROM fused_mask ORDER BY k) t LIMIT 2 OFFSET 1" val expected = Seq(Row(2, "616263"), Row(3, "616263")) for (aqe <- Seq(true, false); fusion <- Seq(false, true)) { withSQLConf( From 5ec0f677a375605ae94f6dfbac699f9f828e910a Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 4 Oct 2026 19:54:54 +0000 Subject: [PATCH 8/9] fix: retain native inputs required by generated filters --- .../apache/comet/rules/CometExecRule.scala | 28 ++++++- .../comet/rules/CometExecRuleSuite.scala | 81 ++++++++++++++++++- 2 files changed, 107 insertions(+), 2 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 0c63b5806aa..f752e0a3799 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -25,7 +25,7 @@ import scala.collection.mutable.ListBuffer import scala.jdk.CollectionConverters._ import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LeafExpression, LessThan, LessThanOrEqual, NamedExpression, Remainder, SortOrder} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, IsNotNull, KnownFloatingPointNormalized, LeafExpression, LessThan, LessThanOrEqual, NamedExpression, PredicateHelper, Remainder, SortOrder} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero @@ -196,6 +196,7 @@ object CometExecRule { case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) extends Rule[SparkPlan] with CometTypeShim + with PredicateHelper with ShimSubqueryBroadcast { private lazy val showTransformations = CometConf.COMET_EXPLAIN_TRANSFORMATIONS.get() @@ -1165,6 +1166,29 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) AttributeSet(QueryPlanSerde.eagerlyEvaluatedChildren(expr).flatMap(eagerReferences)) } + def eagerFilterInputs(filter: FilterExec): AttributeSet = { + // GeneratePredicateHelper loads every referenced projected value before evaluating each + // predicate, but can skip later conjuncts. It moves null-intolerant IsNotNull checks ahead + // of predicates that reference those inputs, including inferred checks for attributes + // inside compound null checks. Match the first check Spark actually emits. + val (notNullPredicates, otherPredicates) = + splitConjunctivePredicates(filter.condition).partition { + case IsNotNull(child) => + isNullIntolerant(child) && child.references.subsetOf(filter.child.outputSet) + case _ => false + } + val nullCheckedInputs = AttributeSet(notNullPredicates.flatMap(_.references)) + otherPredicates.headOption match { + case Some(predicate) => + predicate.references.find(nullCheckedInputs.contains) match { + case Some(firstNullCheckedInput) => AttributeSet(Seq(firstNullCheckedInput)) + case None => predicate.references + } + case None => + notNullPredicates.headOption.map(_.references).getOrElse(AttributeSet.empty) + } + } + def protect( node: SparkPlan, belowLimit: Boolean, @@ -1216,6 +1240,8 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) expr.deterministic && deferredInputs.contains(expr.toAttribute)) project.inputSet -- project.usedInputs -- AttributeSet(eagerOutputs.flatMap(eagerReferences)) + case filter: FilterExec if supportsWholeStage(filter) => + filter.inputSet -- filter.usedInputs -- eagerFilterInputs(filter) case codegen: CodegenSupport if !materializesInput && supportsWholeStage(original) => original.inputSet -- codegen.usedInputs case _ => AttributeSet.empty 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 d38f43f983a..5cfec552230 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -29,7 +29,7 @@ import org.scalatest.PrivateMethodTester._ 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.{Alias, Ascending, Attribute, AttributeReference, DateAdd, Expression, ExpressionInfo, GreaterThan, In, InSet, KnownFloatingPointNormalized, LessThan, Literal, NextDay, Not, SortOrder} +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Ascending, Attribute, AttributeReference, DateAdd, DateDiff, Expression, ExpressionInfo, GreaterThan, In, InSet, IsNotNull, KnownFloatingPointNormalized, LessThan, Literal, NextDay, Not, Or, SortOrder} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, BloomFilterAggregate, Final, Min, Partial, PartialMerge} import org.apache.spark.sql.catalyst.optimizer.{BuildRight, NormalizeNaNAndZero} import org.apache.spark.sql.catalyst.plans.{Inner, LeftSemi} @@ -2475,6 +2475,85 @@ class CometExecRuleSuite extends CometTestBase { } } + test("ANSI next_day filter inputs follow generated predicate order") { + withAnsiNextDayProject { original => + def project: ProjectExec = ProjectExec( + original.projectList :+ + Alias(Literal.create(null, DataTypes.DateType), "nullable_date")(), + original.child) + val date = Literal.create(java.sql.Date.valueOf("2020-01-01"), DataTypes.DateType) + val predicates: Seq[(ProjectExec => Expression, Boolean)] = Seq( + (p => GreaterThan(p.output(1), date), true), + (p => GreaterThan(p.output(2), p.output(1)), true), + (p => And(GreaterThan(p.output(1), date), LessThan(p.output.head, Literal(2))), true), + (p => And(LessThan(p.output.head, Literal(2)), GreaterThan(p.output(1), date)), false), + // A single predicate loads all referenced projected values before its expression code. + (p => Or(LessThan(p.output.head, Literal(2)), GreaterThan(p.output(1), date)), true), + (p => And(IsNotNull(p.output(1)), GreaterThan(p.output(1), date)), true), + (p => And(IsNotNull(p.output(1)), IsNotNull(p.output(2))), true), + (p => And(IsNotNull(p.output(2)), IsNotNull(p.output(1))), false), + // Spark moves this null check ahead of the comparison, even though it appears last. + (p => And(GreaterThan(p.output(1), p.output(2)), IsNotNull(p.output(2))), false), + // A compound null check also causes Spark to inject checks for its input attributes. + ( + p => + And( + GreaterThan(p.output(1), p.output(2)), + IsNotNull(DateAdd(p.output(2), p.output.head))), + false), + (p => IsNotNull(DateDiff(p.output(2), p.output(1))), true)) + for ((predicate, native) <- predicates) { + val child = project + val condition = predicate(child) + val transformed = applyCometExecRule(FilterExec(condition, child)) + withClue(condition.sql) { + assert(countOperators(transformed, classOf[CometProjectExec]) == (if (native) 1 else 0)) + assert(countOperators(transformed, classOf[ProjectExec]) == (if (native) 0 else 1)) + } + } + } + } + + test("ANSI next_day required by a filter stays native and preserves errors") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.ANSI_ENABLED.key -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + val date = java.sql.Date.valueOf("2024-01-01") + val preserve = CometConf.COMET_EXEC_PRESERVE_EVALUATION_MASKS_ENABLED.key + def query(predicate: String): String = + s"""SELECT x, r FROM + |(SELECT next_day(_1, _2) AS x, rand(1L) AS r FROM filter_next_day) + |WHERE $predicate""".stripMargin + + withParquetTable(Seq((date, "Monday")), "filter_next_day") { + for (enabled <- Seq(false, true)) { + withSQLConf(preserve -> enabled.toString) { + val (_, plan) = checkSparkAnswerAndImpl( + sql(query("x >= DATE '2020-01-01'")), + native = Seq("next_day")) + assert(collect(plan) { case p: CometProjectExec => p }.size == 1, plan.toString) + assert(collect(plan) { case f: CometFilterExec => f }.size == 1, plan.toString) + } + } + } + withParquetTable(Seq((date, "NOT_A_DAY")), "filter_next_day") { + // The required projected value fails even when the date comparison would reject it. + val required = sql(query("x < DATE '2020-01-01'")) + assertExpressionImpl(required.queryExecution.executedPlan, Seq("next_day"), Seq.empty) + val errorClass = if (isSpark40Plus) "ILLEGAL_DAY_OF_WEEK" else "_LEGACY_ERROR_TEMP_2000" + val error = checkSparkError(required, errorClass) + assert(error.getMessage.contains("Illegal input for day of week")) + // An earlier conjunct can still reject the row before evaluating next_day. + checkSparkAnswerAndFallbackReason( + query("r < 0 AND x < DATE '2020-01-01'"), + "next_day requires Spark evaluation in a deferred projection") + () + } + } + } + test("strict Project inputs stay native unless their output is deferred") { withAnsiNextDayProject { project => def strictProject: ProjectExec = ProjectExec( From c7c8c7340232d62727fe9c243646ee9ccb7ae681 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Mon, 5 Oct 2026 21:28:13 +0000 Subject: [PATCH 9/9] Preserve native next_day under eager date and array parents --- .../apache/comet/serde/QueryPlanSerde.scala | 4 +- .../datetime/next_day_safe_evaluation.sql | 29 ++++++++++++++ .../comet/rules/CometExecRuleSuite.scala | 39 +++++++++++-------- 3 files changed, 54 insertions(+), 18 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 a69215a921d..84015973c72 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -1212,9 +1212,9 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { case _: TryEval => Seq.empty case order: SortOrder => Seq(order.child) case unary: UnaryExpression => Seq(unary.child) - case _: DateAdd | _: DateDiff | _: BinaryComparison => + case _: DateAdd | _: DateSub | _: DateDiff | _: NextDay | _: BinaryComparison => expr.children.take(if (expr.children.head.nullable) 1 else 2) - case _: Concat | _: Greatest | _: Least => expr.children + case _: Concat | _: Greatest | _: Least | _: CreateArray => expr.children case _ => Seq.empty } diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/next_day_safe_evaluation.sql b/spark/src/test/resources/sql-tests/expressions/datetime/next_day_safe_evaluation.sql index 292711e173a..c7d4b7c8f77 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/next_day_safe_evaluation.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/next_day_safe_evaluation.sql @@ -16,6 +16,7 @@ -- under the License. -- Config: spark.sql.ansi.enabled=true +-- ConfigMatrix: spark.sql.codegen.wholeStage=false,true statement CREATE TABLE next_day_safe(d DATE, dow STRING) USING parquet @@ -38,3 +39,31 @@ SELECT d FROM next_day_safe ORDER BY next_day(d, dow) query expect_native(next_day) SELECT max(next_day(d, dow)) FROM next_day_safe + +-- Each parent consumes next_day even when another array element or the date is null. +query expect_native(next_day) +SELECT date_sub(next_day(d, dow), 1), next_day(next_day(d, dow), 'Friday'), + array(CAST(NULL AS DATE), next_day(d, dow)) +FROM next_day_safe + +statement +CREATE TABLE next_day_eager_invalid(d DATE, dow STRING, null_date DATE) USING parquet + +statement +INSERT INTO next_day_eager_invalid VALUES (DATE '2024-01-01', 'NOT_A_DAY', NULL) + +query expect_error(Illegal input for day of week) +SELECT date_sub(next_day(d, dow), 1) FROM next_day_eager_invalid + +query expect_error(Illegal input for day of week) +SELECT next_day(next_day(d, dow), 'Friday') FROM next_day_eager_invalid + +query expect_error(Illegal input for day of week) +SELECT array(CAST(NULL AS DATE), next_day(d, dow)) FROM next_day_eager_invalid + +-- Binary parents still skip their right input when the left date is null. +query expect_fallback(next_day requires Spark evaluation) +SELECT date_sub(null_date, datediff(next_day(d, dow), d)) FROM next_day_eager_invalid + +query expect_fallback(next_day requires Spark evaluation) +SELECT next_day(null_date, CAST(next_day(d, dow) AS STRING)) FROM next_day_eager_invalid 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 5cfec552230..aee592417d9 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -29,7 +29,7 @@ import org.scalatest.PrivateMethodTester._ 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.{Alias, And, Ascending, Attribute, AttributeReference, DateAdd, DateDiff, Expression, ExpressionInfo, GreaterThan, In, InSet, IsNotNull, KnownFloatingPointNormalized, LessThan, Literal, NextDay, Not, Or, SortOrder} +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Ascending, Attribute, AttributeReference, CreateArray, DateAdd, DateDiff, DateSub, Expression, ExpressionInfo, GreaterThan, In, InSet, IsNotNull, KnownFloatingPointNormalized, LessThan, Literal, NextDay, Not, Or, SortOrder} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, BloomFilterAggregate, Final, Min, Partial, PartialMerge} import org.apache.spark.sql.catalyst.optimizer.{BuildRight, NormalizeNaNAndZero} import org.apache.spark.sql.catalyst.plans.{Inner, LeftSemi} @@ -46,7 +46,7 @@ import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec} import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{DataTypes, DoubleType, FloatType, StructField, StructType} +import org.apache.spark.sql.types.{DataTypes, DateType, DoubleType, FloatType, StructField, StructType} import org.apache.comet.{CometConf, CometCoverageStats, CometExplainInfo, CometSparkSessionExtensions, ConfigEntry, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark42Plus, withFallbackReason} @@ -2556,20 +2556,27 @@ class CometExecRuleSuite extends CometTestBase { test("strict Project inputs stay native unless their output is deferred") { withAnsiNextDayProject { project => - def strictProject: ProjectExec = ProjectExec( - Seq(project.output.head, Alias(DateAdd(project.output(1), Literal(1)), "next")()), - ProjectExec(project.projectList, project.child)) - - val eager = applyCometExecRule(strictProject) - assert(countOperators(eager, classOf[CometProjectExec]) == 2) - assert(countOperators(eager, classOf[ProjectExec]) == 0) - - // A filter can defer the strict parent's result, and therefore its input as well. - val deferred = strictProject - val filtered = - applyCometExecRule(FilterExec(LessThan(deferred.output.head, Literal(2)), deferred)) - assert(countOperators(filtered, classOf[CometProjectExec]) == 0) - assert(countOperators(filtered, classOf[ProjectExec]) == 2) + val eagerParents = Seq[Expression => Expression]( + child => DateAdd(child, Literal(1)), + child => DateSub(child, Literal(1)), + child => NextDay(child, Literal("Friday"), failOnError = true), + child => CreateArray(Seq(Literal.create(null, DateType), child))) + eagerParents.foreach { parent => + def strictProject: ProjectExec = ProjectExec( + Seq(project.output.head, Alias(parent(project.output(1)), "next")()), + ProjectExec(project.projectList, project.child)) + + val eager = applyCometExecRule(strictProject) + assert(countOperators(eager, classOf[CometProjectExec]) == 2, eager.toString) + assert(countOperators(eager, classOf[ProjectExec]) == 0, eager.toString) + + // A filter can defer the strict parent's result, and therefore its input as well. + val deferred = strictProject + val filtered = + applyCometExecRule(FilterExec(LessThan(deferred.output.head, Literal(2)), deferred)) + assert(countOperators(filtered, classOf[CometProjectExec]) == 0, filtered.toString) + assert(countOperators(filtered, classOf[ProjectExec]) == 2, filtered.toString) + } } }