diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 0b0a69a1a75..458f5bb9865 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -618,6 +618,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 5d91aeb53fe..5a996a2d42f 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -322,6 +322,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 9bae4312078..688ae964d7b 100644 --- a/docs/source/user-guide/latest/compatibility/index.md +++ b/docs/source/user-guide/latest/compatibility/index.md @@ -121,6 +121,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 b5abfb88935..31729601503 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 eb05d9a7473..011f03f0a4e 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -215,6 +215,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 623c603af81..357059ebc3d 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -23,9 +23,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 @@ -48,7 +49,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._ @@ -90,6 +91,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. @@ -370,6 +374,9 @@ case class CometExecRule(session: SparkSession) // 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 @@ -931,9 +938,10 @@ case class CometExecRule(session: SparkSession) // 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) @@ -1046,6 +1054,167 @@ case class CometExecRule(session: SparkSession) } } + /** 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) + // 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 + case _ => node + } + + 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() => + 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 || + evaluatedExpressions(originalPlan(node)).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 | _: CometLocalTopKExec | _: 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(evaluatedExpressions(original)) + 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 + // 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. + 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) @@ -1502,6 +1671,8 @@ case class CometExecRule(session: SparkSession) 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 4a4b9eedbe7..6c2754a9fd5 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..554fbf6dbb3 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala @@ -0,0 +1,710 @@ +/* + * 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) == Tuple2(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 => + 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") + 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 == Tuple2(0, 2), initial.toString) + } + if (aqe && agg == "max(k)") { + assert(initialBuffers == Tuple2(2, 0), initial.toString) + } + val plan = df.queryExecution.executedPlan + if (aqe) { + assert(plan.asInstanceOf[AdaptiveSparkPlanExec].isFinalPlan) + 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 + } + 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("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) t LIMIT 2 OFFSET 1" + 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]), + ("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) +}