pingzh commented on code in PR #7117:
URL: https://github.com/apache/datafusion-comet/pull/7117#discussion_r4332583753
##########
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:
Agreed; I've opened #7144 to move these SQL-level regressions into fixtures
as a separate cleanup, related to #6625. The existing tests already cover the
behavior.
The projection fixture should enable `ConstantFolding` explicitly to retain
coverage with the default optimizer. We'll preserve the AQE matrix, fallback
reasons, nullable `datediff` case, and eager native control, and retain focused
Scala assertions where native-scan retention or exact operator counts matter.
##########
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:
The native iterator already [returns
`unsafeRow.copy()`](https://github.com/apache/datafusion-comet/blob/2ece37f1d00cb8a4e0719b51957f7963df8d5426/spark/src/main/scala/org/apache/comet/NativeColumnarToRowConverter.scala#L143),
added in #3367 before this PR. The iterator comments describing zero-copy
output are stale: the returned rows are copied and survive converter closure,
so another copy in the native broadcast transition would be redundant.
As you noted, this PR's `restoreBroadcast` also explicitly uses the JVM
transition. Resolving this thread because the native path already has the
required copy.
--
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]
