comphead commented on code in PR #7117: URL: https://github.com/apache/datafusion-comet/pull/7117#discussion_r4332013624
########## spark/src/main/scala/org/apache/comet/rules/MakeDateEvaluation.scala: ########## @@ -0,0 +1,236 @@ +/* + * 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.{Alias, And, Attribute, AttributeSet, BinaryComparison, Cast, ConditionalExpression, CreateArray, CreateNamedStruct, DateAdd, DateDiff, DateFormatClass, DateSub, EqualNullSafe, Expression, IsNotNull, IsNull, MakeDate, Not, Or, 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.CometConf +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 { Review Comment: #5533 and #5972 (both open) build the same mechanism in general form: a `RequiresSparkEvaluationMask` serde trait, `QueryPlanSerde.evaluationMaskName`, a `CometExecRule.preserveEvaluationMasks` pre-pass and one `spark.comet.exec.preserveEvaluationMasks.enabled` opt-out. This file, `CometMakeDate.unsupportedParentReason` and `spark.comet.exec.makeDate.operatorFallback.enabled` look like a make_date-only copy of that design. The shapes match closely: the tag check as the first `convertNode` case, restoring `originalPlan`, and the eager-child table in `eagerReferences`, which #5972 has as `QueryPlanSerde.eagerlyEvaluatedChildren`. The two passes also disagree about which consumers skip rows. This PR protects inner joins, `SAMPLE`, `mapPartitions` and `Dataset.rdd` roots, which #5972 does not. #5972 gates projection deferral on whole-stage codegen and treats filter-referenced inputs as eager, which this PR does not. Whichever lands second either leaves `unbase64` and `next_day` unprotected under those consumers or makes `make_date` fall back more than it needs to. The two PRs also conflict in `CometConf.scala` and `CometExecRule.scala` by my `git merge-tree` run. Both new keys are `CATEGORY_EXEC` and not `internal()`, so once released they fall under the compatibility guarantee in `config_conventions.md`, and merging them later would need a deprecation cycle. Could `MakeDate` enroll through `RequiresSparkEvaluationMask`, with the extra consumers moved into the shared pass? Agreeing with sunchao on which PR lands the mechanism first would settle the order. The shared pass also returns early when no enrolled expression is in the plan, while `preserve` here has no early exit. ########## spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala: ########## @@ -125,6 +128,427 @@ 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), (6000000) 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), (6000000) AS t(y)") + val right = createSparkPlan(spark, "SELECT * FROM VALUES (2024) AS t(y)") + val date = MakeDate(left.output.head, Literal(1), Literal(1), failOnError = false) + val projection = ProjectExec(Seq(left.output.head, Alias(date, "date")()), left) + def sort(child: SparkPlan): SparkPlan = + SortExec(Seq(SortOrder(child.output.head, Ascending)), false, child) + val join = SortMergeJoinExec( + Seq(left.output.head), + Seq(right.output.head), + Inner, + None, + sort(applyCometExecRule(projection)), + sort(right)) + val converted = applyCometExecRule(join) + for (result <- Seq(converted, applyCometExecRule(converted))) { + assert(countOperators(result, classOf[SortMergeJoinExec]) == 1, result.treeString) + assert(countOperators(result, classOf[CometProjectExec]) == 0, result.treeString) + assert(countOperators(result, classOf[ProjectExec]) == 1, result.treeString) + assert( + new ExtendedExplainInfo() + .getFallbackReasons(result) + .exists(_.contains("make_date requires Spark evaluation below JOIN"))) + } + } + } + + test("make_date remains deferred across conditional projections") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.DATETIME_JAVA8API_ENABLED.key -> "false") { + val selectedDate = java.sql.Date.valueOf("2024-01-01") + val rows = Seq((6000000, null.asInstanceOf[java.sql.Date]), (2024, 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, 1, 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, _2 AS nullable_date, make_date(_1, 1, 1) AS dt + | FROM make_date_deferred + |) + |""".stripMargin) { expected => + assert(expected.size == 2) + assert(expected.map(_.get(1)).toSet == Set(null, 0)) + } + assertDeferredProjection(plan) + } + } + + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val plan = checkProjection(""" + |SELECT r, date_add(dt, 1) AS d, IF(r < 0, 0, r + 1) AS repeated_r FROM ( + | SELECT rand(7L) AS r, make_date(_1, 1, 1) AS dt + | FROM make_date_deferred WHERE _1 < 3000 + |) + |""".stripMargin) { expected => + assert(expected.map(_.getDate(1)) == Seq(java.sql.Date.valueOf("2024-01-02"))) + } + assert(countOperators(plan, classOf[CometProjectExec]) == 2, plan.treeString) + assert(countOperators(plan, classOf[CometNativeScanExec]) == 1, plan.treeString) + assert( + !new ExtendedExplainInfo() + .getFallbackReasons(plan) + .exists(_.contains("make_date requires Spark evaluation below PROJECT"))) + } + } + } + } + + test("make_date aggregate arguments retain Spark execution") { Review Comment: This test and `make_date remains deferred across conditional projections` (line 218) only run SQL over a Parquet table and check answers and fallback reasons, which `CometSqlFileTestSuite` can express: `query expect_fallback(make_date requires Spark evaluation in HashAggregate)` and the `ObjectHashAggregate` equivalent, `-- Config: spark.sql.adaptive.enabled=false`, `-- ConfigMatrix` for the projection cases, and a plain `query` for the native control. Moving them next to `make_date_overflow.sql` would keep the SQL-level regressions in one place. The synthetic-plan tests in this suite and the `rdd`, `mapPartitions` and empty-`IN` tests need Scala and would stay. The `CometNativeScanExec` count assertion would be lost, which is a small loss. ########## spark/src/main/scala/org/apache/comet/serde/datetime.scala: ########## @@ -436,10 +436,70 @@ object CometNextDay extends CometExpressionSerde[NextDay] with CodegenDispatchFa object CometMakeDate extends CometExpressionSerde[MakeDate] { + private val sharedArgumentReason = + "make_date requires Spark evaluation when a computed argument after a nullable " + + "argument cannot preserve Spark's shared-subexpression evaluation." + + override def getUnsupportedReasons(): Seq[String] = Seq(sharedArgumentReason) + + override def getSupportLevel(expr: MakeDate): SupportLevel = { + val firstNullable = expr.children.indexWhere(_.nullable) + val hasComputedMaskedArgument = firstNullable >= 0 && + expr.children.drop(firstNullable + 1).exists(!_.isInstanceOf[LeafExpression]) + if (hasComputedMaskedArgument && !canShortCircuitNulls(expr)) { + Unsupported(Some(sharedArgumentReason)) + } else { + Compatible() + } + } + + /** A native parent must not evaluate a date on rows where Spark skips that child. */ + private[serde] def unsupportedParentReason(expr: Expression): Option[String] = { + var pending: List[(Expression, Option[String])] = (expr, None) :: Nil + while (pending.nonEmpty) { + val (current, parent) = pending.head + pending = pending.tail + if (current.isInstanceOf[MakeDate] && parent.isDefined) { + return Some(s"make_date requires Spark evaluation under ${parent.get}") + } + def pushChildren(mask: Option[String]): Unit = { + current.children.reverseIterator.foreach(child => pending = (child, mask) :: pending) + } + if (parent.isDefined) { + pushChildren(parent) + } else { + current match { + // These native expressions already select the rows passed to their children. + case _: If | _: CaseWhen | _: MakeDate => pushChildren(None) + // COALESCE lowers to CASE and evaluates an earlier child again in its result branch. + case coalesce: Coalesce if coalesce.deterministic => pushChildren(None) + // These parents evaluate every child in Spark, including NULL values. + case _: EqualNullSafe | _: CreateArray | _: CreateNamedStruct => + pushChildren(None) + // These parents evaluate left first and skip right only when left is NULL. + case _: BinaryComparison | _: DateAdd | _: DateSub | _: DateDiff | _: DateFormatClass => + val left = current.children.head + val right = current.children(1) + val rightMask = if (left.nullable) Some(current.nodeName) else None + pending = (left, None) :: (right, rightMask) :: pending + case _: SortOrder => pushChildren(Some(current.nodeName)) + // Empty membership can skip its sole operand in Spark. + case in: In if in.list.isEmpty => pushChildren(Some(current.nodeName)) + case inSet: InSet if inSet.hset.isEmpty => pushChildren(Some(current.nodeName)) + case _ if current.children.size <= 1 => pushChildren(None) + // Other native multi-child expressions may evaluate operands that Spark skips. + case _ => pushChildren(Some(current.nodeName)) Review Comment: With nullable inputs, `WHERE make_date(y, m, d) > DATE '2000-01-01'` should reach `CometFilterExec` as `isnotnull(y) AND isnotnull(m) AND isnotnull(d) AND (make_date(y, m, d) > ...)`, because `MakeDate` is null-intolerant and `InferFiltersFromConstraints` prepends those checks. `CometFilterExec` passes the whole condition to `exprToProto`, and `And` has no case in `unsupportedParentReason`, so I read this arm as masking every child and the filter falling back with `make_date requires Spark evaluation under And`. The rows Spark skips in that shape are the ones with a NULL argument, where the native kernel already returns NULL without validating. The left operand of `And` and `Or` is always evaluated, which `MakeDateEvaluation.eagerReferences` already encodes (`expression.children.take(1)`). Could the guard treat it as eager here too? A `query expect_native(make_date)` case with that `WHERE` on `make_date_boundaries` would pin it, since its NULL-argument rows are the skipped ones. I could not find a `WHERE make_date` query in the PR's tests. I traced this through the sources and did not run it. ########## spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala: ########## @@ -138,7 +138,13 @@ case class CometColumnarToRowExec(child: SparkPlan) numOutputRows) val mode = cometBroadcastExchange.get.mode - val relation = mode.transform(rows, Some(numOutputRows.value)) + // Identity broadcasts retain rows, but both conversion paths reuse the returned row. + // Hashed relations copy row bytes into their own storage while consuming the iterator. + val broadcastRows = mode match { + case IdentityBroadcastMode => rows.map(_.copy()) Review Comment: `CometNativeColumnarToRowExec.doExecuteBroadcast` still calls `mode.transform(rows, ...)` with no copy and then runs `converter.close()` in the `finally`. `NativeColumnarToRowConverter` reuses one `UnsafeRow` that points into native memory, so an `IdentityBroadcastMode` array built there would hold one row after the native memory is released. This PR's `restoreBroadcast` always builds `CometColumnarToRowExec`, so it does not reach that path. Could the same copy go into the native transition, or could this be tracked in an issue so the two broadcast paths stay consistent? I have not reproduced the native case. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
