diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index f66a705afac..cf8baac011d 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -281,7 +281,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci | `hour` | ✅ | Native | | | `last_day` | ✅ | Native | | | `localtimestamp` | ✅ | — | | -| `make_date` | ✅ | Native | Computed arguments preserve year → month → day NULL checks. Expressions fall back when shared subexpressions or a parent could change which arguments Spark evaluates. IN probes stay native for literal/column candidates and array/struct constructors of those values; other computed candidates require Spark evaluation. | +| `make_date` | ✅ | Native | Computed arguments preserve year → month → day NULL checks. Expressions fall back when shared subexpressions or a parent could change which arguments Spark evaluates. IN probes stay native for literal/column candidates and array/struct constructors of those values; other computed candidates require Spark evaluation. Operator-level safeguards and their opt-out are described in [make_date evaluation](#make_date-evaluation). | | `make_dt_interval` | ✅ | Codegen dispatch | | | `make_interval` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; intervals outside Arrow's nanosecond range are tracked by [#5279](https://github.com/apache/datafusion-comet/issues/5279); the native path is opt-in via allowIncompatible ([details](compatibility/expressions/datetime.md)) | | `make_time` | ✅ | — | Spark 4.1+; requires `spark.sql.timeType.enabled=true`, which Spark leaves off by default. Runs natively; remaining TIME type work is tracked by [#4288](https://github.com/apache/datafusion-comet/issues/4288) | @@ -713,6 +713,59 @@ Comet also accelerates a number of Catalyst expressions that have no Spark SQL f This list is illustrative, not exhaustive: the per-function tables are not the complete set of expressions Comet can accelerate. +## make_date evaluation + +With ANSI mode enabled, `make_date(2024, 13, 1)` raises an invalid-date error. If an +input iterator yields a valid row followed by that invalid row, a `LIMIT 1` can stop +before evaluating the second date. A native projection can evaluate a whole batch +before the consuming operator decides which rows it needs. + +By default, Comet keeps affected operators in Spark to preserve skipped date evaluation. +This applies to date projections below `LIMIT`, `SAMPLE`, `FILTER`, joins, +`mapPartitions`, and `Dataset.rdd` consumers; date expressions inside joins, aggregates, +and ordered limits; and projections below a window group limit. A parent projection can +also defer a date output when an `IF`, `COALESCE`, or NULL-skipping operand does not need +it. Comet preserves that evaluation across projections while allowing eager date consumers, +such as `YEAR(dt)`, to remain native when their operands are safe to evaluate. An enclosing +conditional can still defer the extractor's date input. These safeguards also preserve +NULL short-circuiting when a date argument is computed in an earlier projection, including +across the streamed side of a broadcast join. The materialized build side remains independent. +These safeguards also cover required arguments, such as a computed year, on input rows that +the join rejects before producing an output row. +Selecting the date eagerly does not force conditionally used sibling expressions; those +computations also retain Spark evaluation when needed. + +With ANSI disabled, a direct date projection can remain native below these consumers when +all its computations are safe to evaluate early. This includes `make_date(2024, month, day)` +over input columns and years widened from `TINYINT` or `SMALLINT`. An unrestricted integer +year still needs the safeguard: Spark's conversion to Date32 can overflow independently of +ANSI mode. Computed arguments, their deferred producers, and throwing sibling expressions +retain their safeguards. Expression-level parent guards remain in effect. +For example, if `dt` is a nullable date alias, `unix_date(dt) + element_at(values, i)` +keeps its consuming operator in Spark when a NULL date can skip an invalid array index. +The safe date-producing projection can remain native. + +An unfiltered `COUNT(make_date(...))` can remain native when its sole argument is nullable, +because Spark checks that argument on every input row. Filtered counts and counts with +multiple arguments retain the safeguard when they can skip date evaluation. + +Only affected parts of the plan fall back. An ordinary native scan can remain native. +An adjacent collect limit and date projection can also stay native: Comet caps the projection's +input before evaluating dates, so unread rows cannot raise errors. This requires a supported +native local limit; otherwise the Spark safeguard remains. The cap does not bypass safeguards +for a lower argument-producing projection or intervening filters and joins. +Exchanges materialize their inputs and partition keys independently, so an outer limit does +not force an eager date partition key to Spark. Local sorts and aggregates can also separate +date evaluation from an outer limit. + +Users who accept eagerly evaluating dates, projected arguments, and sibling expressions can set +`spark.comet.exec.makeDate.operatorFallback.enabled=false` to allow native operators for +these shapes. The default is `true`. Disabling it can expose invalid-date or argument errors +that Spark would skip under ANSI mode. This setting does not change +`make_date`'s invalid-date handling or the expression-level safeguards for skipped +arguments and operands. For example, an empty `IN` predicate containing `make_date` +still falls back when needed. + ## See also - [Comet Compatibility Guide](compatibility/index.md) - known incompatibilities and edge cases for supported expressions. diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6003703a630..8af6cc02d6b 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -341,6 +341,17 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(true) + val COMET_EXEC_MAKE_DATE_OPERATOR_FALLBACK_ENABLED: ConfigEntry[Boolean] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.makeDate.operatorFallback.enabled") + .category(CATEGORY_EXEC) + .doc( + "Keep affected make_date operators in Spark when native batching can evaluate rows " + + "that Spark skips. When false, allow native operators even if make_date can raise " + + "an error on a row Spark would not evaluate. This does not disable expression-level " + + "evaluation safeguards or change make_date's invalid-date handling.") + .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 d79927f0c04..9616b7d492d 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -445,6 +445,9 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } def convertNode(op: SparkPlan): SparkPlan = op match { + case op if MakeDateEvaluation.reason(op).isDefined => + withFallbackReason(op, MakeDateEvaluation.reason(op).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 @@ -751,6 +754,25 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) */ private def limitCollectProjectionInput(op: SparkPlan): SparkPlan = op match { case limit: CollectLimitExec if limit.limit >= 0 => + def capProjection(project: SparkPlan, original: SparkPlan): Option[SparkPlan] = { + val input = project.children.head + val cap = input match { + case bounded: CometLocalLimitExec if bounded.limit <= limit.limit => Some(bounded) + case _ => convertToComet(LocalLimitExec(limit.limit, input), CometLocalLimitExec) + } + for { + bounded <- cap + rebuilt <- convertToComet(original.withNewChildren(Seq(bounded)), CometProjectExec) + } yield { + // AQE can link the current projection to a materialized logical query stage. + // Its originalPlan may have an older link; preserve the current direct/inherited tags. + Seq(SparkPlan.LOGICAL_PLAN_TAG, SparkPlan.LOGICAL_PLAN_INHERITED_TAG).foreach { tag => + project.getTagValue(tag).foreach(rebuilt.setTagValue(tag, _)) + } + rebuilt + } + } + limit.child match { case project: CometProjectExec => // Selecting or renaming columns cannot throw. Keep Project -> Filter adjacent so @@ -766,29 +788,18 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) if (onlyColumns || alreadyCapped) { op } else { - // Use normal config and serde gating for the synthetic limit. If conversion fails, - // its fallback tags are discarded with it and the original plan is retained. - val bounded = for { - cap <- convertToComet( - LocalLimitExec(limit.limit, project.child), - CometLocalLimitExec) - rebuilt <- convertToComet( - project.originalPlan.withNewChildren(Seq(cap)), - CometProjectExec) - } yield { - // AQE can link the current projection to a materialized logical query stage. - // Its originalPlan may have an older link; serde also adds tags, so copyTagsFrom - // (which only copies into an empty tag map) would not preserve the current link. - project.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).foreach { link => - rebuilt.setTagValue(SparkPlan.LOGICAL_PLAN_TAG, link) - } - project.getTagValue(SparkPlan.LOGICAL_PLAN_INHERITED_TAG).foreach { link => - rebuilt.setTagValue(SparkPlan.LOGICAL_PLAN_INHERITED_TAG, link) - } - limit.withNewChildren(Seq(rebuilt)) - } - bounded.getOrElse(op) + capProjection(project, project.originalPlan) + .map(rebuilt => limit.withNewChildren(Seq(rebuilt))) + .getOrElse(op) } + case project: ProjectExec => + // The date guard stays in place until both the cap and projection convert. This + // candidate excludes other lazy consumers and guarded producers below the projection. + MakeDateEvaluation + .collectLimitProjection(project, limit.limit) + .flatMap(candidate => capProjection(candidate, candidate)) + .map(rebuilt => MakeDateEvaluation.withCappedProjection(limit, rebuilt)) + .getOrElse(op) case _ => op } case _ => op @@ -1041,15 +1052,16 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) plan } } else { + val planWithDateEvaluation = MakeDateEvaluation.preserve(plan) // The rewrite removes the sorts of each join it converts, so put back any sort a kept // operator above still needs. This runs before transform(), which converts the sorts it // adds like any other SortExec. val planWithJoinRewritten = if (CometConf.COMET_FORCE_SHJ.get()) { - RewriteJoin.restoreRequiredOrdering(plan.transformUp { case p => - RewriteJoin.rewrite(p, conf) + RewriteJoin.restoreRequiredOrdering(planWithDateEvaluation.transformUp { case p => + if (MakeDateEvaluation.reason(p).isDefined) p else RewriteJoin.rewrite(p, conf) }) } else { - plan + planWithDateEvaluation } // Tag Partial aggregates that must not be converted to Comet because a diff --git a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala new file mode 100644 index 00000000000..3a6f0d6c196 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala @@ -0,0 +1,477 @@ +/* + * 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.rules + +import java.time.LocalDate + +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, And, Attribute, AttributeSet, BinaryComparison, Cast, ConditionalExpression, CreateArray, CreateNamedStruct, DateAdd, DateDiff, DateFormatClass, DateSub, EqualNullSafe, Expression, GetDateField, Greatest, HashExpression, IsNotNull, IsNull, Least, Literal, MakeDate, Multiply, Not, Or, SortOrder, Subtract, UnixDate} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Count} +import org.apache.spark.sql.catalyst.optimizer.BuildLeft +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.comet._ +import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec, QueryStageExec} +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} +import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec} +import org.apache.spark.sql.types.{ByteType, IntegerType, ShortType} + +import org.apache.comet.{CometConf, CometExplainInfo} +import org.apache.comet.serde.CometMakeDate +import org.apache.comet.shims.ShimCometWindowGroupLimit + +/** Keeps make_date on Spark's row path where native batching can change whether it throws. */ +private[rules] object MakeDateEvaluation { + private val sparkEvaluation = TreeNodeTag[String]("comet.makeDateSparkEvaluation") + private val collectProjectionLimit = TreeNodeTag[Int]("comet.makeDateCollectProjectionLimit") + private val collectLimitReason = "make_date requires Spark evaluation below LIMIT" + + // Return a copy: failed cap conversion must leave the original guarded pipeline untouched. + def collectLimitProjection(project: ProjectExec, limit: Int): Option[ProjectExec] = + if (enabled(project) && project.getTagValue(collectProjectionLimit).contains(limit)) { + val candidate = project.copy() + candidate.copyTagsFrom(project) + clearCollectLimitReason(candidate) + Some(candidate) + } else None + + def withCappedProjection(limit: CollectLimitExec, project: SparkPlan): CollectLimitExec = { + val capped = limit.copy(child = project) + capped.copyTagsFrom(limit) + clearCollectLimitReason(capped) + capped + } + + private def clearCollectLimitReason(plan: SparkPlan): Unit = { + if (plan.getTagValue(sparkEvaluation).contains(collectLimitReason)) { + plan.unsetTagValue(sparkEvaluation) + } + plan.unsetTagValue(collectProjectionLimit) + plan.getTagValue(CometExplainInfo.FALLBACK_REASONS).foreach { reasons => + val remaining = reasons - collectLimitReason + if (remaining.isEmpty) plan.unsetTagValue(CometExplainInfo.FALLBACK_REASONS) + else plan.setTagValue(CometExplainInfo.FALLBACK_REASONS, remaining) + } + } + + private def enabled(plan: SparkPlan): Boolean = + CometConf.COMET_EXEC_MAKE_DATE_OPERATOR_FALLBACK_ENABLED.get(plan.conf) + + def reason(plan: SparkPlan): Option[String] = + if (enabled(plan)) plan.getTagValue(sparkEvaluation) else None + + private def containsMakeDate(expressions: Seq[Expression]): Boolean = + expressions.exists(_.exists(_.isInstanceOf[MakeDate])) + + // Spark's epoch-day conversion can overflow even with ANSI disabled. Exclude the two + // partially representable boundary years, so every valid month/day fits in Date32. + private val firstCompleteYear = LocalDate.ofEpochDay(Int.MinValue.toLong).getYear + 1 + private val lastCompleteYear = LocalDate.ofEpochDay(Int.MaxValue.toLong).getYear - 1 + + private def isSmallIntegerWidening(cast: Cast): Boolean = + cast.dataType == IntegerType && + (cast.child.dataType == ByteType || cast.child.dataType == ShortType) + + // This only describes computations over their input values. An attribute may still name a + // deferred producer, so argument lineage and sibling provenance must be tracked separately. + private def canEvaluateProjectionEagerly(expression: Expression): Boolean = expression match { + case Alias(child, _) => canEvaluateProjectionEagerly(child) + case date: MakeDate if !date.failOnError => + val boundedYear = date.year match { + case Literal(year: Int, IntegerType) => + year >= firstCompleteYear && year <= lastCompleteYear + case cast: Cast if isSmallIntegerWidening(cast) => true + case _ => false + } + boundedYear && date.children.forall { + case cast: Cast if isSmallIntegerWidening(cast) => + CometMakeDate.canEvaluateEagerly(cast.child) + case child => CometMakeDate.canEvaluateEagerly(child) + } + case _ => CometMakeDate.canEvaluateEagerly(expression) + } + + private def containsDeferredAggregateDate(expression: Expression): Boolean = expression match { + // COUNT checks its sole nullable argument on every input row. Filters, other nullable + // arguments, and a non-nullable argument can let Spark skip evaluating the date. + case AggregateExpression(Count(Seq(child)), _, _, None, _) if child.nullable => false + case _: MakeDate => true + case _ => expression.children.exists(containsDeferredAggregateDate) + } + + private case class ProjectionDemand( + boundary: String, + deferred: AttributeSet, + // Deferred inputs whose producers need guarding even when they contain no make_date. + guardedInputs: AttributeSet, + // Date argument lineage retained until a streamed join can skip its input rows. + dateInputs: AttributeSet) + + // Attributes may stand for deferred expressions in a lower projection, not materialized values. + private def harmlessValue(expression: Expression): Boolean = + expression.references.isEmpty && CometMakeDate.canEvaluateEagerly(expression) + + /** References that a consumer requires without exposing a potentially skipped sibling. */ + private def eagerReferences(expression: Expression): AttributeSet = expression match { + case attribute: Attribute => AttributeSet(Seq(attribute)) + case _ => + val children = expression match { + case date: MakeDate => + val firstNullable = date.children.indexWhere(_.nullable) + date.children.take(if (firstNullable < 0) date.children.size else firstNullable + 1) + case conditional: ConditionalExpression => + val required = conditional.alwaysEvaluatedInputs + if (conditional.children.filterNot(required.contains).forall(harmlessValue)) required + else Nil + case _: Alias | _: Cast | _: IsNull | _: IsNotNull | _: Not | _: UnixDate | + _: GetDateField => + expression.children + case _: EqualNullSafe | _: CreateArray | _: CreateNamedStruct | _: Greatest | _: Least | + _: HashExpression[_] => + expression.children + case _: And | _: Or => + if (harmlessValue(expression.children(1))) expression.children.take(1) else Nil + case _: Add | _: Subtract | _: Multiply | _: BinaryComparison | _: DateAdd | _: DateSub | + _: DateDiff | _: DateFormatClass => + val left = expression.children.head + val right = expression.children(1) + if (!left.nullable) expression.children + // A right-hand attribute can inline a throwing expression from an earlier projection. + // Only self-contained harmless values are safe without tracing that attribute's source. + else if (harmlessValue(right)) Seq(left) + else Nil + // Unknown expressions may skip inputs or catch errors raised by their children. + case _ => Nil + } + AttributeSet(children.flatMap(eagerReferences)) + } + + private def copyLogicalLinks(from: SparkPlan, to: SparkPlan): Unit = { + if (from ne to) { + // setLogicalLink promotes inherited links to direct links and rewrites descendants. + // Keep each tag's role and retain the fallback plan's links when the source has none. + Seq(SparkPlan.LOGICAL_PLAN_TAG, SparkPlan.LOGICAL_PLAN_INHERITED_TAG).foreach { tag => + from.getTagValue(tag).foreach(to.setTagValue(tag, _)) + } + } + } + + private def nativeBroadcastInput(plan: SparkPlan): Option[SparkPlan] = plan match { + case _: CometBroadcastExchangeExec => Some(plan) + case stage: BroadcastQueryStageExec => nativeBroadcastInput(stage.plan).map(_ => stage) + case reuse: ReusedExchangeExec => nativeBroadcastInput(reuse.child).map(_ => reuse) + case sink: CometSinkPlaceHolder => nativeBroadcastInput(sink.child) + case wrapper: CometScanWrapper => nativeBroadcastInput(wrapper.originalPlan) + case _ => None + } + + private def restoreBroadcast(plan: SparkPlan, reason: String): SparkPlan = + nativeBroadcastInput(plan) + .map { input => + // Spark joins need a HashedRelation (or rows for a nested-loop join), not Arrow buffers. + // The existing transition decodes the cached broadcast and applies its Spark broadcast + // mode, including through AQE/reuse wrappers, without re-executing the build input. + val restored = CometColumnarToRowExec(input) + restored.setTagValue(sparkEvaluation, reason) + copyLogicalLinks(plan, restored) + restored + } + .getOrElse(plan) + + /** + * Run before bottom-up conversion and join rewriting. Reasons propagate back to the consuming + * operator, keeping the row pipeline together instead of inserting an eager Arrow batch between + * a throwing projection and its lazy consumer. Exchanges materialize their input independently; + * local sorts and aggregates do so only after their consumer requests input. + */ + def preserve(plan: SparkPlan): SparkPlan = { + if (!enabled(plan)) return plan + + def visit( + node: SparkPlan, + rowDemand: Option[String], + projectionDemand: Option[ProjectionDemand], + directCollectLimit: Option[Int] = None): (SparkPlan, Option[String], AttributeSet) = { + // Candidate tags describe this walk's parent context, not a previous AQE plan. + node.unsetTagValue(collectProjectionLimit) + node match { + // Materialized stages are independent producers. Their plans have already been guarded. + case _: QueryStageExec | _: ReusedExchangeExec => return (node, None, AttributeSet.empty) + // These wrappers do not have the child shape of their original Spark operator. + case wrapper: CometScanWrapper => + val (child, childReason, dateOutputs) = + visit(wrapper.originalPlan, rowDemand, projectionDemand, directCollectLimit) + return if (childReason.isDefined) (child, childReason, dateOutputs) + else (wrapper.copy(originalPlan = child), None, dateOutputs) + case sink: CometSinkPlaceHolder => + val (child, childReason, dateOutputs) = + visit(sink.child, rowDemand, projectionDemand, directCollectLimit) + return if (childReason.isDefined) (child, childReason, dateOutputs) + else (sink.withNewChildren(Seq(child)), None, dateOutputs) + case _ => + } + + val original = node match { + case scan: CometScanExec => + scan.wrapped.copy( + partitionFilters = scan.partitionFilters, + dataFilters = scan.dataFilters) + case scan: CometNativeScanExec => scan.sparkFallback(Nil) + case native: CometExec => native.originalPlan + case shuffle: CometShuffleExchangeExec => shuffle.originalPlan + case broadcast: CometBroadcastExchangeExec => broadcast.originalPlan + case _ => node + } + val evaluatesDate = containsMakeDate(original.expressions) + val exchange = original.isInstanceOf[Exchange] + val join = original.isInstanceOf[BaseJoinExec] + // Broadcast joins can inline streamed expressions into their consumer. Their build side + // materializes rows; shuffled hash and sort-merge joins also isolate the streamed input. + val streamedChildIndex = original match { + case hash: BroadcastHashJoinExec if hash.supportCodegen => + Some(if (hash.buildSide == BuildLeft) 1 else 0) + case nestedLoop: BroadcastNestedLoopJoinExec if nestedLoop.supportCodegen => + Some(if (nestedLoop.buildSide == BuildLeft) 1 else 0) + case _ => None + } + val localMaterialization = original match { + case _: SortExec | _: HashAggregateExec | _: ObjectHashAggregateExec => true + case _: TakeOrderedAndProjectExec => true + case _ => false + } + val startsLimit = original match { + case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec => true + case topK: TakeOrderedAndProjectExec => + SortOrder.orderingSatisfies(topK.child.outputOrdering, topK.sortOrder) + case _ + if ShimCometWindowGroupLimit.windowGroupLimitClass.exists(_.isInstance(original)) => + true + case _ => false + } + val childDemand = original match { + case _ if exchange => None + // A join can leave a local sort or aggregate completely unpolled. + case _ if join || rowDemand.contains("JOIN") => Some("JOIN") + case _ if startsLimit => Some("LIMIT") + case _: SampleExec => Some("SAMPLE") + case _: MapPartitionsExec => Some("a mapPartitions function") + case _ if localMaterialization => None + case _ => rowDemand + } + val childProjectionDemand = original match { + // Generated filters can defer deterministic projection outputs until a row is accepted. + case filter: FilterExec => + Some( + ProjectionDemand( + "FILTER", + filter.child.outputSet, + projectionDemand.map(_.guardedInputs).getOrElse(AttributeSet.empty), + projectionDemand.map(_.dateInputs).getOrElse(AttributeSet.empty))) + case project: ProjectExec => + val deferred = projectionDemand.map(_.deferred).getOrElse(AttributeSet.empty) + val eager = project.projectList + .filter(output => !output.deterministic || !deferred.contains(output.toAttribute)) + .flatMap(eagerReferences) + // ProjectExec forces repeated inputs and nondeterministic outputs. Other inputs can + // remain inlined in a consumer's conditional, even across intervening aliases. + val skipped = project.child.outputSet -- project.usedInputs -- AttributeSet(eager) + val inheritedInputs = + projectionDemand.map(_.guardedInputs).getOrElse(AttributeSet.empty) + // A consumer can require the date but still skip another output of this projection. + val guardedOutputs = inheritedInputs ++ + (if (evaluatesDate) deferred else AttributeSet.empty) + val directDateInputs = AttributeSet(project.projectList.flatMap { output => + output.collect { case date: MakeDate => date }.flatMap(_.references) + }) + val inputsToGuard = + directDateInputs ++ AttributeSet(project.projectList.flatMap { output => + if ((output.deterministic || rowDemand.isDefined) && + guardedOutputs.contains(output.toAttribute)) { + output.references.toSeq + } else { + Nil + } + }) + // Keep argument lineage through eager aliases too. A join below this projection can + // discard an input row before this projection ever requires the argument's value. + val dateInputs = + directDateInputs ++ AttributeSet(project.projectList.flatMap { output => + if (projectionDemand.exists(_.dateInputs.contains(output.toAttribute))) { + output.references.toSeq + } else { + Nil + } + }) + // A lazy row consumer can skip even the date's required arguments on unread rows. + val guardedChildInputs = + inputsToGuard.intersect(if (rowDemand.isDefined) project.child.outputSet else skipped) + if (skipped.isEmpty && guardedChildInputs.isEmpty && dateInputs.isEmpty) { + None + } else { + Some( + ProjectionDemand( + projectionDemand.map(_.boundary).getOrElse("PROJECT"), + skipped, + guardedChildInputs, + dateInputs)) + } + case _ if exchange || localMaterialization => None + case _ => projectionDemand + } + val childCollectLimit = original match { + case limit: CollectLimitExec + if limit.limit >= 0 && rowDemand.isEmpty && projectionDemand.isEmpty => + Some(limit.limit) + case _ => None + } + val visited = node.children.zipWithIndex.map { case (child, index) => + val demand = if (!join) { + childProjectionDemand + } else if (streamedChildIndex.contains(index)) { + childProjectionDemand.map { demand => + // Keys and conditions force their inputs on each visited row. An outer row consumer + // can still skip that row altogether, so retain its guarded inputs in that case. + val required = if (rowDemand.isDefined) AttributeSet.empty else original.references + val streamedOutputs = original.children(index).outputSet + val forwarded = streamedOutputs -- required + demand.copy( + deferred = demand.deferred.intersect(forwarded), + guardedInputs = (demand.guardedInputs ++ demand.dateInputs).intersect(forwarded), + // A key is required only on this join's input rows. An earlier join can still + // discard rows before producing that key, so retain its argument lineage. + dateInputs = demand.dateInputs.intersect(streamedOutputs)) + } + } else { + None + } + visit(child, childDemand, demand, childCollectLimit) + } + val childReason = visited.flatMap(_._2).headOption + val childDateOutputs = AttributeSet(visited.zipWithIndex.flatMap { + case ((_, _, outputs), index) if !join || streamedChildIndex.contains(index) => outputs + case _ => AttributeSet.empty + }) + // Date provenance flows up through aliases while deferred evaluation flows down. This + // also finds a skipped sibling computed in a projection above the date's producer. + val dateDerivedOutputs = if (exchange || localMaterialization) { + AttributeSet.empty + } else { + original match { + case project: ProjectExec => + AttributeSet( + project.projectList + .filter { output => + containsMakeDate(Seq(output)) || output.references + .intersect(childDateOutputs) + .nonEmpty + } + .map(_.toAttribute)) + case _ => childDateOutputs.intersect(original.outputSet) + } + } + // An exchange evaluates its partition keys for every input before a consumer reads rows. + val safeProjection = original match { + case project: ProjectExec => project.projectList.forall(canEvaluateProjectionEagerly) + case _ => false + } + val rowReason = if (evaluatesDate && !exchange && !safeProjection) { + rowDemand.map(boundary => s"make_date requires Spark evaluation below $boundary") + } else None + // A safe date producer can remain native, but its NULL output can still hide a throwing + // arithmetic operand in the consumer. The expression serde only sees an attribute here. + val maskedDateSibling = childDateOutputs.nonEmpty && original.expressions.exists(_.exists { + case arithmetic @ (_: Add | _: Subtract | _: Multiply) => + val left = arithmetic.children.head + val right = arithmetic.children(1) + left.nullable && left.references.intersect(childDateOutputs).nonEmpty && + !CometMakeDate.canEvaluateEagerly(right) + case _ => false + }) + val operatorReason = original match { + case _ if maskedDateSibling => + Some(s"make_date requires Spark evaluation in ${original.nodeName}") + case project: ProjectExec => + projectionDemand + .filter(demand => + project.projectList.exists { output => + val computation = output match { + case alias: Alias => alias.child + case other => other + } + val deferredDate = output.deterministic && + demand.deferred.contains(output.toAttribute) && containsMakeDate(Seq(output)) && + !canEvaluateProjectionEagerly(output) + val deferredComputation = (output.deterministic || rowDemand.isDefined) && + (demand.guardedInputs.contains(output.toAttribute) || + (dateDerivedOutputs.nonEmpty && demand.deferred.contains( + output.toAttribute))) && + !canEvaluateProjectionEagerly(computation) + deferredDate || deferredComputation + }) + .map(demand => s"make_date requires Spark evaluation below ${demand.boundary}") + case aggregate: BaseAggregateExec + if aggregate.expressions.exists(containsDeferredAggregateDate) => + Some(s"make_date requires Spark evaluation in ${original.nodeName}") + // Join keys/conditions and top-k outputs can be skipped inside the operator. + case _: BaseJoinExec | _: TakeOrderedAndProjectExec if evaluatesDate => + Some(s"make_date requires Spark evaluation in ${original.nodeName}") + case _ => None + } + val ownReason = reason(node).orElse(rowReason).orElse(operatorReason) + val fallbackReason = ownReason.orElse(if (exchange) None else childReason) + val children = visited.map { case (child, _, _) => + if (join && fallbackReason.isDefined) restoreBroadcast(child, fallbackReason.get) + else child + } + val rebuilt = node match { + case _: RowToColumnarExec | _: CometSparkToColumnarExec if childReason.isDefined => + children.head + case _: ColumnarToRowExec | _: CometColumnarToRowExec | _: CometNativeColumnarToRowExec + if childReason.isDefined && !children.head.supportsColumnar => + children.head + case native: CometExec if fallbackReason.isDefined || children != node.children => + native.sparkFallback(children) + case _ if (original ne node) && (fallbackReason.isDefined || children != node.children) => + original.withNewChildren(children) + case _ => node.withNewChildren(children) + } + copyLogicalLinks(node, rebuilt) + rebuilt.unsetTagValue(collectProjectionLimit) + if (original.isInstanceOf[ProjectExec] && rowReason.contains(collectLimitReason) && + operatorReason.isEmpty && childReason.isEmpty && + reason(node).forall(_ == collectLimitReason)) { + directCollectLimit.foreach(rebuilt.setTagValue(collectProjectionLimit, _)) + } + fallbackReason.foreach(rebuilt.setTagValue(sparkEvaluation, _)) + (rebuilt, fallbackReason, dateDerivedOutputs) + } + + // Dataset.rdd adds an object-producing root, and its caller may read only a few objects. + // EliminateSerialization can leave a typed operator, filter, or project at that root. + def producesObjects(op: SparkPlan): Boolean = op match { + case _: ObjectProducerExec => true + case _: FilterExec | _: ProjectExec => producesObjects(op.children.head) + case _ => false + } + val rootDemand = if (producesObjects(plan)) Some("code reading Dataset.rdd") else None + visit(plan, rootDemand, None)._1 + } +} 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 5a867ec8387..8f7ca482209 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -456,7 +456,7 @@ object CometMakeDate extends CometExpressionSerde[MakeDate] { // Native parents can evaluate operands for the whole batch. These expressions are safe // on rows where Spark skips them after a NULL operand or an IN match. - private def canEvaluateEagerly(expr: Expression): Boolean = expr match { + private[comet] def canEvaluateEagerly(expr: Expression): Boolean = expr match { case _: Literal | _: Attribute => true case _: CreateArray | _: CreateNamedStruct => expr.children.forall(canEvaluateEagerly) diff --git a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala index e61a57a6418..c9e97bbfddd 100644 --- a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala @@ -26,7 +26,10 @@ import org.apache.spark.sql.{CometTestBase, DataFrame, Row, SaveMode} import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute import org.apache.spark.sql.catalyst.expressions.{Days, Hours, Literal} import org.apache.spark.sql.catalyst.util.DateTimeUtils +import org.apache.spark.sql.comet.{CometLocalLimitExec, CometProjectExec} +import org.apache.spark.sql.execution.ProjectExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec import org.apache.spark.sql.functions.col import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataTypes, StructField, StructType} @@ -56,6 +59,122 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH .getOrElse( fail(s"No SparkThrowable in cause chain: ${causeChain(error).map(_.getClass.getName)}")) + test("make_date does not evaluate an invalid row beyond a partial consumer") { + import testImplicits._ + + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> + "org.apache.spark.sql.catalyst.optimizer.PushProjectionThroughLimit", + CometConf.COMET_BATCH_SIZE.key -> "8192") { + withTempPath { path => + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark + .range(2) + .coalesce(1) + .selectExpr( + "CAST(CASE WHEN id = 0 THEN 1 ELSE 13 END AS INT) AS m", + "CASE WHEN id = 0 THEN '2024' ELSE 'bad' END AS year_text") + .write + .parquet(path.getCanonicalPath) + } + withTempView("make_date_limit") { + spark.read.parquet(path.getCanonicalPath).createOrReplaceTempView("make_date_limit") + for (codegen <- Seq(false, true)) { + withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegen.toString) { + val query = "SELECT unix_date(make_date(2024, m, 1)) FROM make_date_limit LIMIT 1" + var expected = Seq.empty[Row] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + expected = sql(query).collect().toSeq + } + assert(expected == Seq(Row(19723))) + val df = sql(query) + val plan = stripAQEPlan(df.queryExecution.executedPlan) + assert(plan.collect { case p: ProjectExec => p }.isEmpty, plan.treeString) + assert(plan.collect { case p: CometProjectExec => p }.size == 1, plan.treeString) + assert(plan.collect { case p: CometLocalLimitExec => p }.size == 1, plan.treeString) + assert( + !new ExtendedExplainInfo() + .getFallbackReasons(plan) + .exists(_.contains("make_date")), + plan.treeString) + assert(df.collect().toSeq == expected) + + val unlimited = query.stripSuffix(" LIMIT 1") + val partialConsumers: Seq[(String, DataFrame => Seq[Row])] = Seq( + ("RDD.take", df => df.rdd.take(1).toSeq), + ( + "mapPartitions", + df => df.as[Int].mapPartitions(_.take(1)).collect().toSeq.map(Row(_)))) + partialConsumers.foreach { case (name, consume) => + withClue(s"$name codegen=$codegen: ") { + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + expected = consume(sql(unlimited)) + } + assert(expected == Seq(Row(19723))) + assert(consume(sql(unlimited)) == expected) + } + } + + // LIMIT can also skip an eager date argument produced by an earlier projection. + val yearQuery = """ + |SELECT r, unix_date(make_date(y, 1, 1)), r + 1 FROM ( + | SELECT rand(7L) AS r, CAST(year_text AS INT) AS y FROM make_date_limit + |) LIMIT 1 + |""".stripMargin + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + expected = sql(yearQuery).collect().toSeq + } + assert(expected.size == 1 && expected.head.getInt(1) == 19723) + val yearDf = sql(yearQuery) + assert(yearDf.collect().toSeq == expected) + val yearPlan = stripAQEPlan(yearDf.queryExecution.executedPlan) + assert(yearPlan.collect { case p: ProjectExec => p }.size == 2, yearPlan.treeString) + assert( + yearPlan.collect { case p: CometProjectExec => p }.isEmpty, + yearPlan.treeString) + + // A join key needs the cast on visited rows, but LIMIT still skips later rows. + val joinedYearQuery = """ + |SELECT /*+ BROADCAST(b) */ p.r, unix_date(make_date(p.y, 1, 1)), p.r + 1 FROM ( + | SELECT rand(7L) AS r, CAST(year_text AS INT) AS y FROM make_date_limit + |) p JOIN (SELECT 2024 AS k UNION ALL SELECT 2025 AS k) b ON p.y = b.k + |LIMIT 1 + |""".stripMargin + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + expected = sql(joinedYearQuery).collect().toSeq + } + assert(expected.size == 1 && expected.head.getInt(1) == 19723) + val joinedYearDf = sql(joinedYearQuery) + assert(joinedYearDf.collect().toSeq == expected) + val joinedYearPlan = stripAQEPlan(joinedYearDf.queryExecution.executedPlan) + assert( + joinedYearPlan.collect { + case p: ProjectExec if p.output.exists(_.name == "r") => p + }.size == 2, + joinedYearPlan.treeString) + assert( + joinedYearPlan.collect { case join: BroadcastHashJoinExec => join }.size == 1, + joinedYearPlan.treeString) + + // A full collection must still evaluate the second row in the native expression. + val full = sql(unlimited) + checkCometOperators(stripAQEPlan(full.queryExecution.executedPlan)) + val (sparkError, cometError) = checkSparkAnswerMaybeThrows(full) + Seq("Spark" -> sparkError, "Comet" -> cometError).foreach { case (engine, error) => + withClue(s"$engine codegen=$codegen: ") { + assert( + error.exists(e => causeChain(e).exists(_.toString.contains("MonthOfYear")))) + } + } + } + } + } + } + } + } + test("next_day and make_date ANSI errors match Spark exceptions") { withSQLConf( SQLConf.ANSI_ENABLED.key -> "true", 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 e67ec19c7b3..ffbe1d0e1b0 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -29,12 +29,13 @@ import org.scalatest.PrivateMethodTester._ import org.apache.logging.log4j.Level import org.apache.spark.sql._ import org.apache.spark.sql.catalyst.{FunctionIdentifier, InternalRow} -import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, Descending, Expression, ExpressionInfo, GreaterThan, In, InSet, KnownFloatingPointNormalized, Literal, Not, RuntimeReplaceable, SortOrder} -import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, BloomFilterAggregate, Final, Min, Partial, PartialMerge} +import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, Attribute, AttributeReference, BoundReference, Descending, Expression, ExpressionInfo, GreaterThan, In, InSet, KnownFloatingPointNormalized, Literal, MakeDate, Not, RuntimeReplaceable, SortOrder} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, BloomFilterAggregate, Count, Final, Min, Partial, PartialMerge} import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke} -import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, NormalizeNaNAndZero} +import org.apache.spark.sql.catalyst.plans.Inner import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan} -import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, IdentityBroadcastMode, RoundRobinPartitioning} +import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, IdentityBroadcastMode, RoundRobinPartitioning, SinglePartition} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.CometExec.nativeLeaves @@ -43,13 +44,16 @@ import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, LogicalQueryStage, QueryStageExec, ShuffleQueryStageExec, SimpleCost, SimpleCostEvaluator} import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.datasources.v2.BatchScanExec -import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec} +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec, ShuffleExchangeExec} +import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, HashedRelationBroadcastMode, SortMergeJoinExec} +import org.apache.spark.sql.functions.expr import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataTypes, DoubleType, FloatType, StructField, StructType} import org.apache.comet.{CometConf, CometCoverageStats, CometExplainInfo, CometSparkSessionExtensions, ConfigEntry, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark42Plus, withFallbackReason} import org.apache.comet.serde.{CometAggregateExpressionSerde, Compatible, ExprOuterClass, OperatorOuterClass, QueryPlanSerde, Unsupported} +import org.apache.comet.shims.ShimCometWindowGroupLimit import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator} /** @@ -151,6 +155,1011 @@ class CometExecRuleSuite extends CometTestBase { }.sum } + test("make_date preserves row demand below lazy consumers, including reused native plans") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val input = createSparkPlan(spark, "SELECT * FROM VALUES (2024), (2025) AS t(y)") + val sample = createSparkPlan( + spark, + "SELECT * FROM (SELECT * FROM range(8)) TABLESAMPLE (50 PERCENT)").collectFirst { + case op: SampleExec => op + }.get + val consumers: Seq[(String, SparkPlan => SparkPlan)] = Seq( + ("LIMIT", child => CollectLimitExec(1, child)), + ("LIMIT", child => LocalLimitExec(1, child)), + ("LIMIT", child => GlobalLimitExec(1, child)), + ("SAMPLE", child => sample.withNewChildren(Seq(child))), + ("FILTER", child => FilterExec(GreaterThan(child.output.head, Literal(0)), child))) + for (ansi <- Seq(false, true); boundedYear <- Seq(false, true)) { + def project: SparkPlan = ProjectExec( + Seq( + input.output.head, + Alias( + MakeDate( + if (boundedYear) Literal(2024) else input.output.head, + Literal(1), + Literal(1), + ansi), + "date")()), + input) + val native = applyCometExecRule(project) + assert(countOperators(native, classOf[CometProjectExec]) == 1, native.treeString) + for ((boundary, consume) <- consumers; reused <- Seq(false, true)) { + val child = if (reused) applyCometExecRule(project) else project + val consumer = consume(child) + val converted = applyCometExecRule(consumer) + for (result <- Seq(converted, applyCometExecRule(converted))) { + // A reused local-table sink has no native child operator for a new cap. + val capped = consumer.isInstanceOf[CollectLimitExec] && !reused + val fallback = (ansi || !boundedYear) && !capped + withClue( + s"ansi=$ansi boundedYear=$boundedYear reused=$reused $boundary\n${result.treeString}") { + assert(countOperators(result, classOf[ProjectExec]) == (if (fallback) 1 else 0)) + assert( + countOperators(result, classOf[CometProjectExec]) == (if (fallback) 0 else 1)) + assert(new ExtendedExplainInfo() + .getFallbackReasons(result) + .exists( + _.contains(s"make_date requires Spark evaluation below $boundary")) == fallback) + } + } + } + + // These consumers evaluate their complete input independently of the outer limit. + val materializers: Seq[SparkPlan => SparkPlan] = Seq( + child => SortExec(Seq(SortOrder(child.output.head, Ascending)), false, child), + child => ShuffleExchangeExec(SinglePartition, child)) + materializers.foreach { materialize => + val result = applyCometExecRule(LocalLimitExec(1, materialize(project))) + assert(countOperators(result, classOf[CometProjectExec]) == 1, result.treeString) + assert( + !new ExtendedExplainInfo() + .getFallbackReasons(result) + .exists(_.contains("make_date requires Spark evaluation below LIMIT"))) + } + } + } + } + + test("make_date ANSI-off projections retain guards for throwing inputs and siblings") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") { + val rows = Seq((2024.toShort, Some(1), 1), (2024.toShort, Some(13), 0)) + withTempPath { path => + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.createDataFrame(rows).coalesce(1).write.parquet(path.getCanonicalPath) + } + withParquetTable(path.getCanonicalPath, "make_date_non_ansi") { + // A literal or widened small-integer year cannot overflow the Date32 representation. + // The invalid month is harmless with ANSI off, including below a lazy row consumer. + for (year <- Seq("2024", "_1")) { + val df = sql(s"SELECT make_date($year, _2, 1) FROM make_date_non_ansi LIMIT 10") + checkSparkAnswerAndOperator(df, includeClasses = Seq(classOf[CometProjectExec])) + } + + val deferred = sql(""" + |SELECT r, IF(r < 0, dt, DATE '2024-01-01'), r + 1 FROM ( + | SELECT rand(7L) AS r, make_date(2024, _2, 1) AS dt FROM make_date_non_ansi + |) + |""".stripMargin) + checkSparkAnswerAndOperator(deferred, includeClasses = Seq(classOf[CometProjectExec])) + assert( + countOperators(deferred.queryExecution.executedPlan, classOf[CometProjectExec]) == 2) + + // A harmless date producer can still yield NULL and hide a throwing operand in + // its consumer. Repeated r keeps the date and arithmetic in separate projections. + val consumers = Seq( + ("unix_date(dt) + element_at(array(1), i)", true), + ("unix_date(dt) - element_at(array(1), i)", true), + ("unix_date(dt) * element_at(array(1), i)", true), + ("unix_date(dt) + 1", false), + ("unix_date(dt) + i", false), + ("i + unix_date(dt)", false)) + for ((consumer, fallback) <- consumers) { + val df = sql(s""" + |SELECT r, r + 1, $consumer AS result FROM ( + | SELECT rand(7L) AS r, make_date(2024, _2, 1) AS dt, _3 AS i + | FROM make_date_non_ansi + |) + |""".stripMargin) + checkSparkAnswer(df) + val plan = df.queryExecution.executedPlan + withClue(s"consumer=$consumer\n${plan.treeString}") { + assert(countOperators(plan, classOf[ProjectExec]) == (if (fallback) 1 else 0)) + assert(countOperators(plan, classOf[CometProjectExec]) == (if (fallback) 1 else 2)) + assert(countOperators(plan, classOf[CometNativeScanExec]) == 1) + } + } + + // Index zero throws even with ANSI off. A limit can skip a throwing date argument, + // and a NULL date can skip an unrelated right operand in Spark's arithmetic. + val guarded = Seq( + "SELECT make_date(2024, 1, element_at(array(1), _3)) FROM make_date_non_ansi LIMIT 1", + "SELECT unix_date(make_date(2024, _2, 1)) + element_at(array(1), _3) FROM make_date_non_ansi") + guarded.foreach { query => + val df = sql(query) + val capped = query.endsWith("LIMIT 1") + if (capped) { + checkSparkAnswerAndOperator(df, includeClasses = Seq(classOf[CometProjectExec])) + } else { + checkSparkAnswer(df) + } + val plan = df.queryExecution.executedPlan + assert(countOperators(plan, classOf[ProjectExec]) == (if (capped) 0 else 1)) + assert(countOperators(plan, classOf[CometProjectExec]) == (if (capped) 1 else 0)) + if (capped) assert(countOperators(plan, classOf[CometLocalLimitExec]) == 1) + } + } + + } + + // The date's day attribute looks harmless locally, but its lower producer can throw. + withParquetTable(Seq((Some(1), 1), (None, 0)), "make_date_non_ansi_input") { + val df = sql(""" + |SELECT r, make_date(2024, m, d), r + 1 FROM ( + | SELECT rand(7L) AS r, _1 AS m, element_at(array(1), _2) AS d + | FROM make_date_non_ansi_input + |) + |""".stripMargin) + checkSparkAnswer(df) + val plan = df.queryExecution.executedPlan + assert(countOperators(plan, classOf[ProjectExec]) == 2, plan.treeString) + assert(countOperators(plan, classOf[CometProjectExec]) == 0, plan.treeString) + } + } + } + + test("make_date collect limit reuses the native projection input cap") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") { + withTempPath { path => + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark + .createDataFrame(Seq((2024, 1, "1"), (2024, 13, "bad"))) + .coalesce(1) + .write + .parquet(path.getCanonicalPath) + } + withParquetTable(path.getCanonicalPath, "make_date_limit_cap") { + for (capEnabled <- Seq(false, true)) { + withSQLConf(CometConf.COMET_EXEC_LOCAL_LIMIT_ENABLED.key -> capEnabled.toString) { + val df = sql("SELECT make_date(_1, _2, 1) FROM make_date_limit_cap LIMIT 1") + if (capEnabled) { + checkSparkAnswerAndOperator(df, includeClasses = Seq(classOf[CometProjectExec])) + } else { + checkSparkAnswer(df) + } + val converted = stripAQEPlan(df.queryExecution.executedPlan) + for (plan <- Seq(converted, applyCometExecRule(converted))) { + assert(countOperators(plan, classOf[ProjectExec]) == (if (capEnabled) 0 else 1)) + assert( + countOperators(plan, classOf[CometProjectExec]) == (if (capEnabled) 1 else 0)) + assert( + countOperators(plan, classOf[CometLocalLimitExec]) == (if (capEnabled) 1 + else 0)) + assert(new ExtendedExplainInfo() + .getFallbackReasons(plan) + .exists( + _.contains("make_date requires Spark evaluation below LIMIT")) != capEnabled) + } + + // The cap must still evaluate a bad date when it belongs to the consumed prefix. + val (sparkError, cometError) = checkSparkAnswerMaybeThrows( + sql("SELECT make_date(_1, _2, 1) FROM make_date_limit_cap LIMIT 2")) + Seq(sparkError, cometError).foreach { error => + assert( + error.exists(e => causeChain(e).exists(_.toString.contains("MonthOfYear")))) + } + } + } + + // A cap immediately below the date does not protect this lower producing projection. + // Repeated r keeps both projections, and LIMIT skips the second row's invalid cast. + val producer = sql(""" + |SELECT r, make_date(y, m, 1), r + 1 FROM ( + | SELECT rand(7L) AS r, _1 AS y, CAST(_3 AS INT) AS m FROM make_date_limit_cap + |) LIMIT 1 + |""".stripMargin) + checkSparkAnswer(producer) + val plan = producer.queryExecution.executedPlan + assert(countOperators(plan, classOf[ProjectExec]) == 2, plan.treeString) + assert(countOperators(plan, classOf[CometProjectExec]) == 0, plan.treeString) + } + } + } + } + + test("make_date shuffle keys remain eager below a limit") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_ENABLED.key -> "true") { + for (invalid <- Seq(false, true)) { + val rows = Seq.fill(20)((2024, 1, 1)) :+ ((2024, if (invalid) 13 else 1, 1)) + withParquetTable(rows, "make_date_shuffle") { + val df = spark + .table("make_date_shuffle") + .repartition(8, expr("make_date(_1, _2, _3)")) + .limit(10) + val plan = stripAQEPlan(df.queryExecution.executedPlan) + val exchanges = plan.collect { case exchange: CometShuffleExchangeExec => exchange } + withClue(plan.treeString) { + assert(exchanges.size == 1) + assert(exchanges.head.shuffleType == CometNativeShuffle) + assert( + exchanges.head.originalPlan.expressions.exists(_.exists(_.isInstanceOf[MakeDate]))) + } + if (invalid) { + // LIMIT cannot hide an invalid key, even when enough valid input rows precede it. + val (sparkError, cometError) = checkSparkAnswerMaybeThrows(df) + Seq("Spark" -> sparkError, "Comet" -> cometError).foreach { case (engine, error) => + withClue(s"$engine: ") { + assert( + error.exists(e => causeChain(e).exists(_.toString.contains("MonthOfYear")))) + } + } + } else { + checkSparkAnswer(df) + } + } + } + } + } + + test("make_date join input demand reaches local sorts before hash join rewriting") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { + val left = createSparkPlan(spark, "SELECT * FROM VALUES (2024), (2025) AS t(y)") + val right = createSparkPlan(spark, "SELECT * FROM VALUES (2024) AS t(y)") + val date = MakeDate(left.output.head, Literal(1), Literal(1), failOnError = false) + val projection = ProjectExec(Seq(left.output.head, Alias(date, "date")()), left) + def sort(child: SparkPlan): SparkPlan = + SortExec(Seq(SortOrder(child.output.head, Ascending)), false, child) + val join = SortMergeJoinExec( + Seq(left.output.head), + Seq(right.output.head), + Inner, + None, + sort(applyCometExecRule(projection)), + sort(right)) + val converted = applyCometExecRule(join) + for (result <- Seq(converted, applyCometExecRule(converted))) { + assert(countOperators(result, classOf[SortMergeJoinExec]) == 1, result.treeString) + assert(countOperators(result, classOf[CometProjectExec]) == 0, result.treeString) + assert(countOperators(result, classOf[ProjectExec]) == 1, result.treeString) + assert( + new ExtendedExplainInfo() + .getFallbackReasons(result) + .exists(_.contains("make_date requires Spark evaluation below JOIN"))) + } + } + } + + test("make_date remains deferred across conditional projections") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.DATETIME_JAVA8API_ENABLED.key -> "false") { + val selectedDate = java.sql.Date.valueOf("2024-01-01") + val rows = Seq((2024, 13, None), (2024, 1, Some(selectedDate))) + withParquetTable(rows, "make_date_deferred") { + def checkProjection(query: String)(checkSparkRows: Seq[Row] => Unit): SparkPlan = { + var expected = Seq.empty[Row] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + expected = sql(query).collect().toSeq + } + checkSparkRows(expected) + val df = sql(query) + checkCometAnswer(df, expected) + stripAQEPlan(df.queryExecution.executedPlan) + } + + def assertDeferredProjection(plan: SparkPlan): Unit = { + assert(countOperators(plan, classOf[ProjectExec]) == 2, plan.treeString) + assert(countOperators(plan, classOf[CometProjectExec]) == 0, plan.treeString) + assert(countOperators(plan, classOf[CometNativeScanExec]) == 1, plan.treeString) + assert( + new ExtendedExplainInfo() + .getFallbackReasons(plan) + .exists(_.contains("make_date requires Spark evaluation below PROJECT"))) + } + + val expressions = Seq( + "IF(r < 0, dt, DATE '2024-01-01')", + "COALESCE(IF(r >= 0, DATE '2024-01-01', NULL), dt)") + for (aqe <- Seq(false, true); expression <- expressions) { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + // Reusing nondeterministic r prevents the optimizer from merging the projections. + val query = s""" + |SELECT r, $expression AS d FROM ( + | SELECT rand(7L) AS r, make_date(_1, _2, 1) AS dt FROM make_date_deferred + |) + |""".stripMargin + val plan = checkProjection(query) { expected => + assert(expected.map(_.getDate(1)) == Seq.fill(2)(selectedDate)) + } + withClue(s"aqe=$aqe expression=$expression\n${plan.treeString}") { + assertDeferredProjection(plan) + } + } + } + + for (aqe <- Seq(false, true)) { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + val plan = checkProjection(""" + |SELECT r, datediff(nullable_date, dt) AS d, r + 1 AS repeated_r FROM ( + | SELECT rand(7L) AS r, _3 AS nullable_date, make_date(_1, _2, 1) AS dt + | FROM make_date_deferred + |) + |""".stripMargin) { expected => + assert(expected.size == 2) + assert(expected.map(_.get(1)).toSet == Set(null, 0)) + } + assertDeferredProjection(plan) + } + } + + // A nullable date can also hide a throwing sibling produced by the lower projection. + withParquetTable(Seq((None, "bad"), (Some(2024), "1")), "make_date_sibling") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val consumers = Seq( + "unix_date(dt) + rhs", + "date_add(dt, rhs)", + "(dt IS NULL) OR (rhs > 0)", + "IF(dt IS NULL, 0, rhs)") + for (expression <- consumers) { + val plan = checkProjection(s""" + |SELECT r, $expression AS result, r + 1 AS repeated_r FROM ( + | SELECT rand(7L) AS r, make_date(_1, 1, 1) AS dt, CAST(_2 AS INT) AS rhs + | FROM make_date_sibling + |) + |""".stripMargin) { expected => + assert(expected.size == 2) + } + assertDeferredProjection(plan) + } + + // Selecting dt forces the date eagerly, but the IF branch can still skip rhs. + val selected = checkProjection(""" + |SELECT r, dt, IF(dt IS NULL, 0, rhs) AS result, r + 1 AS repeated_r FROM ( + | SELECT rand(7L) AS r, make_date(_1, 1, 1) AS dt, CAST(_2 AS INT) AS rhs + | FROM make_date_sibling + |) + |""".stripMargin) { expected => + assert(expected.map(row => (row.getDate(1), row.getInt(2))).toSet == + Set((null, 0), (selectedDate, 1))) + } + assertDeferredProjection(selected) + + // A deferred source column stays native; it cannot expose a skipped error. + val safeSibling = checkProjection(""" + |SELECT r, dt, IF(dt IS NULL, 0, rhs) AS result, r + 1 AS repeated_r FROM ( + | SELECT rand(7L) AS r, make_date(_1, 1, 1) AS dt, _1 AS rhs + | FROM make_date_sibling + |) + |""".stripMargin) { expected => + assert(expected.map(row => (row.getDate(1), row.getInt(2))).toSet == + Set((null, 0), (selectedDate, 2024))) + } + assert( + countOperators(safeSibling, classOf[CometProjectExec]) == 2, + safeSibling.treeString) + assert( + countOperators(safeSibling, classOf[CometNativeScanExec]) == 1, + safeSibling.treeString) + + // Repeated r and r2 retain all three projections. The date and skipped cast can + // originate on either side of the intervening alias, so neither placement is safe. + val siblingProducers = Seq( + ( + "cast below date", + 0, + """ + |SELECT r, r + 1 AS repeated_r, rand(11L) AS r2, + | make_date(y, 1, 1) AS dt, raw_rhs AS rhs FROM ( + | SELECT rand(7L) AS r, _1 AS y, CAST(_2 AS INT) AS raw_rhs + | FROM make_date_sibling + |) + |""".stripMargin), + ( + "cast above date", + 1, + """ + |SELECT r, r + 1 AS repeated_r, rand(11L) AS r2, + | original_date AS dt, CAST(s AS INT) AS rhs FROM ( + | SELECT rand(7L) AS r, make_date(_1, 1, 1) AS original_date, _2 AS s + | FROM make_date_sibling + |) + |""".stripMargin)) + for ((placement, nativeProjects, producer) <- siblingProducers) { + val plan = checkProjection(s""" + |SELECT r, repeated_r, r2, r2 + 1, dt, IF(dt IS NULL, 0, rhs) AS result + |FROM ($producer) + |""".stripMargin) { expected => + assert(expected.map(row => (row.getDate(4), row.getInt(5))).toSet == + Set((null, 0), (selectedDate, 1))) + } + withClue(s"$placement\n${plan.treeString}") { + assert(countOperators(plan, classOf[ProjectExec]) == 3 - nativeProjects) + assert(countOperators(plan, classOf[CometProjectExec]) == nativeProjects) + assert(countOperators(plan, classOf[CometNativeScanExec]) == 1) + } + } + } + } + + // Each expression references dt once: repeated uses would force Spark to evaluate it + // eagerly and would not exercise the eager-reference analysis for that parent. + val eagerExpressions = Seq( + "year(dt)", + "month(dt)", + "date_add(dt, 1)", + "greatest(dt, DATE '2024-02-01')", + "hash(dt, 1)", + "unix_date(dt) + 1") + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + // YEAR evaluates its child only when its enclosing IF selects that branch. + val maskedYear = checkProjection(""" + |SELECT r, IF(r < 0, YEAR(dt), 1) AS result, r + 1 AS repeated_r FROM ( + | SELECT rand(7L) AS r, make_date(_1, _2, 1) AS dt FROM make_date_deferred + |) + |""".stripMargin) { expected => + assert(expected.map(_.getInt(1)) == Seq(1, 1)) + } + assertDeferredProjection(maskedYear) + + eagerExpressions.foreach { expression => + val plan = checkProjection(s""" + |SELECT r, $expression AS result, IF(r < 0, 0, r + 1) AS repeated_r FROM ( + | SELECT rand(7L) AS r, make_date(_1, _2, 1) AS dt + | FROM make_date_deferred WHERE _2 BETWEEN 1 AND 12 + |) + |""".stripMargin) { expected => + assert(expected.size == 1) + } + withClue(s"expression=$expression\n${plan.treeString}") { + assert(countOperators(plan, classOf[CometProjectExec]) == 2) + assert(countOperators(plan, classOf[CometNativeScanExec]) == 1) + assert( + !new ExtendedExplainInfo() + .getFallbackReasons(plan) + .exists(_.contains("make_date requires Spark evaluation below PROJECT"))) + } + } + } + } + } + } + + test("make_date keeps deferred arguments in their producing projections") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val rows = Seq((None, None, "bad"), (Some(2024), Some(1), "1")) + withParquetTable(rows, "make_date_arguments") { + def checkProjection( + query: String, + expectedDates: Set[Any], + dateIndex: Int = 1): SparkPlan = { + var expected = Seq.empty[Row] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + expected = sql(query).collect().toSeq + } + assert(expected.map(_.get(dateIndex)).toSet == expectedDates) + val df = sql(query) + checkCometAnswer(df, expected) + val plan = stripAQEPlan(df.queryExecution.executedPlan) + assert(countOperators(plan, classOf[CometNativeScanExec]) == 1, plan.treeString) + plan + } + + val cases = Seq( + ("_1 AS y, CAST(_3 AS INT) AS m", "make_date(y, m, 1)", "", true), + ("_2 AS m, CAST(_3 AS INT) AS d", "make_date(2024, m, d)", "", true), + // Reading a column is safe even if make_date skips the argument. + ("_1 AS y, _2 AS m", "make_date(y, m, 1)", "", false), + // The year is always evaluated, so its producer can remain native. + ( + "CAST(_3 AS INT) + 2023 AS y, _2 AS m", + "make_date(y, m, 1)", + "WHERE _1 IS NOT NULL", + false)) + for ((arguments, date, filter, deferred) <- cases) { + // Repeated r keeps the projections separate without forcing the date arguments. + val query = s""" + |SELECT r, unix_date($date) AS d, r + 1 AS repeated_r FROM ( + | SELECT rand(7L) AS r, $arguments FROM make_date_arguments $filter + |) + |""".stripMargin + val expectedDates: Set[Any] = if (filter.isEmpty) Set(null, 19723) else Set(19723) + val plan = checkProjection(query, expectedDates) + withClue(s"arguments=$arguments date=$date\n${plan.treeString}") { + assert(countOperators(plan, classOf[ProjectExec]) == (if (deferred) 2 else 0)) + assert(countOperators(plan, classOf[CometProjectExec]) == (if (deferred) 0 else 2)) + val reasons = new ExtendedExplainInfo().getFallbackReasons(plan) + assert(reasons.exists(_.contains("make_date requires Spark evaluation")) == deferred) + } + } + + // Broadcast joins can forward a streamed-side expression without evaluating it. The raw + // string join key does not force the month cast when the upper date sees a null year. + for (buildLeft <- Seq(false, true); deferred <- Seq(true, false)) { + val month = if (deferred) "CAST(_3 AS INT)" else "_2" + val probe = s"""( + | SELECT rand(7L) AS r, _1 AS y, _3 AS s, $month AS m + | FROM make_date_arguments + |) p""".stripMargin + val build = "(SELECT 'bad' AS k UNION ALL SELECT '1' AS k) b" + val inputs = if (buildLeft) s"$build JOIN $probe" else s"$probe JOIN $build" + val joined = checkProjection( + s""" + |SELECT /*+ BROADCAST(b) */ p.r, unix_date(make_date(p.y, p.m, 1)) AS d, + | p.r + 1 AS repeated_r + |FROM $inputs ON p.s = b.k + |""".stripMargin, + Set(null, 19723)) + withClue(s"buildLeft=$buildLeft deferred=$deferred\n${joined.treeString}") { + // Exclude the literal projections used to construct the broadcast input. + val sparkProjects = joined.collect { + case p: ProjectExec if p.output.exists(_.name == "r") => p + } + val nativeProjects = joined.collect { + case p: CometProjectExec if p.output.exists(_.name == "r") => p + } + assert(sparkProjects.size == (if (deferred) 2 else 0)) + assert(nativeProjects.size == (if (deferred) 0 else 2)) + assert( + countOperators(joined, classOf[BroadcastHashJoinExec]) == + (if (deferred) 1 else 0)) + assert( + countOperators(joined, classOf[CometBroadcastHashJoinExec]) == + (if (deferred) 0 else 1)) + val buildSides = joined.collect { + case join: BroadcastHashJoinExec => join.buildSide + case join: CometBroadcastHashJoinExec => join.buildSide + } + assert(buildSides == Seq(if (buildLeft) BuildLeft else BuildRight)) + } + } + + // A join key forces the cast before the date can skip its month argument. + val keyRequired = checkProjection( + """ + |SELECT /*+ BROADCAST(b) */ p.r, unix_date(make_date(p.y, p.m, 1)), p.r + 1 FROM ( + | SELECT rand(7L) AS r, _1 AS y, + | CAST(COALESCE(CAST(_2 AS STRING), '1') AS INT) AS m + | FROM make_date_arguments + |) p JOIN (SELECT 1 AS k UNION ALL SELECT 2 AS k) b ON p.m = b.k + |""".stripMargin, + Set(null, 19723)) + assert( + keyRequired.collect { + case p: CometProjectExec if p.output.exists(_.name == "r") => p + }.size == 2, + keyRequired.treeString) + assert( + countOperators(keyRequired, classOf[CometBroadcastHashJoinExec]) == 1, + keyRequired.treeString) + + // Repeated r and r2 retain an intervening projection that only aliases the month. + val aliased = checkProjection( + """ + |SELECT r, repeated_r, r2, r2 + 1, unix_date(make_date(y, month_alias, 1)) FROM ( + | SELECT r, r + 1 AS repeated_r, rand(11L) AS r2, y, m AS month_alias FROM ( + | SELECT rand(7L) AS r, _1 AS y, CAST(_3 AS INT) AS m FROM make_date_arguments + | ) + |) + |""".stripMargin, + Set(null, 19723), + dateIndex = 4) + assert(countOperators(aliased, classOf[ProjectExec]) == 3, aliased.treeString) + assert(countOperators(aliased, classOf[CometProjectExec]) == 0, aliased.treeString) + } + } + } + + test("make_date arguments remain deferred for unmatched broadcast join rows") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withParquetTable(Seq((1, "2024"), (2, "bad")), "make_date_join_probe") { + withParquetTable(Seq(Tuple1(1)), "make_date_join_build") { + def checkJoined(query: String, dateIndex: Int, projects: Int, scans: Int): SparkPlan = { + var expected = Seq.empty[Row] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + expected = sql(query).collect().toSeq + } + assert(expected.size == 1 && expected.head.getInt(dateIndex) == 19723) + val df = sql(query) + checkCometAnswer(df, expected) + val plan = stripAQEPlan(df.queryExecution.executedPlan) + withClue(plan.treeString) { + assert(plan.collect { + case p: ProjectExec if p.output.exists(_.name == "r") => p + }.size == projects) + assert(plan.collect { + case p: CometProjectExec if p.output.exists(_.name == "r") => p + }.isEmpty) + assert(countOperators(plan, classOf[CometNativeScanExec]) == scans) + } + plan + } + + val probe = """( + | SELECT _1 AS k, rand(7L) AS r, CAST(_2 AS INT) AS y + | FROM make_date_join_probe + |) p""".stripMargin + val build = "make_date_join_build b" + for (buildLeft <- Seq(false, true)) { + val inputs = if (buildLeft) s"$build JOIN $probe" else s"$probe JOIN $build" + val plan = checkJoined( + s""" + |SELECT /*+ BROADCAST(b) */ p.r, p.r + 1, unix_date(make_date(p.y, 1, 1)) + |FROM $inputs ON p.k = b._1 + |""".stripMargin, + dateIndex = 2, + projects = 2, + scans = 2) + withClue(s"buildLeft=$buildLeft\n${plan.treeString}") { + val joins = plan.collect { case join: BroadcastHashJoinExec => join } + assert(joins.map(_.buildSide) == Seq(if (buildLeft) BuildLeft else BuildRight)) + } + } + + // Repeated r2 retains the alias projection above the join that discards the bad row. + val aliased = checkJoined( + s""" + |SELECT r, repeated_r, r2, r2 + 1, unix_date(make_date(date_year, 1, 1)) FROM ( + | SELECT /*+ BROADCAST(b) */ p.r, p.r + 1 AS repeated_r, rand(11L) AS r2, + | p.y AS date_year FROM $probe JOIN $build ON p.k = b._1 + |) + |""".stripMargin, + dateIndex = 4, + projects = 3, + scans = 2) + assert(countOperators(aliased, classOf[BroadcastHashJoinExec]) == 1, aliased.treeString) + + // The outer join requires y only after the inner join accepts a row. Null-safe equality + // keeps an inferred isNotNull(y) filter from evaluating the cast before either join. + // Spark also inserts a projection between the joins to drop the inner join keys. + val nested = checkJoined( + s""" + |SELECT /*+ BROADCAST(b, c) */ p.r, p.r + 1, unix_date(make_date(p.y, 1, 1)) + |FROM $probe JOIN $build ON p.k = b._1 + |JOIN (SELECT _1 + 2023 AS y FROM make_date_join_build) c ON p.y <=> c.y + |""".stripMargin, + dateIndex = 2, + projects = 3, + scans = 3) + assert(countOperators(nested, classOf[BroadcastHashJoinExec]) == 2, nested.treeString) + } + } + } + } + + test("make_date single-argument count evaluates its nullable argument eagerly") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val rows = Seq((0, Some(2024), 1), (0, None, 1), (1, Some(2025), 1)) + withParquetTable(rows, "make_date_count") { + val query = "SELECT _1, count(make_date(_2, _3, 1)), sum(_3) " + + "FROM make_date_count GROUP BY _1" + val df = sql(query) + checkCometAnswer(df, Seq(Row(0, 1L, 2L), Row(1, 1L, 1L))) + val plan = stripAQEPlan(df.queryExecution.executedPlan) + assert(countOperators(plan, classOf[HashAggregateExec]) == 0, plan.treeString) + assert(countOperators(plan, classOf[CometHashAggregateExec]) == 2, plan.treeString) + + // A physical plan bypasses the optimizer's COUNT(non-nullable) -> COUNT(1) rewrite. + // Spark does not evaluate this invalid date, so it must not reach a native update. + val partial = createSparkPlan(spark, query).collectFirst { + case aggregate: HashAggregateExec + if aggregate.aggregateExpressions.forall(_.mode == Partial) => + aggregate + }.get + val skippedDate = MakeDate(Literal(2024), Literal(13), Literal(1), failOnError = true) + assert(!skippedDate.nullable) + val nonNullableCount = + partial.copy(aggregateExpressions = partial.aggregateExpressions.map { + case expression if expression.aggregateFunction.isInstanceOf[Count] => + expression.copy(aggregateFunction = Count(Seq(skippedDate))) + case expression => expression + }) + assert(MakeDateEvaluation.reason(MakeDateEvaluation.preserve(nonNullableCount)).isDefined) + } + } + } + + test("make_date aggregate arguments retain Spark execution") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.USE_OBJECT_HASH_AGG.key -> "true") { + val rows = Seq((2024, 1, Some(1)), (2024, 13, None), (2025, 1, Some(1))) + withParquetTable(rows, "make_date_aggregate") { + val cases = Seq( + ( + "unix_date(min(make_date(_1, _2, 1)) FILTER (WHERE _2 <= 12))", + classOf[HashAggregateExec], + Row(19723)), + ( + "sort_array(collect_list(unix_date(make_date(_1, _2, 1))) " + + "FILTER (WHERE _2 <= 12))", + classOf[ObjectHashAggregateExec], + Row(Seq(19723, 20089))), + ( + "count(make_date(_1, _2, 1)) FILTER (WHERE _2 <= 12)", + classOf[HashAggregateExec], + Row(2L)), + ("count(_3, make_date(_1, _2, 1))", classOf[HashAggregateExec], Row(2L))) + cases.foreach { case (expression, aggregateClass, expected) => + val df = sql(s"SELECT $expression FROM make_date_aggregate") + val plan = stripAQEPlan(df.queryExecution.executedPlan) + assert( + plan.exists { op => + aggregateClass.isInstance(op) && + op.expressions.exists(_.exists(_.isInstanceOf[MakeDate])) + }, + plan.treeString) + val name = aggregateClass.getSimpleName.stripSuffix("Exec") + assert( + new ExtendedExplainInfo() + .getFallbackReasons(plan) + .exists(_.contains(s"make_date requires Spark evaluation in $name")), + plan.treeString) + checkCometAnswer(df, Seq(expected)) + } + } + } + } + + test("make_date top-k respects input ordering and deferred output projection") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + // Ordered top-k takes one row per partition before the final selection. + SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "1", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true") { + val template = createSparkPlan( + spark, + "SELECT * FROM VALUES (1), (13) AS t(m) ORDER BY m LIMIT 1").collectFirst { + case topK: TakeOrderedAndProjectExec => topK + }.get + val input = template.child + val date = Alias(MakeDate(Literal(2024), input.output.head, Literal(1), true), "date")() + def project(child: SparkPlan): ProjectExec = ProjectExec(input.output :+ date, child) + val order = template.sortOrder + for (sorted <- Seq(false, true)) { + val child = project(if (sorted) SortExec(order, false, input) else input) + assert(SortOrder.orderingSatisfies(child.outputOrdering, order) == sorted) + val result = applyCometExecRule(template.copy(projectList = child.output, child = child)) + if (sorted) { + assert(result.isInstanceOf[TakeOrderedAndProjectExec], result.treeString) + assert(countOperators(result, classOf[CometProjectExec]) == 0, result.treeString) + assert( + new ExtendedExplainInfo() + .getFallbackReasons(result) + .exists(_.contains("make_date requires Spark evaluation below LIMIT"))) + assert( + result.executeCollect().map(row => (row.getInt(0), row.getInt(1))).toSeq == + Seq((1, 19723))) + } else { + assert(result.isInstanceOf[CometTakeOrderedAndProjectExec], result.treeString) + assert(countOperators(result, classOf[CometProjectExec]) == 1, result.treeString) + } + } + val deferred = applyCometExecRule(template.copy(projectList = input.output :+ date)) + assert(deferred.isInstanceOf[TakeOrderedAndProjectExec], deferred.treeString) + assert( + new ExtendedExplainInfo() + .getFallbackReasons(deferred) + .exists(_.contains("make_date requires Spark evaluation in TakeOrderedAndProject"))) + assert( + deferred.executeCollect().map(row => (row.getInt(0), row.getInt(1))).toSeq == + Seq((1, 19723))) + } + } + + test("make_date below a window group limit retains Spark evaluation") { + val windowLimitClass = ShimCometWindowGroupLimit.windowGroupLimitClass + assume(windowLimitClass.isDefined, "WindowGroupLimit requires Spark 3.5+") + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", + CometConf.COMET_EXEC_WINDOW_GROUP_LIMIT_ENABLED.key -> "true") { + val plan = createSparkPlan( + spark, + """ + |SELECT p, y FROM ( + | SELECT p, y, row_number() OVER (PARTITION BY p ORDER BY y) AS rn + | FROM VALUES (0, 2024), (0, 2025) AS t(p, y) + |) WHERE rn <= 1 + |""".stripMargin) + val limit = plan.collectFirst { case op if windowLimitClass.get.isInstance(op) => op }.get + val input = limit.children.head + val year = input.output.find(_.name == "y").get + val date = Alias(MakeDate(year, Literal(1), Literal(1), false), "date")() + val projected = ProjectExec(input.output :+ date, input) + val result = applyCometExecRule(limit.withNewChildren(Seq(projected))) + assert(windowLimitClass.get.isInstance(result), result.treeString) + assert( + result.exists { + case p: ProjectExec => p.projectList.exists(_.exists(_.isInstanceOf[MakeDate])) + case _ => false + }, + result.treeString) + assert( + new ExtendedExplainInfo() + .getFallbackReasons(result) + .exists(_.contains("make_date requires Spark evaluation below LIMIT")), + result.treeString) + } + } + + test("make_date operator opt-out retains invalid-date errors and expression guards") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.ANSI_ENABLED.key -> "true", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true") { + val input = createSparkPlan(spark, "SELECT * FROM VALUES (2024) AS t(y)") + val date = Alias(MakeDate(input.output.head, Literal(1), Literal(1), true), "date")() + // Tag the Spark plan before disabling the setting. + val guarded = MakeDateEvaluation.preserve(LocalLimitExec(1, ProjectExec(Seq(date), input))) + assert(countOperators(guarded, classOf[ProjectExec]) == 1, guarded.treeString) + assert(MakeDateEvaluation.reason(guarded).isDefined, guarded.treeString) + withSQLConf(CometConf.COMET_EXEC_MAKE_DATE_OPERATOR_FALLBACK_ENABLED.key -> "false") { + assert(MakeDateEvaluation.reason(guarded).isEmpty) + val converted = applyCometExecRule(guarded) + assert(countOperators(converted, classOf[CometProjectExec]) == 1, converted.treeString) + assert(countOperators(converted, classOf[CometLocalLimitExec]) == 1, converted.treeString) + withParquetTable(Seq((2024, 1, 1), (2025, 1, 1)), "make_date_opt_out") { + checkSparkAnswerAndOperator( + sql("SELECT unix_date(make_date(_1, _2, _3)) FROM make_date_opt_out LIMIT 1"), + Seq(classOf[CometProjectExec], classOf[CometCollectLimitExec])) + val (_, joinPlan) = checkSparkAnswerAndOperator( + sql(""" + |SELECT /*+ BROADCAST(b) */ unix_date(p.d), b._1 FROM ( + | SELECT _1 AS y, make_date(_1, _2, _3) AS d FROM make_date_opt_out + |) p JOIN make_date_opt_out b ON p.y = b._1 + |""".stripMargin), + Seq(classOf[CometProjectExec], classOf[CometBroadcastHashJoinExec])) + val join = joinPlan.collectFirst { case j: CometBroadcastHashJoinExec => j }.get + assert( + join.left.exists { + case p: CometProjectExec => + p.originalPlan.expressions.exists(_.exists(_.isInstanceOf[MakeDate])) + case _ => false + }, + joinPlan.treeString) + + val invalidDate = sql("SELECT make_date(_1, _2 + 12, _3) FROM make_date_opt_out") + checkCometOperators(stripAQEPlan(invalidDate.queryExecution.executedPlan)) + val (sparkError, cometError) = checkSparkAnswerMaybeThrows(invalidDate) + Seq("Spark" -> sparkError, "Comet" -> cometError).foreach { case (engine, error) => + withClue(s"$engine: ") { + assert(error.exists(e => causeChain(e).exists(_.toString.contains("MonthOfYear")))) + } + } + } + + // A nullable left operand must still skip a throwing date on the right of arithmetic. + withParquetTable(Seq((None, 13), (Some(1), 1)), "make_date_opt_out_expression") { + val df = sql(""" + |SELECT _1 + unix_date(make_date(2024, _2, 1)) FROM make_date_opt_out_expression + |""".stripMargin) + val plan = stripAQEPlan(df.queryExecution.executedPlan) + assert(countOperators(plan, classOf[ProjectExec]) == 1, plan.treeString) + assert( + new ExtendedExplainInfo() + .getFallbackReasons(plan) + .exists(_.contains("make_date requires Spark evaluation under Add")), + plan.treeString) + checkCometAnswer(df, Seq(Row(null), Row(19724))) + } + } + } + } + + test("make_date fallback reuses native broadcast rows in Spark's broadcast representation") { + for (nestedLoop <- Seq(false, true)) { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val input = createSparkPlan(spark, "SELECT * FROM VALUES (2024) AS t(y)") + val date = MakeDate(input.output.head, Literal(1), Literal(1), failOnError = false) + val probe = + applyCometExecRule(ProjectExec(Seq(input.output.head, Alias(date, "date")()), input)) + val build = applyCometExecRule( + createSparkPlan(spark, "SELECT CAST(id AS INT) AS y FROM range(2024, 2028, 1, 2)")) + assert(build.isInstanceOf[CometNativeExec], build.treeString) + val key = build.output.head + val mode = + if (nestedLoop) IdentityBroadcastMode + else HashedRelationBroadcastMode(Seq(BoundReference(0, key.dataType, key.nullable))) + val sparkExchange = BroadcastExchangeExec(mode, build) + val exchange = CometBroadcastExchangeExec(sparkExchange, build.output, mode, build) + val reused = ReusedExchangeExec(build.output.map(_.newInstance()), exchange) + val inputs = Seq[SparkPlan]( + exchange, + reused, + BroadcastQueryStageExec(0, exchange, exchange.canonicalized), + BroadcastQueryStageExec(1, reused, reused.canonicalized)) + inputs.zipWithIndex.foreach { case (broadcast, index) => + // Model a native plan supplied by an earlier planning pass; its parent must restore it. + val native = if (nestedLoop) { + val original = BroadcastNestedLoopJoinExec(probe, broadcast, BuildRight, Inner, None) + CometBroadcastNestedLoopJoinExec.createExec( + OperatorOuterClass.Operator.getDefaultInstance, + original) + } else { + val original = BroadcastHashJoinExec( + Seq(probe.output.head), + Seq(broadcast.output.head), + Inner, + BuildRight, + None, + probe, + broadcast) + CometBroadcastHashJoinExec.createExec( + OperatorOuterClass.Operator.getDefaultInstance, + original) + } + val converted = applyCometExecRule(native) + for (result <- Seq(converted, applyCometExecRule(converted))) { + val join = result.asInstanceOf[BaseJoinExec] + assert(join.isInstanceOf[BroadcastNestedLoopJoinExec] == nestedLoop) + val restored = join.right.asInstanceOf[CometColumnarToRowExec] + assert(restored.output == broadcast.output) + assert(restored.child eq broadcast) + assert(countOperators(join.left, classOf[CometProjectExec]) == 0) + } + if (index == 0) { + val restored = converted + .asInstanceOf[BaseJoinExec] + .right + .asInstanceOf[CometColumnarToRowExec] + val relation = restored.executeBroadcast[Any]().value + val rows = if (nestedLoop) { + relation.asInstanceOf[Array[InternalRow]].toVector + } else { + // HashedRelation is private[execution]; inspect Spark's broadcast through its keys/get + // methods. + val keys = PrivateMethod[Iterator[InternalRow]](Symbol("keys")) + val get = PrivateMethod[Iterator[InternalRow]](Symbol("get")) + val hashed = relation.asInstanceOf[AnyRef] + (hashed invokePrivate keys()).flatMap { key => + (hashed invokePrivate get(key)).map(_.copy()).toVector + }.toVector + } + // Restore the complete build relation without re-executing its native input. + withClue(s"nestedLoop=$nestedLoop: ") { + assert(rows.map(_.getInt(0)).sorted == Seq(2024, 2025, 2026, 2027)) + } + } + } + } + } + } + /** * Build a Spark plan containing an `ObjectHashAggregateExec` and hand it to `f`. `collect_list` * is a `TypedImperativeAggregate`, so Spark plans it as `ObjectHashAggregateExec` rather than diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometDatetimeExpressionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometDatetimeExpressionBenchmark.scala index 1ff82a6d450..634fc230664 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometDatetimeExpressionBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometDatetimeExpressionBenchmark.scala @@ -20,7 +20,12 @@ package org.apache.spark.sql.benchmark import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.catalyst.expressions.MakeDate +import org.apache.spark.sql.catalyst.optimizer.BuildRight import org.apache.spark.sql.catalyst.util.DateTimeTestUtils.{withDefaultTimeZone, LA} +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometProjectExec} +import org.apache.spark.sql.execution.ProjectExec +import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf @@ -30,11 +35,108 @@ import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus /** * Benchmark to measure Comet execution performance. To run this benchmark: * `SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometDatetimeExpressionBenchmark` + * Append `-- --make-date-only` to run only the make_date operator fallback cases. * Results will be written to "spark/benchmarks/CometDatetimeExpressionBenchmark-**results.txt". */ // spotless:on object CometDatetimeExpressionBenchmark extends CometBenchmarkBase { + def makeDateOperatorFallbackBenchmark(values: Int): Unit = { + withTempPath { dir => + withTempTable("make_date_probe", "make_date_build") { + val probePath = s"${dir.getCanonicalPath}/probe" + val buildPath = s"${dir.getCanonicalPath}/build" + spark + .range(values.toLong) + .selectExpr( + "CAST(id % 1024 AS INT) AS k", + "CAST(2000 + id % 25 AS INT) AS y", + "CAST(1 + id % 12 AS INT) AS m", + "CAST(1 + id % 28 AS INT) AS d") + .write + .parquet(probePath) + spark.range(1024).selectExpr("CAST(id AS INT) AS k").write.parquet(buildPath) + spark.read.parquet(probePath).createOrReplaceTempView("make_date_probe") + spark.read.parquet(buildPath).createOrReplaceTempView("make_date_build") + + val queries = Seq( + ( + "MakeDate - LIMIT 10", + "SELECT make_date(y, m, d) AS dt FROM make_date_probe LIMIT 10", + true), + ( + "MakeDate - broadcast join probe projection", + """SELECT /*+ BROADCAST(b) */ p.dt, b.k + |FROM (SELECT k, make_date(y, m, d) AS dt FROM make_date_probe) p + |JOIN make_date_build b ON p.k = b.k""".stripMargin, + false)) + val modes = Seq( + ("Spark", false, true), + ("Comet (default)", true, true), + ("Comet (operator fallback disabled)", true, false)) + queries.foreach { case (name, query, limited) => + val benchmark = new Benchmark( + name, + if (limited) 10L else values.toLong, + minNumIters = 5, + output = output) + benchmark.out.println( + s"Input: $values probe rows, 1024 build rows, ${spark.sparkContext.master}, AQE off") + benchmark.out.println(query) + modes.foreach { case (mode, comet, operatorFallback) => + val configs = Seq( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_ENABLED.key -> comet.toString, + CometConf.COMET_EXEC_ENABLED.key -> comet.toString, + CometConf.COMET_EXEC_MAKE_DATE_OPERATOR_FALLBACK_ENABLED.key -> + operatorFallback.toString) + def consume(): Unit = { + val df = spark.sql(query) + // Collect only ten rows for LIMIT. The noop writer consumes the full join output, + // preserving the date projection that count(*) would prune from the query. + if (limited) { + val _ = df.collect() + } else { + df.noop() + } + } + withSQLConf(configs: _*) { + // Capture the executed action, including the noop writer's actual query plan. + val plans = capturePlans { consume() } + require(plans.nonEmpty, s"No executed plan captured for $name / $mode") + benchmark.out.println(s"$mode executed plans:") + plans.foreach(plan => benchmark.out.println(plan.treeString)) + if (!limited) { + val probes = collectAcross(plans) { + case join: BroadcastHashJoinExec => + if (join.buildSide == BuildRight) join.left else join.right + case join: CometBroadcastHashJoinExec => + if (join.buildSide == BuildRight) join.left else join.right + } + require( + probes.exists(_.exists { + case p: ProjectExec => + p.projectList.exists(_.exists(_.isInstanceOf[MakeDate])) + case p: CometProjectExec => + p.projectList.exists(_.exists(_.isInstanceOf[MakeDate])) + case _ => false + }), + s"make_date must remain in a projection below the broadcast join for $mode") + } + } + benchmark.addCase(mode) { _ => + withSQLConf(configs: _*) { + consume() + } + } + } + benchmark.run() + } + } + } + } + def dateTruncExprBenchmark(values: Int): Unit = { withTempPath { dir => withTempTable("parquetV1Table") { @@ -226,6 +328,11 @@ object CometDatetimeExpressionBenchmark extends CometBenchmarkBase { override def runCometBenchmark(mainArgs: Array[String]): Unit = { val values = 1024 * 1024; + runBenchmark("MakeDate operator fallback") { + makeDateOperatorFallbackBenchmark(values) + } + if (mainArgs.contains("--make-date-only")) return + for (timeZone <- Seq("UTC", "America/Los_Angeles")) { withSQLConf("spark.sql.parquet.datetimeRebaseModeInWrite" -> "CORRECTED") { runBenchmarkWithTable(s"UnixTimestamp(timestamp) - $timeZone", values) { v =>