sunchao commented on code in PR #5394:
URL: https://github.com/apache/datafusion-comet/pull/5394#discussion_r3815326115
##########
spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala:
##########
@@ -797,4 +799,164 @@ class CometExecRuleSuite extends CometTestBase {
}
}
+ /**
+ * Run `sql` with plan-only mode enabled and assert nothing was offloaded to
native. `useV1`
+ * toggles between `USE_V1_SOURCE_LIST=parquet` (V1 `CometScanExec` path) and
+ * `USE_V1_SOURCE_LIST=""` (V2 `CometBatchScanExec` path).
+ */
+ private def runPlanOnlyAndAssertReverted(
+ sql: String,
+ useV1: Boolean = true,
+ aqe: Boolean = true): Unit = {
+ withSQLConf(
+ SQLConf.USE_V1_SOURCE_LIST.key -> (if (useV1) "parquet" else ""),
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+ val executed = spark.sql(sql).queryExecution.executedPlan
+ val cometNodes = stripAQEPlan(executed).collect { case p: CometPlan => p
}
+ assert(
+ cometNodes.isEmpty,
+ s"plan-only mode must not offload; found Comet operators: $cometNodes")
+ }
+ }
+
+ for {
+ useV1 <- Seq(true, false)
+ aqe <- Seq(true, false)
+ } {
+ val label = s"${if (useV1) "V1" else "V2"} scan, AQE=$aqe"
+ test(s"plan-only mode: $label") {
+ withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+ runPlanOnlyAndAssertReverted(
+ "SELECT _2, count(*) FROM tbl GROUP BY _2",
+ useV1 = useV1,
+ aqe = aqe)
+ }
+ }
+ }
+
+ test("plan-only mode: scalar subquery is also reverted") {
+ withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+ runPlanOnlyAndAssertReverted("SELECT _1 FROM tbl WHERE _1 > (SELECT
max(_2) FROM tbl)")
+ }
+ }
+
+ test("plan-only mode: same query with the config off runs on Comet") {
+ withSQLConf(
+ SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "false") {
+ withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+ val plan =
+ spark.sql("SELECT _2, count(*) FROM tbl GROUP BY
_2").queryExecution.executedPlan
+ val cometNodes = stripAQEPlan(plan).collect { case p: CometPlan => p }
+ assert(cometNodes.nonEmpty, "expected Comet operators when plan-only
mode is disabled")
+ }
+ }
+ }
+
+ private val PLAN_ONLY_PREFIX = "[Comet plan-only]"
+
+ /** Runs `f` and returns the `[Comet plan-only]` reports that
`CometExecRule` logged. */
+ private def capturePlanOnlyReports(f: => Unit): Seq[String] = {
+ val appender = new LogAppender("Comet plan-only reports")
+ withLogAppender(
+ appender,
+ loggerNames = Seq(classOf[CometExecRule].getName),
+ level = Some(Level.WARN)) {
+ f
+ }
+ appender.loggingEvents
+ .map(_.getMessage.getFormattedMessage)
+ .filter(_.startsWith(PLAN_ONLY_PREFIX))
+ .toSeq
+ }
+
+ /** The `Comet accelerated N out of M eligible operators` counts in a
plan-only report. */
+ private def coverageOf(report: String): (Int, Int) = {
+ val pattern = """Comet accelerated (\d+) out of (\d+) eligible
operators""".r
+ pattern
+ .findFirstMatchIn(report)
+ .map(m => (m.group(1).toInt, m.group(2).toInt))
+ .getOrElse(fail(s"report has no coverage summary:\n$report"))
+ }
+
+ // The outer query is planned after any subquery it contains, so a report
slot owned by the
+ // first plan Spark prepares would describe the subquery and never the query
being evaluated.
+ for (aqe <- Seq(true, false)) {
+ test(s"plan-only mode: report describes the outer query, not just a
subquery (AQE=$aqe)") {
+ withSQLConf(
+ SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+ withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+ val reports = capturePlanOnlyReports {
+ spark.sql("SELECT _1 FROM tbl WHERE _1 > (SELECT max(_2) FROM
tbl)").collect()
+ }
+ assert(reports.nonEmpty, "expected a plan-only report")
+ // The outer plan's Filter appears in no subquery plan, so its
presence proves the
+ // outer query was reported and not suppressed by the subquery's
earlier planning.
+ assert(
+ reports.exists(_.contains("Filter")),
+ s"no report describes the outer
query:\n${reports.mkString("\n\n")}")
+ assert(
+ reports.distinct.size == reports.size,
+ s"the same plan was reported more than
once:\n${reports.mkString("\n\n")}")
+ // Expected: one report for the subquery plan, one for the outer
plan. AQE applies the
+ // rule again per stage and per re-optimization; those must not add
reports.
+ assert(
+ reports.size <= 4,
+ s"expected a report per planned plan, got ${reports.size}:\n" +
+ reports.mkString("\n\n"))
+ }
+ }
+ }
+ }
+
+ test("plan-only mode: coverage accounts for post-columnar stage reversion") {
+ withSQLConf(
+ SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+ // AQE off so that Spark applies the post-columnar rules to the whole
plan exactly once,
+ // which is what the preview does, making the two directly comparable.
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "false") {
+ withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
+ val query = "SELECT _2, count(*), sum(_1) FROM tbl GROUP BY _2"
+
+ // Comet accelerates part of this plan when the stage is left alone,
so a preview that
+ // stopped before the post-columnar rules would report a non-zero
count below.
+ withSQLConf(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key ->
"false") {
+ val df = sql(query)
+ df.collect()
+
assert(CometCoverageStats.forPlan(df.queryExecution.executedPlan).cometOperators
> 0)
+ }
+
+ withSQLConf(
+ CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") {
+ // What Comet really executes with reversion enabled.
+ val df = sql(query)
+ df.collect()
+ val executed =
CometCoverageStats.forPlan(df.queryExecution.executedPlan)
+
+ val reports =
withSQLConf(CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
Review Comment:
[P1] Preserve the captured report across Spark 3.x withSQLConf
Could the `withSQLConf` block be moved inside `capturePlanOnlyReports`, or
could these assertions run inside the configuration block? Spark 3.4 and 3.5
define `withSQLConf(...)(f: => Unit): Unit`, unlike the generic Spark 4.x
helper, so `reports` is inferred as `Unit` here. The exact-head [Spark
3.4](https://github.com/apache/datafusion-comet/actions/runs/32280513769/job/96158604209)
and [Spark
3.5](https://github.com/apache/datafusion-comet/actions/runs/32280513769/job/96158604194)
checks both fail test compilation on the subsequent `size`, `mkString`, and
`head` calls. This prevents the supported Spark 3.x test builds from compiling.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -115,12 +115,91 @@ object CometExecRule {
*/
val SKIP_COMET_BROADCAST_TAG:
org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit] =
org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit]("comet.skipCometBroadcast")
+
+ /**
+ * A bounded set of keys, used for plan-only reporting state. Evicts in LRU
order once `limit`
+ * keys are held, so a long-lived driver retains a fixed amount of reporting
state. Same
+ * synchronized-`LinkedHashMap` pattern used by
`IcebergPlanDataInjector.commonCache`.
+ */
+ private class BoundedKeySet(limit: Int) {
+ private val keys: java.util.Map[String, java.lang.Boolean] =
+ java.util.Collections.synchronizedMap(
+ new java.util.LinkedHashMap[String, java.lang.Boolean](16, 0.75f,
true) {
+ override def removeEldestEntry(
+ eldest: java.util.Map.Entry[String, java.lang.Boolean]): Boolean
= size() > limit
+ })
+
+ /** Adds `key`, returning true if it was not already present. */
+ def add(key: String): Boolean = keys.put(key, java.lang.Boolean.TRUE) ==
null
+
+ def contains(key: String): Boolean = keys.containsKey(key)
+ }
+
+ private val PLAN_ONLY_REPORTED_LIMIT = 1024
+
+ /** `executionId:planFingerprint` keys that plan-only mode has already
reported. */
+ private val planOnlyReportedPlans = new
BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT)
+
+ /** Execution IDs whose plan-only report came from the query-stage-prep
rule. */
+ private val planOnlyPrepReportedIds = new
BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT)
+
+ /**
+ * Whether plan-only mode should report `plan`, recording that it did so.
+ *
+ * Spark applies this rule many times during one SQL execution, and only
some of those
+ * applications correspond to a plan the user is asking about:
+ *
+ * - Each scalar subquery and DPP subquery is prepared as its own
top-level plan, and that
+ * happens *before* the outer plan reaches the conversion rules. Keying
the report on the
+ * execution ID alone therefore let a nested subquery consume the slot
and suppressed the
+ * outer plan, which is the plan being evaluated. Keying on the
execution ID *and* the plan
+ * gives the outer plan its own report and each separately prepared
subquery theirs.
+ * - Under AQE the rule also runs once per query stage (as a columnar
rule) and again on every
+ * re-optimization (as a query-stage-prep rule). Those are re-planning
of a plan already
+ * reported, so a `plan` containing query stages is skipped, and once
the query-stage-prep
+ * rule has reported an execution the columnar applications for it stay
quiet.
+ *
+ * @param queryStagePrep
+ * whether the calling rule instance is registered as a query-stage-prep
rule.
+ */
+ private[comet] def shouldReportPlanOnly(
+ executionId: Option[String],
+ plan: SparkPlan,
+ queryStagePrep: Boolean): Boolean = {
+ executionId match {
+ case None =>
Review Comment:
[P2] Deduplicate adaptive reports when the execution ID is absent
Could the no-ID path retain plan-scoped reporting state instead of returning
`true` for every invocation? The public `df.rdd.count()` path can build and
execute AQE stages without installing `spark.sql.execution.id`. A Spark 3.5.2
probe using this exact decision logic and both rule registrations produced five
report decisions for `SELECT id % 2 AS k, count(*) AS n FROM range(20) GROUP BY
id % 2`: initial preparation, the adaptive wrapper, the exchange stage,
adaptive re-optimization, and the final stage. A fresh `collect()` produced
one. Because this branch bypasses both the stage check and deduplication,
plan-only mode rebuilds previews and emits overlapping coverage summaries for
those ordinary RDD-backed workloads, contrary to the documented suppression of
stage/re-optimization reports. Please cover `df.rdd.count()` and planning via
`executedPlan` before an action in the reporting tests.
--
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]
