From 0dc09ecca85194b8f4644c30287506ea79230a85 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 28 Oct 2026 00:31:03 +0000 Subject: [PATCH 01/14] fix: preserve make_date row demand across operators --- .../apache/comet/rules/CometExecRule.scala | 10 +- .../comet/rules/MakeDateEvaluation.scala | 195 +++++++++++ .../comet/CometTemporalExpressionSuite.scala | 74 +++++ .../comet/rules/CometExecRuleSuite.scala | 304 +++++++++++++++++- 4 files changed, 576 insertions(+), 7 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala 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..ee317f7799a 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 @@ -1041,15 +1044,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..b471ca51226 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala @@ -0,0 +1,195 @@ +/* + * 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 org.apache.spark.sql.catalyst.expressions.{Expression, MakeDate, SortOrder} +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 + +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") + + def reason(plan: SparkPlan): Option[String] = + plan.getTagValue(sparkEvaluation) + + private def containsMakeDate(expressions: Seq[Expression]): Boolean = + expressions.exists(_.exists(_.isInstanceOf[MakeDate])) + + 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 = { + def visit( + node: SparkPlan, + rowDemand: Option[String], + projectionDemand: Option[String]): (SparkPlan, Option[String]) = { + node match { + // Materialized stages are independent producers. Their plans have already been guarded. + case _: QueryStageExec | _: ReusedExchangeExec => return (node, None) + // These wrappers do not have the child shape of their original Spark operator. + case wrapper: CometScanWrapper => + val (child, childReason) = visit(wrapper.originalPlan, rowDemand, projectionDemand) + return if (childReason.isDefined) (child, childReason) + else (wrapper.copy(originalPlan = child), None) + case sink: CometSinkPlaceHolder => + val (child, childReason) = visit(sink.child, rowDemand, projectionDemand) + return if (childReason.isDefined) (child, childReason) + else (sink.withNewChildren(Seq(child)), None) + 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 exchange = original.isInstanceOf[Exchange] + val join = original.isInstanceOf[BaseJoinExec] + 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 _: FilterExec => Some("FILTER") + case _ if exchange || localMaterialization || join => None + case _ => projectionDemand + } + val visited = node.children.map(visit(_, childDemand, childProjectionDemand)) + val childReason = visited.flatMap(_._2).headOption + val ownReason = reason(node).orElse { + if (containsMakeDate(original.expressions)) { + rowDemand + .map(boundary => s"make_date requires Spark evaluation below $boundary") + .orElse(original match { + case project: ProjectExec + if containsMakeDate(project.projectList.filter(_.deterministic)) => + projectionDemand.map(boundary => + s"make_date requires Spark evaluation below $boundary") + // Join keys/conditions and aggregate updates can be skipped inside the operator. + case _: BaseJoinExec | _: BaseAggregateExec | _: TakeOrderedAndProjectExec => + Some(s"make_date requires Spark evaluation in ${original.nodeName}") + case _ => None + }) + } else None + } + 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) + fallbackReason.foreach(rebuilt.setTagValue(sparkEvaluation, _)) + (rebuilt, fallbackReason) + } + + // 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/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala index e61a57a6418..4e74279b074 100644 --- a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala @@ -26,6 +26,8 @@ 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.CometProjectExec +import org.apache.spark.sql.execution.ProjectExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.functions.col import org.apache.spark.sql.internal.SQLConf @@ -56,6 +58,78 @@ 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") + .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 }.nonEmpty, plan.treeString) + assert(plan.collect { case p: CometProjectExec => p }.isEmpty, 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) + } + } + + // 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..7f006c35c54 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.{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, 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.{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,15 @@ 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.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 +154,299 @@ 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)) { + def project: SparkPlan = ProjectExec( + Seq( + input.output.head, + Alias(MakeDate(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 converted = applyCometExecRule(consume(child)) + for (result <- Seq(converted, applyCometExecRule(converted))) { + withClue(s"ansi=$ansi reused=$reused $boundary\n${result.treeString}") { + assert(countOperators(result, classOf[ProjectExec]) == 1) + assert(countOperators(result, classOf[CometProjectExec]) == 0) + assert( + new ExtendedExplainInfo() + .getFallbackReasons(result) + .exists(_.contains(s"make_date requires Spark evaluation below $boundary"))) + } + } + } + + // 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 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 aggregate arguments retain Spark execution") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.USE_OBJECT_HASH_AGG.key -> "true") { + withParquetTable(Seq((2024, 1), (2024, 13), (2025, 1)), "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)))) + 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 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 From 4d8be556f866a5682e595aad30faa7a144fc2a34 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 28 Oct 2026 00:35:25 +0000 Subject: [PATCH 02/14] fix: preserve deferred make_date projections --- .../comet/rules/MakeDateEvaluation.scala | 68 ++++++++-- .../org/apache/comet/serde/datetime.scala | 2 +- .../comet/rules/CometExecRuleSuite.scala | 117 ++++++++++++++++++ 3 files changed, 179 insertions(+), 8 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala index b471ca51226..0f0c8982ee4 100644 --- a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala +++ b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala @@ -19,7 +19,7 @@ package org.apache.comet.rules -import org.apache.spark.sql.catalyst.expressions.{Expression, MakeDate, SortOrder} +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, And, Attribute, AttributeSet, BinaryComparison, Cast, ConditionalExpression, CreateArray, CreateNamedStruct, DateAdd, DateDiff, DateFormatClass, DateSub, EqualNullSafe, Expression, Greatest, HashExpression, IsNotNull, IsNull, Least, MakeDate, Multiply, Not, Or, SortOrder, Subtract, UnixDate} import org.apache.spark.sql.catalyst.trees.TreeNodeTag import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec @@ -29,6 +29,7 @@ import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregat import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} import org.apache.spark.sql.execution.joins.BaseJoinExec +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. */ @@ -41,6 +42,43 @@ private[rules] object MakeDateEvaluation { private def containsMakeDate(expressions: Seq[Expression]): Boolean = expressions.exists(_.exists(_.isInstanceOf[MakeDate])) + private case class ProjectionDemand(boundary: String, deferred: 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 conditional: ConditionalExpression => + val required = conditional.alwaysEvaluatedInputs + if (conditional.children.filterNot(required.contains).forall(harmlessValue)) required + else Nil + case _: Alias | _: Cast | _: IsNull | _: IsNotNull | _: Not | _: UnixDate => + 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. @@ -83,7 +121,7 @@ private[rules] object MakeDateEvaluation { def visit( node: SparkPlan, rowDemand: Option[String], - projectionDemand: Option[String]): (SparkPlan, Option[String]) = { + projectionDemand: Option[ProjectionDemand]): (SparkPlan, Option[String]) = { node match { // Materialized stages are independent producers. Their plans have already been guarded. case _: QueryStageExec | _: ReusedExchangeExec => return (node, None) @@ -138,7 +176,20 @@ private[rules] object MakeDateEvaluation { } val childProjectionDemand = original match { // Generated filters can defer deterministic projection outputs until a row is accepted. - case _: FilterExec => Some("FILTER") + case filter: FilterExec => Some(ProjectionDemand("FILTER", filter.child.outputSet)) + 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) + if (skipped.isEmpty) { + None + } else { + Some(ProjectionDemand(projectionDemand.map(_.boundary).getOrElse("PROJECT"), skipped)) + } case _ if exchange || localMaterialization || join => None case _ => projectionDemand } @@ -149,10 +200,13 @@ private[rules] object MakeDateEvaluation { rowDemand .map(boundary => s"make_date requires Spark evaluation below $boundary") .orElse(original match { - case project: ProjectExec - if containsMakeDate(project.projectList.filter(_.deterministic)) => - projectionDemand.map(boundary => - s"make_date requires Spark evaluation below $boundary") + case project: ProjectExec => + projectionDemand + .filter(demand => + project.projectList.exists(output => + output.deterministic && demand.deferred.contains(output.toAttribute) && + containsMakeDate(Seq(output)))) + .map(demand => s"make_date requires Spark evaluation below ${demand.boundary}") // Join keys/conditions and aggregate updates can be skipped inside the operator. case _: BaseJoinExec | _: BaseAggregateExec | _: TakeOrderedAndProjectExec => Some(s"make_date requires Spark evaluation in ${original.nodeName}") 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/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala index 7f006c35c54..5afe10c1518 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -241,6 +241,123 @@ class CometExecRuleSuite extends CometTestBase { } } + 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) + } + } + } + + // 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( + "date_add(dt, 1)", + "greatest(dt, DATE '2024-02-01')", + "hash(dt, 1)", + "unix_date(dt) + 1") + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + 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 aggregate arguments retain Spark execution") { withSQLConf( SQLConf.ANSI_ENABLED.key -> "true", From 8ebd569c4ea47b6bd146acaa4783275317aaa0f5 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 28 Oct 2026 00:37:19 +0000 Subject: [PATCH 03/14] feat: add a supported make_date operator fallback opt-out --- docs/source/user-guide/latest/expressions.md | 29 ++++++++- .../scala/org/apache/comet/CometConf.scala | 11 ++++ .../comet/rules/MakeDateEvaluation.scala | 8 ++- .../comet/rules/CometExecRuleSuite.scala | 65 +++++++++++++++++++ 4 files changed, 111 insertions(+), 2 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index f66a705afac..aa0809945a0 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,33 @@ 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 +to remain native when their operands are safe to evaluate. + +Only affected parts of the plan fall back. An ordinary native scan can remain native. +Exchanges materialize their inputs independently, and local sorts and aggregates can also +separate date evaluation from an outer limit. + +Users who accept evaluating dates on rows Spark would skip can set +`spark.comet.exec.makeDate.operatorFallback.enabled=false` to allow native operators for +these shapes. The default is `true`. Disabling it can produce an invalid-date error where +Spark would return rows successfully 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/MakeDateEvaluation.scala b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala index 0f0c8982ee4..8e8abf9365f 100644 --- a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala +++ b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala @@ -29,6 +29,7 @@ import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregat import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} import org.apache.spark.sql.execution.joins.BaseJoinExec +import org.apache.comet.CometConf import org.apache.comet.serde.CometMakeDate import org.apache.comet.shims.ShimCometWindowGroupLimit @@ -36,8 +37,11 @@ import org.apache.comet.shims.ShimCometWindowGroupLimit private[rules] object MakeDateEvaluation { private val sparkEvaluation = TreeNodeTag[String]("comet.makeDateSparkEvaluation") + private def enabled(plan: SparkPlan): Boolean = + CometConf.COMET_EXEC_MAKE_DATE_OPERATOR_FALLBACK_ENABLED.get(plan.conf) + def reason(plan: SparkPlan): Option[String] = - plan.getTagValue(sparkEvaluation) + if (enabled(plan)) plan.getTagValue(sparkEvaluation) else None private def containsMakeDate(expressions: Seq[Expression]): Boolean = expressions.exists(_.exists(_.isInstanceOf[MakeDate])) @@ -118,6 +122,8 @@ private[rules] object MakeDateEvaluation { * 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], 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 5afe10c1518..d5e75acb360 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -481,6 +481,71 @@ class CometExecRuleSuite extends CometTestBase { } } + 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( From f07f57bc5089400c34ac6661f838c46e31fd1e36 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 28 Oct 2026 00:39:18 +0000 Subject: [PATCH 04/14] bench: compare make_date operator fallback costs --- .../CometDatetimeExpressionBenchmark.scala | 107 ++++++++++++++++++ 1 file changed, 107 insertions(+) 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 => From 33e1cd082c385578c2bf0710aae50131bd2a14a7 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 28 Oct 2026 01:41:36 +0000 Subject: [PATCH 05/14] fix: preserve deferred make_date input projections --- docs/source/user-guide/latest/expressions.md | 9 +- .../comet/rules/MakeDateEvaluation.scala | 82 ++++++++++++++----- .../comet/CometTemporalExpressionSuite.scala | 22 ++++- .../comet/rules/CometExecRuleSuite.scala | 68 +++++++++++++++ 4 files changed, 156 insertions(+), 25 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index aa0809945a0..684c7519a0f 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -726,16 +726,17 @@ This applies to date projections below `LIMIT`, `SAMPLE`, `FILTER`, joins, 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 -to remain native when their operands are safe to evaluate. +to remain native when their operands are safe to evaluate. These safeguards also preserve +NULL short-circuiting when a date argument is computed in an earlier projection. Only affected parts of the plan fall back. An ordinary native scan can remain native. Exchanges materialize their inputs independently, and local sorts and aggregates can also separate date evaluation from an outer limit. -Users who accept evaluating dates on rows Spark would skip can set +Users who accept eagerly evaluating dates and their projected arguments can set `spark.comet.exec.makeDate.operatorFallback.enabled=false` to allow native operators for -these shapes. The default is `true`. Disabling it can produce an invalid-date error where -Spark would return rows successfully under ANSI mode. This setting does not change +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. diff --git a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala index 8e8abf9365f..c4f5e617519 100644 --- a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala +++ b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala @@ -46,7 +46,10 @@ private[rules] object MakeDateEvaluation { private def containsMakeDate(expressions: Seq[Expression]): Boolean = expressions.exists(_.exists(_.isInstanceOf[MakeDate])) - private case class ProjectionDemand(boundary: String, deferred: AttributeSet) + private case class ProjectionDemand( + boundary: String, + deferred: AttributeSet, + makeDateInputs: AttributeSet) // Attributes may stand for deferred expressions in a lower projection, not materialized values. private def harmlessValue(expression: Expression): Boolean = @@ -57,6 +60,9 @@ private[rules] object MakeDateEvaluation { 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 @@ -182,7 +188,12 @@ private[rules] object MakeDateEvaluation { } 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)) + case filter: FilterExec => + Some( + ProjectionDemand( + "FILTER", + filter.child.outputSet, + projectionDemand.map(_.makeDateInputs).getOrElse(AttributeSet.empty))) case project: ProjectExec => val deferred = projectionDemand.map(_.deferred).getOrElse(AttributeSet.empty) val eager = project.projectList @@ -191,10 +202,31 @@ private[rules] object MakeDateEvaluation { // 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) - if (skipped.isEmpty) { + val inheritedInputs = + projectionDemand.map(_.makeDateInputs).getOrElse(AttributeSet.empty) + val dateInputs = AttributeSet(project.projectList.flatMap { output => + val direct = output.collect { case date: MakeDate => date }.flatMap(_.references) + // Follow aliases to the computation, which need not itself contain make_date. + val inherited = + if ((output.deterministic || rowDemand.isDefined) && + inheritedInputs.contains(output.toAttribute)) { + output.references.toSeq + } else { + Nil + } + direct ++ inherited + }) + // A lazy row consumer can skip even the date's required arguments on unread rows. + val deferredDateInputs = + dateInputs.intersect(if (rowDemand.isDefined) project.child.outputSet else skipped) + if (skipped.isEmpty && deferredDateInputs.isEmpty) { None } else { - Some(ProjectionDemand(projectionDemand.map(_.boundary).getOrElse("PROJECT"), skipped)) + Some( + ProjectionDemand( + projectionDemand.map(_.boundary).getOrElse("PROJECT"), + skipped, + deferredDateInputs)) } case _ if exchange || localMaterialization || join => None case _ => projectionDemand @@ -202,23 +234,33 @@ private[rules] object MakeDateEvaluation { val visited = node.children.map(visit(_, childDemand, childProjectionDemand)) val childReason = visited.flatMap(_._2).headOption val ownReason = reason(node).orElse { - if (containsMakeDate(original.expressions)) { - rowDemand - .map(boundary => s"make_date requires Spark evaluation below $boundary") - .orElse(original match { - case project: ProjectExec => - projectionDemand - .filter(demand => - project.projectList.exists(output => - output.deterministic && demand.deferred.contains(output.toAttribute) && - containsMakeDate(Seq(output)))) - .map(demand => s"make_date requires Spark evaluation below ${demand.boundary}") - // Join keys/conditions and aggregate updates can be skipped inside the operator. - case _: BaseJoinExec | _: BaseAggregateExec | _: TakeOrderedAndProjectExec => - Some(s"make_date requires Spark evaluation in ${original.nodeName}") - case _ => None - }) + val evaluatesDate = containsMakeDate(original.expressions) + val rowReason = if (evaluatesDate) { + rowDemand.map(boundary => s"make_date requires Spark evaluation below $boundary") } else None + rowReason.orElse(original match { + 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)) + val deferredArgument = (output.deterministic || rowDemand.isDefined) && + demand.makeDateInputs.contains(output.toAttribute) && + !CometMakeDate.canEvaluateEagerly(computation) + deferredDate || deferredArgument + }) + .map(demand => s"make_date requires Spark evaluation below ${demand.boundary}") + // Join keys/conditions and aggregate updates can be skipped inside the operator. + case _: BaseJoinExec | _: BaseAggregateExec | _: TakeOrderedAndProjectExec + if evaluatesDate => + Some(s"make_date requires Spark evaluation in ${original.nodeName}") + case _ => None + }) } val fallbackReason = ownReason.orElse(if (exchange) None else childReason) val children = visited.map { case (child, _) => diff --git a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala index 4e74279b074..5d845b60e5d 100644 --- a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala @@ -72,7 +72,9 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH spark .range(2) .coalesce(1) - .selectExpr("CAST(CASE WHEN id = 0 THEN 1 ELSE 13 END AS INT) AS m") + .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) } @@ -113,6 +115,24 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH } } + // 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 full collection must still evaluate the second row in the native expression. val full = sql(unlimited) checkCometOperators(stripAQEPlan(full.queryExecution.executedPlan)) 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 d5e75acb360..e69498aa1e8 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -358,6 +358,74 @@ class CometExecRuleSuite extends CometTestBase { } } + 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) + } + } + + // 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 aggregate arguments retain Spark execution") { withSQLConf( SQLConf.ANSI_ENABLED.key -> "true", From 0db23962399be74e629409a9993388133ad17a22 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 28 Oct 2026 05:45:23 +0000 Subject: [PATCH 06/14] fix: preserve deferred siblings of eager make_date outputs --- docs/source/user-guide/latest/expressions.md | 6 +- .../comet/rules/MakeDateEvaluation.scala | 72 +++++++++++++------ .../comet/rules/CometExecRuleSuite.scala | 67 +++++++++++++++++ 3 files changed, 120 insertions(+), 25 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 684c7519a0f..86449560fea 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -727,13 +727,15 @@ and ordered limits; and projections below a window group limit. A parent project 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 to remain native when their operands are safe to evaluate. These safeguards also preserve -NULL short-circuiting when a date argument is computed in an earlier projection. +NULL short-circuiting when a date argument is computed in an earlier projection. Selecting +the date eagerly does not force conditionally used sibling expressions; those computations +also retain Spark evaluation when needed. Only affected parts of the plan fall back. An ordinary native scan can remain native. Exchanges materialize their inputs independently, and local sorts and aggregates can also separate date evaluation from an outer limit. -Users who accept eagerly evaluating dates and their projected arguments can set +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 diff --git a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala index c4f5e617519..9ca4a2cf8cf 100644 --- a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala +++ b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala @@ -49,7 +49,8 @@ private[rules] object MakeDateEvaluation { private case class ProjectionDemand( boundary: String, deferred: AttributeSet, - makeDateInputs: AttributeSet) + // Deferred inputs whose producers need guarding even when they contain no make_date. + guardedInputs: AttributeSet) // Attributes may stand for deferred expressions in a lower projection, not materialized values. private def harmlessValue(expression: Expression): Boolean = @@ -133,19 +134,20 @@ private[rules] object MakeDateEvaluation { def visit( node: SparkPlan, rowDemand: Option[String], - projectionDemand: Option[ProjectionDemand]): (SparkPlan, Option[String]) = { + projectionDemand: Option[ProjectionDemand]): (SparkPlan, Option[String], AttributeSet) = { node match { // Materialized stages are independent producers. Their plans have already been guarded. - case _: QueryStageExec | _: ReusedExchangeExec => return (node, None) + 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) = visit(wrapper.originalPlan, rowDemand, projectionDemand) - return if (childReason.isDefined) (child, childReason) - else (wrapper.copy(originalPlan = child), None) + val (child, childReason, dateOutputs) = + visit(wrapper.originalPlan, rowDemand, projectionDemand) + return if (childReason.isDefined) (child, childReason, dateOutputs) + else (wrapper.copy(originalPlan = child), None, dateOutputs) case sink: CometSinkPlaceHolder => - val (child, childReason) = visit(sink.child, rowDemand, projectionDemand) - return if (childReason.isDefined) (child, childReason) - else (sink.withNewChildren(Seq(child)), None) + val (child, childReason, dateOutputs) = visit(sink.child, rowDemand, projectionDemand) + return if (childReason.isDefined) (child, childReason, dateOutputs) + else (sink.withNewChildren(Seq(child)), None, dateOutputs) case _ => } @@ -160,6 +162,7 @@ private[rules] object MakeDateEvaluation { case broadcast: CometBroadcastExchangeExec => broadcast.originalPlan case _ => node } + val evaluatesDate = containsMakeDate(original.expressions) val exchange = original.isInstanceOf[Exchange] val join = original.isInstanceOf[BaseJoinExec] val localMaterialization = original match { @@ -193,7 +196,7 @@ private[rules] object MakeDateEvaluation { ProjectionDemand( "FILTER", filter.child.outputSet, - projectionDemand.map(_.makeDateInputs).getOrElse(AttributeSet.empty))) + projectionDemand.map(_.guardedInputs).getOrElse(AttributeSet.empty))) case project: ProjectExec => val deferred = projectionDemand.map(_.deferred).getOrElse(AttributeSet.empty) val eager = project.projectList @@ -203,13 +206,16 @@ private[rules] object MakeDateEvaluation { // remain inlined in a consumer's conditional, even across intervening aliases. val skipped = project.child.outputSet -- project.usedInputs -- AttributeSet(eager) val inheritedInputs = - projectionDemand.map(_.makeDateInputs).getOrElse(AttributeSet.empty) - val dateInputs = AttributeSet(project.projectList.flatMap { output => + 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 inputsToGuard = AttributeSet(project.projectList.flatMap { output => val direct = output.collect { case date: MakeDate => date }.flatMap(_.references) // Follow aliases to the computation, which need not itself contain make_date. val inherited = if ((output.deterministic || rowDemand.isDefined) && - inheritedInputs.contains(output.toAttribute)) { + guardedOutputs.contains(output.toAttribute)) { output.references.toSeq } else { Nil @@ -217,24 +223,42 @@ private[rules] object MakeDateEvaluation { direct ++ inherited }) // A lazy row consumer can skip even the date's required arguments on unread rows. - val deferredDateInputs = - dateInputs.intersect(if (rowDemand.isDefined) project.child.outputSet else skipped) - if (skipped.isEmpty && deferredDateInputs.isEmpty) { + val guardedChildInputs = + inputsToGuard.intersect(if (rowDemand.isDefined) project.child.outputSet else skipped) + if (skipped.isEmpty && guardedChildInputs.isEmpty) { None } else { Some( ProjectionDemand( projectionDemand.map(_.boundary).getOrElse("PROJECT"), skipped, - deferredDateInputs)) + guardedChildInputs)) } case _ if exchange || localMaterialization || join => None case _ => projectionDemand } val visited = node.children.map(visit(_, childDemand, childProjectionDemand)) val childReason = visited.flatMap(_._2).headOption + val childDateOutputs = AttributeSet(visited.flatMap(_._3)) + // 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 || join) { + 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) + } + } val ownReason = reason(node).orElse { - val evaluatesDate = containsMakeDate(original.expressions) val rowReason = if (evaluatesDate) { rowDemand.map(boundary => s"make_date requires Spark evaluation below $boundary") } else None @@ -249,10 +273,12 @@ private[rules] object MakeDateEvaluation { } val deferredDate = output.deterministic && demand.deferred.contains(output.toAttribute) && containsMakeDate(Seq(output)) - val deferredArgument = (output.deterministic || rowDemand.isDefined) && - demand.makeDateInputs.contains(output.toAttribute) && + val deferredComputation = (output.deterministic || rowDemand.isDefined) && + (demand.guardedInputs.contains(output.toAttribute) || + (dateDerivedOutputs.nonEmpty && demand.deferred.contains( + output.toAttribute))) && !CometMakeDate.canEvaluateEagerly(computation) - deferredDate || deferredArgument + deferredDate || deferredComputation }) .map(demand => s"make_date requires Spark evaluation below ${demand.boundary}") // Join keys/conditions and aggregate updates can be skipped inside the operator. @@ -263,7 +289,7 @@ private[rules] object MakeDateEvaluation { }) } val fallbackReason = ownReason.orElse(if (exchange) None else childReason) - val children = visited.map { case (child, _) => + val children = visited.map { case (child, _, _) => if (join && fallbackReason.isDefined) restoreBroadcast(child, fallbackReason.get) else child } @@ -281,7 +307,7 @@ private[rules] object MakeDateEvaluation { } copyLogicalLinks(node, rebuilt) fallbackReason.foreach(rebuilt.setTagValue(sparkEvaluation, _)) - (rebuilt, fallbackReason) + (rebuilt, fallbackReason, dateDerivedOutputs) } // Dataset.rdd adds an object-producing root, and its caller may read only a few objects. 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 e69498aa1e8..0ed471a51da 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -324,6 +324,73 @@ class CometExecRuleSuite extends CometTestBase { } 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) + } + } } } From d2ba4421976c3ddcc322ff6b4cd2bd1136d6f093 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 28 Oct 2026 17:38:49 +0000 Subject: [PATCH 07/14] fix: preserve deferred make_date arguments through broadcast joins --- docs/source/user-guide/latest/expressions.md | 7 ++- .../comet/rules/MakeDateEvaluation.scala | 40 ++++++++++-- .../comet/CometTemporalExpressionSuite.scala | 24 +++++++ .../comet/rules/CometExecRuleSuite.scala | 62 ++++++++++++++++++- 4 files changed, 124 insertions(+), 9 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 86449560fea..9de45654889 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -727,9 +727,10 @@ and ordered limits; and projections below a window group limit. A parent project 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 to remain native when their operands are safe to evaluate. These safeguards also preserve -NULL short-circuiting when a date argument is computed in an earlier projection. Selecting -the date eagerly does not force conditionally used sibling expressions; those computations -also retain Spark evaluation when needed. +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. +Selecting the date eagerly does not force conditionally used sibling expressions; those +computations also retain Spark evaluation when needed. Only affected parts of the plan fall back. An ordinary native scan can remain native. Exchanges materialize their inputs independently, and local sorts and aggregates can also diff --git a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala index 9ca4a2cf8cf..3091b76f445 100644 --- a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala +++ b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala @@ -20,6 +20,7 @@ package org.apache.comet.rules import org.apache.spark.sql.catalyst.expressions.{Add, Alias, And, Attribute, AttributeSet, BinaryComparison, Cast, ConditionalExpression, CreateArray, CreateNamedStruct, DateAdd, DateDiff, DateFormatClass, DateSub, EqualNullSafe, Expression, Greatest, HashExpression, IsNotNull, IsNull, Least, MakeDate, Multiply, Not, Or, SortOrder, Subtract, UnixDate} +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 @@ -27,7 +28,7 @@ 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 +import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec} import org.apache.comet.CometConf import org.apache.comet.serde.CometMakeDate @@ -165,6 +166,15 @@ private[rules] object MakeDateEvaluation { 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 @@ -234,15 +244,35 @@ private[rules] object MakeDateEvaluation { skipped, guardedChildInputs)) } - case _ if exchange || localMaterialization || join => None + case _ if exchange || localMaterialization => None case _ => projectionDemand } - val visited = node.children.map(visit(_, childDemand, childProjectionDemand)) + 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 forwarded = original.children(index).outputSet -- required + demand.copy( + deferred = demand.deferred.intersect(forwarded), + guardedInputs = demand.guardedInputs.intersect(forwarded)) + } + } else { + None + } + visit(child, childDemand, demand) + } val childReason = visited.flatMap(_._2).headOption - val childDateOutputs = AttributeSet(visited.flatMap(_._3)) + 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 || join) { + val dateDerivedOutputs = if (exchange || localMaterialization) { AttributeSet.empty } else { original match { diff --git a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala index 5d845b60e5d..276a00ea6b5 100644 --- a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala @@ -29,6 +29,7 @@ import org.apache.spark.sql.catalyst.util.DateTimeUtils import org.apache.spark.sql.comet.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} @@ -133,6 +134,29 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH 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)) 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 0ed471a51da..f4ccbb2241d 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -32,7 +32,7 @@ import org.apache.spark.sql.catalyst.{FunctionIdentifier, InternalRow} 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, Final, Min, Partial, PartialMerge} import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke} -import org.apache.spark.sql.catalyst.optimizer.{BuildRight, 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, SinglePartition} @@ -476,6 +476,66 @@ class CometExecRuleSuite extends CometTestBase { } } + // 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( """ From 9cfad7e4bd66b583bab7b089047db09d59a1c357 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 28 Oct 2026 18:47:14 +0000 Subject: [PATCH 08/14] fix: preserve make_date inputs skipped by broadcast joins --- docs/source/user-guide/latest/expressions.md | 2 + .../comet/rules/MakeDateEvaluation.scala | 42 +++++++--- .../comet/rules/CometExecRuleSuite.scala | 80 +++++++++++++++++++ 3 files changed, 112 insertions(+), 12 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 9de45654889..d7cdf86c821 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -729,6 +729,8 @@ it. Comet preserves that evaluation across projections while allowing eager date to remain native when their operands are safe to evaluate. 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. diff --git a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala index 3091b76f445..6d31ca3d43f 100644 --- a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala +++ b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala @@ -51,7 +51,9 @@ private[rules] object MakeDateEvaluation { boundary: String, deferred: AttributeSet, // Deferred inputs whose producers need guarding even when they contain no make_date. - guardedInputs: AttributeSet) + 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 = @@ -206,7 +208,8 @@ private[rules] object MakeDateEvaluation { ProjectionDemand( "FILTER", filter.child.outputSet, - projectionDemand.map(_.guardedInputs).getOrElse(AttributeSet.empty))) + 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 @@ -220,29 +223,40 @@ private[rules] object MakeDateEvaluation { // A consumer can require the date but still skip another output of this projection. val guardedOutputs = inheritedInputs ++ (if (evaluatesDate) deferred else AttributeSet.empty) - val inputsToGuard = AttributeSet(project.projectList.flatMap { output => - val direct = output.collect { case date: MakeDate => date }.flatMap(_.references) - // Follow aliases to the computation, which need not itself contain make_date. - val inherited = + 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 } - direct ++ inherited - }) + }) + // 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) { + if (skipped.isEmpty && guardedChildInputs.isEmpty && dateInputs.isEmpty) { None } else { Some( ProjectionDemand( projectionDemand.map(_.boundary).getOrElse("PROJECT"), skipped, - guardedChildInputs)) + guardedChildInputs, + dateInputs)) } case _ if exchange || localMaterialization => None case _ => projectionDemand @@ -255,10 +269,14 @@ private[rules] object MakeDateEvaluation { // 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 forwarded = original.children(index).outputSet -- required + val streamedOutputs = original.children(index).outputSet + val forwarded = streamedOutputs -- required demand.copy( deferred = demand.deferred.intersect(forwarded), - guardedInputs = demand.guardedInputs.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 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 f4ccbb2241d..86042044bbf 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -553,6 +553,86 @@ class CometExecRuleSuite extends CometTestBase { } } + 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 aggregate arguments retain Spark execution") { withSQLConf( SQLConf.ANSI_ENABLED.key -> "true", From 8d766d7cc11e0656afeeb00e127c59396be7d29d Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 28 Oct 2026 19:11:59 +0000 Subject: [PATCH 09/14] fix: retain native make_date shuffle keys below limits --- docs/source/user-guide/latest/expressions.md | 5 ++- .../comet/rules/MakeDateEvaluation.scala | 3 +- .../comet/rules/CometExecRuleSuite.scala | 40 +++++++++++++++++++ 3 files changed, 45 insertions(+), 3 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index d7cdf86c821..f92b305701d 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -735,8 +735,9 @@ Selecting the date eagerly does not force conditionally used sibling expressions computations also retain Spark evaluation when needed. Only affected parts of the plan fall back. An ordinary native scan can remain native. -Exchanges materialize their inputs independently, and local sorts and aggregates can also -separate date evaluation from an outer limit. +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 diff --git a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala index 6d31ca3d43f..0d2ab9e3168 100644 --- a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala +++ b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala @@ -307,7 +307,8 @@ private[rules] object MakeDateEvaluation { } } val ownReason = reason(node).orElse { - val rowReason = if (evaluatesDate) { + // An exchange evaluates its partition keys for every input before a consumer reads rows. + val rowReason = if (evaluatesDate && !exchange) { rowDemand.map(boundary => s"make_date requires Spark evaluation below $boundary") } else None rowReason.orElse(original match { 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 86042044bbf..af8adc0047e 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -46,6 +46,7 @@ import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAg import org.apache.spark.sql.execution.datasources.v2.BatchScanExec 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} @@ -209,6 +210,45 @@ class CometExecRuleSuite extends CometTestBase { } } + 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", From 8e1961a99028dea6357f7a66d7efe4c72e8a67b0 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 28 Oct 2026 19:15:44 +0000 Subject: [PATCH 10/14] fix: retain native count for eager make_date arguments --- docs/source/user-guide/latest/expressions.md | 4 ++ .../comet/rules/MakeDateEvaluation.scala | 17 +++++-- .../comet/rules/CometExecRuleSuite.scala | 46 +++++++++++++++++-- 3 files changed, 61 insertions(+), 6 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index f92b305701d..47c67d60a14 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -734,6 +734,10 @@ 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. +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. 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 diff --git a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala index 0d2ab9e3168..c8c03d63938 100644 --- a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala +++ b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala @@ -20,6 +20,7 @@ package org.apache.comet.rules import org.apache.spark.sql.catalyst.expressions.{Add, Alias, And, Attribute, AttributeSet, BinaryComparison, Cast, ConditionalExpression, CreateArray, CreateNamedStruct, DateAdd, DateDiff, DateFormatClass, DateSub, EqualNullSafe, Expression, Greatest, HashExpression, IsNotNull, IsNull, Least, 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._ @@ -47,6 +48,14 @@ private[rules] object MakeDateEvaluation { private def containsMakeDate(expressions: Seq[Expression]): Boolean = expressions.exists(_.exists(_.isInstanceOf[MakeDate])) + 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, @@ -330,9 +339,11 @@ private[rules] object MakeDateEvaluation { deferredDate || deferredComputation }) .map(demand => s"make_date requires Spark evaluation below ${demand.boundary}") - // Join keys/conditions and aggregate updates can be skipped inside the operator. - case _: BaseJoinExec | _: BaseAggregateExec | _: TakeOrderedAndProjectExec - if evaluatesDate => + 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 }) 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 af8adc0047e..ccf97c3f29a 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -30,7 +30,7 @@ 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, 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, Final, Min, Partial, PartialMerge} +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.{BuildLeft, BuildRight, NormalizeNaNAndZero} import org.apache.spark.sql.catalyst.plans.Inner @@ -673,12 +673,47 @@ class CometExecRuleSuite extends CometTestBase { } } + 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") { - withParquetTable(Seq((2024, 1), (2024, 13), (2025, 1)), "make_date_aggregate") { + 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))", @@ -688,7 +723,12 @@ class CometExecRuleSuite extends CometTestBase { "sort_array(collect_list(unix_date(make_date(_1, _2, 1))) " + "FILTER (WHERE _2 <= 12))", classOf[ObjectHashAggregateExec], - Row(Seq(19723, 20089)))) + 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) From 121e3279285882751a2de90f0303f5f1ad5f9556 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 28 Oct 2026 19:18:23 +0000 Subject: [PATCH 11/14] fix: follow eager date extractors through projections --- docs/source/user-guide/latest/expressions.md | 5 +++-- .../org/apache/comet/rules/MakeDateEvaluation.scala | 5 +++-- .../org/apache/comet/rules/CometExecRuleSuite.scala | 12 ++++++++++++ 3 files changed, 18 insertions(+), 4 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 47c67d60a14..3313d8b1807 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -725,8 +725,9 @@ 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 -to remain native when their operands are safe to evaluate. These safeguards also preserve +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 diff --git a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala index c8c03d63938..911bd0da6b4 100644 --- a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala +++ b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala @@ -19,7 +19,7 @@ package org.apache.comet.rules -import org.apache.spark.sql.catalyst.expressions.{Add, Alias, And, Attribute, AttributeSet, BinaryComparison, Cast, ConditionalExpression, CreateArray, CreateNamedStruct, DateAdd, DateDiff, DateFormatClass, DateSub, EqualNullSafe, Expression, Greatest, HashExpression, IsNotNull, IsNull, Least, MakeDate, Multiply, Not, Or, SortOrder, Subtract, UnixDate} +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, 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 @@ -80,7 +80,8 @@ private[rules] object MakeDateEvaluation { val required = conditional.alwaysEvaluatedInputs if (conditional.children.filterNot(required.contains).forall(harmlessValue)) required else Nil - case _: Alias | _: Cast | _: IsNull | _: IsNotNull | _: Not | _: UnixDate => + case _: Alias | _: Cast | _: IsNull | _: IsNotNull | _: Not | _: UnixDate | + _: GetDateField => expression.children case _: EqualNullSafe | _: CreateArray | _: CreateNamedStruct | _: Greatest | _: Least | _: HashExpression[_] => 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 ccf97c3f29a..b7ac5a8f81a 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -437,11 +437,23 @@ class CometExecRuleSuite extends CometTestBase { // 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 ( From aa190d81cdfdde9bbaf651759a8ea3ec43044f7a Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Wed, 28 Oct 2026 20:32:23 +0000 Subject: [PATCH 12/14] test: make shuffle fixture tuple explicit for Scala 2.12 --- .../test/scala/org/apache/comet/rules/CometExecRuleSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 b7ac5a8f81a..4e8af005fc0 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -218,7 +218,7 @@ class CometExecRuleSuite extends CometTestBase { 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) + 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") From 35fa3cb6956658da2bf5e3c162bb8f9d2e5285a4 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Thu, 29 Oct 2026 00:38:06 +0000 Subject: [PATCH 13/14] fix: keep safe non-ANSI date projections native --- docs/source/user-guide/latest/expressions.md | 7 ++ .../comet/rules/MakeDateEvaluation.scala | 45 ++++++++++- .../comet/rules/CometExecRuleSuite.scala | 77 +++++++++++++++++-- 3 files changed, 117 insertions(+), 12 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 3313d8b1807..7f2f2b5bcce 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -735,6 +735,13 @@ 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. + 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. diff --git a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala index 911bd0da6b4..35690ee1b7d 100644 --- a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala +++ b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala @@ -19,7 +19,9 @@ package org.apache.comet.rules -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, MakeDate, Multiply, Not, Or, SortOrder, Subtract, UnixDate} +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 @@ -30,6 +32,7 @@ import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec, QuerySt 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 import org.apache.comet.serde.CometMakeDate @@ -48,6 +51,34 @@ private[rules] object MakeDateEvaluation { 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. @@ -318,7 +349,11 @@ private[rules] object MakeDateEvaluation { } val ownReason = reason(node).orElse { // An exchange evaluates its partition keys for every input before a consumer reads rows. - val rowReason = if (evaluatesDate && !exchange) { + 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 rowReason.orElse(original match { @@ -331,12 +366,14 @@ private[rules] object MakeDateEvaluation { case other => other } val deferredDate = output.deterministic && - demand.deferred.contains(output.toAttribute) && containsMakeDate(Seq(output)) + 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))) && - !CometMakeDate.canEvaluateEagerly(computation) + !canEvaluateProjectionEagerly(computation) deferredDate || deferredComputation }) .map(demand => s"make_date requires Spark evaluation below ${demand.boundary}") 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 4e8af005fc0..d5e0ca71f73 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -171,11 +171,17 @@ class CometExecRuleSuite extends CometTestBase { ("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)) { + for (ansi <- Seq(false, true); boundedYear <- Seq(false, true)) { def project: SparkPlan = ProjectExec( Seq( input.output.head, - Alias(MakeDate(input.output.head, Literal(1), Literal(1), ansi), "date")()), + 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) @@ -183,13 +189,16 @@ class CometExecRuleSuite extends CometTestBase { val child = if (reused) applyCometExecRule(project) else project val converted = applyCometExecRule(consume(child)) for (result <- Seq(converted, applyCometExecRule(converted))) { - withClue(s"ansi=$ansi reused=$reused $boundary\n${result.treeString}") { - assert(countOperators(result, classOf[ProjectExec]) == 1) - assert(countOperators(result, classOf[CometProjectExec]) == 0) + val fallback = ansi || !boundedYear + withClue( + s"ansi=$ansi boundedYear=$boundedYear reused=$reused $boundary\n${result.treeString}") { + assert(countOperators(result, classOf[ProjectExec]) == (if (fallback) 1 else 0)) assert( - new ExtendedExplainInfo() - .getFallbackReasons(result) - .exists(_.contains(s"make_date requires Spark evaluation below $boundary"))) + 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) } } } @@ -210,6 +219,58 @@ class CometExecRuleSuite extends CometTestBase { } } + 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)) + withParquetTable(rows, "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) + checkSparkAnswer(deferred) + assert( + countOperators(deferred.queryExecution.executedPlan, classOf[CometProjectExec]) == 2) + + // 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) + checkSparkAnswer(df) + assert(countOperators(df.queryExecution.executedPlan, classOf[ProjectExec]) == 1) + assert(countOperators(df.queryExecution.executedPlan, classOf[CometProjectExec]) == 0) + } + } + + // 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 shuffle keys remain eager below a limit") { withSQLConf( SQLConf.ANSI_ENABLED.key -> "true", From 47f66c44227d2c5ccbd7430c4d03e5a5c820d1e6 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Thu, 29 Oct 2026 00:48:23 +0000 Subject: [PATCH 14/14] fix: reuse collect-limit caps for guarded date projections --- docs/source/user-guide/latest/expressions.md | 4 + .../apache/comet/rules/CometExecRule.scala | 52 ++++--- .../comet/rules/MakeDateEvaluation.scala | 130 ++++++++++++------ .../comet/CometTemporalExpressionSuite.scala | 9 +- .../comet/rules/CometExecRuleSuite.scala | 124 +++++++++++++---- 5 files changed, 227 insertions(+), 92 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 7f2f2b5bcce..ab604528222 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -747,6 +747,10 @@ because Spark checks that argument on every input row. Filtered counts and count 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. 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 ee317f7799a..9616b7d492d 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -754,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 @@ -769,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 diff --git a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala index 35690ee1b7d..a3fbf4c06f0 100644 --- a/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala +++ b/spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala @@ -34,13 +34,43 @@ 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 +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) @@ -178,18 +208,22 @@ private[rules] object MakeDateEvaluation { def visit( node: SparkPlan, rowDemand: Option[String], - projectionDemand: Option[ProjectionDemand]): (SparkPlan, Option[String], AttributeSet) = { + 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) + 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) + 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 _ => @@ -302,6 +336,12 @@ private[rules] object MakeDateEvaluation { 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 @@ -322,7 +362,7 @@ private[rules] object MakeDateEvaluation { } else { None } - visit(child, childDemand, demand) + visit(child, childDemand, demand, childCollectLimit) } val childReason = visited.flatMap(_._2).headOption val childDateOutputs = AttributeSet(visited.zipWithIndex.flatMap { @@ -347,45 +387,43 @@ private[rules] object MakeDateEvaluation { case _ => childDateOutputs.intersect(original.outputSet) } } - val ownReason = reason(node).orElse { - // 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 - rowReason.orElse(original match { - 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 - }) + // 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 + val operatorReason = original match { + 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) @@ -404,6 +442,12 @@ private[rules] object MakeDateEvaluation { 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) } diff --git a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala index 276a00ea6b5..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,7 @@ 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.CometProjectExec +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 @@ -91,10 +91,11 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH assert(expected == Seq(Row(19723))) val df = sql(query) val plan = stripAQEPlan(df.queryExecution.executedPlan) - assert(plan.collect { case p: ProjectExec => p }.nonEmpty, plan.treeString) - assert(plan.collect { case p: CometProjectExec => p }.isEmpty, plan.treeString) + 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() + !new ExtendedExplainInfo() .getFallbackReasons(plan) .exists(_.contains("make_date")), plan.treeString) 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 d5e0ca71f73..d339ccd3480 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -187,9 +187,12 @@ class CometExecRuleSuite extends CometTestBase { 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 converted = applyCometExecRule(consume(child)) + val consumer = consume(child) + val converted = applyCometExecRule(consumer) for (result <- Seq(converted, applyCometExecRule(converted))) { - val fallback = ansi || !boundedYear + // 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)) @@ -225,34 +228,47 @@ class CometExecRuleSuite extends CometTestBase { 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)) - withParquetTable(rows, "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])) + 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(""" + 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) - checkSparkAnswer(deferred) - assert( - countOperators(deferred.queryExecution.executedPlan, classOf[CometProjectExec]) == 2) - - // 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) - checkSparkAnswer(df) - assert(countOperators(df.queryExecution.executedPlan, classOf[ProjectExec]) == 1) - assert(countOperators(df.queryExecution.executedPlan, classOf[CometProjectExec]) == 0) + checkSparkAnswerAndOperator(deferred, includeClasses = Seq(classOf[CometProjectExec])) + assert( + countOperators(deferred.queryExecution.executedPlan, classOf[CometProjectExec]) == 2) + + // 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. @@ -271,6 +287,68 @@ class CometExecRuleSuite extends CometTestBase { } } + 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",