sunchao commented on code in PR #5394:
URL: https://github.com/apache/datafusion-comet/pull/5394#discussion_r3867696864
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -115,12 +116,143 @@ 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 for which AQE has begun cutting the plan into query
stages. */
+ private val planOnlyStagedExecutions = new
BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT)
+
+ /**
+ * Set on the root of every plan that plan-only mode has reported. Catalyst
copies a node's tags
+ * onto the node that replaces it (`TreeNode.copyTagsFrom`), so the mark
survives the rewrites
+ * Spark applies between one application of this rule and the next, which is
what lets a later
+ * application recognize a plan it has already described.
+ */
+ private val PLAN_ONLY_REPORTED: TreeNodeTag[Unit] =
TreeNodeTag[Unit]("comet.planOnlyReported")
+
+ /**
+ * Whether `plan` is the plan of a query stage AQE has just cut, which
reaches the columnar rule
+ * rooted at the `Exchange` the cut was made at.
+ *
+ * With AQE off a plan rooted at an `Exchange` is an ordinary plan -
`df.repartition(n)`, say -
+ * and must still be reported, hence the `aqeEnabled` guard.
+ */
+ private def isQueryStage(
+ plan: SparkPlan,
+ queryStagePrep: Boolean,
+ aqeEnabled: Boolean): Boolean = {
+ aqeEnabled && !queryStagePrep && plan.isInstanceOf[Exchange]
+ }
+
+ /**
+ * Whether `plan` is an application of this rule to a plan already reported,
rather than to a
+ * plan the user is asking about. The state is on the plan itself, so this
holds whether or not
+ * the application carries a SQL execution ID - `df.rdd.count()` and reading
`executedPlan`
+ * without an action both plan, and in the first case execute AQE stages,
with no execution ID
+ * set.
+ *
+ * `AdaptiveSparkPlanExec` is a leaf as far as `exists` is concerned, so the
wrapper is matched
+ * by name rather than through its contents. A re-optimized plan holds
`QueryStageExec` nodes
+ * for the stages already materialized. A final plan for a query AQE never
had to cut into
+ * stages holds neither, and is recognized by the mark left when it was
first reported.
+ *
+ * @param queryStagePrep
+ * whether the calling rule instance is registered as a query-stage-prep
rule.
+ */
+ private def isReapplication(plan: SparkPlan): Boolean = {
+ plan.exists(p =>
+ p.isInstanceOf[QueryStageExec] || p.isInstanceOf[AdaptiveSparkPlanExec]
||
+ p.getTagValue(PLAN_ONLY_REPORTED).isDefined)
+ }
+
+ /**
+ * Whether plan-only mode should report `plan`, marking it reported if so.
+ *
+ * Spark applies this rule many times while executing one query, and only
some of those
+ * applications correspond to a plan the user is asking about. Under AQE one
query reaches the
+ * rules at least five times:
+ *
+ * - the initial plan, through the query-stage-prep rule - the one to
report;
+ * - the same plan again as a columnar rule, now wrapped in
`AdaptiveSparkPlanExec`;
+ * - each query stage as it is created, again as a columnar rule;
+ * - the re-optimized plan after each stage materializes, through the prep
rule;
+ * - the final plan once every stage has materialized, as a columnar rule.
+ *
+ * Three mechanisms sort those out, because no one of them covers every
shape:
+ *
+ * - 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. A
single report slot per
+ * SQL execution therefore let a nested subquery consume the slot and
suppressed the outer
+ * plan, which is the plan being evaluated. Marking plans individually
gives the outer plan
+ * its own report and each separately prepared subquery theirs; see
[[isReapplication]].
+ * - A mark cannot survive AQE replacing the physical tree wholesale,
which is what happens
+ * when a stage materializes empty and the plan collapses to an empty
relation: the new plan
+ * shares no nodes with the one reported and holds no query stages
either. Once AQE has cut
+ * a stage for an execution, though, everything that arrives afterwards
is AQE re-planning
+ * something already reported, and subqueries are all compiled before
the first stage is
+ * cut, so suppressing the rest of the execution costs no report.
+ * - A subquery referenced from more than one place in the outer plan is
prepared once per
+ * reference, as a separate but identical plan each time. Neither of the
above catches
+ * those, so within one SQL execution the plan's structural hash dedupes
them.
+ *
+ * @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,
+ aqeEnabled: Boolean): Boolean = {
+ if (isQueryStage(plan, queryStagePrep, aqeEnabled)) {
+ // Record that AQE has started executing this query before dropping the
stage itself.
+ executionId.foreach(planOnlyStagedExecutions.add)
+ false
+ } else if (isReapplication(plan)) {
+ false
+ } else if (executionId.exists(planOnlyStagedExecutions.contains)) {
Review Comment:
[P2] Suppress empty AQE reports without an execution ID
The new guard only retains state when an execution ID exists. With AQE and
`spark.comet.explain.planOnly.enabled=true`, an empty grouped query through the
JVM bridge used by [PySpark's
`df.rdd`](https://github.com/apache/spark/blob/v4.1.3/python/pyspark/sql/classic/dataframe.py#L157-L161)
still emits the initial 83% report followed by a misleading 0% `EmptyRelation`
report:
```sql
SELECT id % 2 AS k, count(*) AS n
FROM range(20)
WHERE id < 0
GROUP BY id % 2
```
I reproduced this on the current head with Spark 4.1.3 by invoking
`javaToPython()` and counting the resulting RDD; AQE disabled emits one report.
The action has no SQL execution ID, and the replacement tree has no reporting
tag. Could suppression survive this AQE replacement without requiring an
execution ID, with a regression covering this RDD path?
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -115,12 +116,143 @@ 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 for which AQE has begun cutting the plan into query
stages. */
+ private val planOnlyStagedExecutions = new
BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT)
+
+ /**
+ * Set on the root of every plan that plan-only mode has reported. Catalyst
copies a node's tags
+ * onto the node that replaces it (`TreeNode.copyTagsFrom`), so the mark
survives the rewrites
+ * Spark applies between one application of this rule and the next, which is
what lets a later
+ * application recognize a plan it has already described.
+ */
+ private val PLAN_ONLY_REPORTED: TreeNodeTag[Unit] =
TreeNodeTag[Unit]("comet.planOnlyReported")
+
+ /**
+ * Whether `plan` is the plan of a query stage AQE has just cut, which
reaches the columnar rule
+ * rooted at the `Exchange` the cut was made at.
+ *
+ * With AQE off a plan rooted at an `Exchange` is an ordinary plan -
`df.repartition(n)`, say -
+ * and must still be reported, hence the `aqeEnabled` guard.
+ */
+ private def isQueryStage(
+ plan: SparkPlan,
+ queryStagePrep: Boolean,
+ aqeEnabled: Boolean): Boolean = {
+ aqeEnabled && !queryStagePrep && plan.isInstanceOf[Exchange]
+ }
+
+ /**
+ * Whether `plan` is an application of this rule to a plan already reported,
rather than to a
+ * plan the user is asking about. The state is on the plan itself, so this
holds whether or not
+ * the application carries a SQL execution ID - `df.rdd.count()` and reading
`executedPlan`
+ * without an action both plan, and in the first case execute AQE stages,
with no execution ID
+ * set.
+ *
+ * `AdaptiveSparkPlanExec` is a leaf as far as `exists` is concerned, so the
wrapper is matched
+ * by name rather than through its contents. A re-optimized plan holds
`QueryStageExec` nodes
+ * for the stages already materialized. A final plan for a query AQE never
had to cut into
+ * stages holds neither, and is recognized by the mark left when it was
first reported.
+ *
+ * @param queryStagePrep
+ * whether the calling rule instance is registered as a query-stage-prep
rule.
+ */
+ private def isReapplication(plan: SparkPlan): Boolean = {
+ plan.exists(p =>
+ p.isInstanceOf[QueryStageExec] || p.isInstanceOf[AdaptiveSparkPlanExec]
||
+ p.getTagValue(PLAN_ONLY_REPORTED).isDefined)
+ }
+
+ /**
+ * Whether plan-only mode should report `plan`, marking it reported if so.
+ *
+ * Spark applies this rule many times while executing one query, and only
some of those
+ * applications correspond to a plan the user is asking about. Under AQE one
query reaches the
+ * rules at least five times:
+ *
+ * - the initial plan, through the query-stage-prep rule - the one to
report;
+ * - the same plan again as a columnar rule, now wrapped in
`AdaptiveSparkPlanExec`;
+ * - each query stage as it is created, again as a columnar rule;
+ * - the re-optimized plan after each stage materializes, through the prep
rule;
+ * - the final plan once every stage has materialized, as a columnar rule.
+ *
+ * Three mechanisms sort those out, because no one of them covers every
shape:
+ *
+ * - 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. A
single report slot per
+ * SQL execution therefore let a nested subquery consume the slot and
suppressed the outer
+ * plan, which is the plan being evaluated. Marking plans individually
gives the outer plan
+ * its own report and each separately prepared subquery theirs; see
[[isReapplication]].
+ * - A mark cannot survive AQE replacing the physical tree wholesale,
which is what happens
+ * when a stage materializes empty and the plan collapses to an empty
relation: the new plan
+ * shares no nodes with the one reported and holds no query stages
either. Once AQE has cut
+ * a stage for an execution, though, everything that arrives afterwards
is AQE re-planning
+ * something already reported, and subqueries are all compiled before
the first stage is
+ * cut, so suppressing the rest of the execution costs no report.
+ * - A subquery referenced from more than one place in the outer plan is
prepared once per
+ * reference, as a separate but identical plan each time. Neither of the
above catches
+ * those, so within one SQL execution the plan's structural hash dedupes
them.
+ *
+ * @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,
+ aqeEnabled: Boolean): Boolean = {
+ if (isQueryStage(plan, queryStagePrep, aqeEnabled)) {
+ // Record that AQE has started executing this query before dropping the
stage itself.
+ executionId.foreach(planOnlyStagedExecutions.add)
+ false
+ } else if (isReapplication(plan)) {
+ false
+ } else if (executionId.exists(planOnlyStagedExecutions.contains)) {
+ false
+ } else {
+ plan.setTagValue(PLAN_ONLY_REPORTED, ())
+ // Node tags are not part of a plan's structural hash, so the mark set
above does not
+ // perturb the key. Without an execution ID there is nothing to scope
the state to, and the
+ // checks above have already ruled out the repeat applications AQE
makes, so report.
+ executionId.forall(id =>
planOnlyReportedPlans.add(s"$id:${plan.hashCode()}"))
Review Comment:
[P3] Canonicalize nested subqueries before deduplicating reports
Projecting the same nested scalar subquery twice and calling `.collect()`
with AQE disabled produces four reports under one execution ID: the inner
aggregate, the outer aggregate twice, and the main query. For a Parquet table
`tbl` with integer columns `_1` and `_2`, the reproducer is:
```sql
SELECT _1,
(SELECT max(_2) FROM tbl
WHERE _1 > (SELECT min(_2) FROM tbl)) AS a,
(SELECT max(_2) FROM tbl
WHERE _1 > (SELECT min(_2) FROM tbl)) AS b
FROM tbl
```
On Spark 4.1.3 at the current head, the two outer-aggregate reports are
byte-for-byte identical, and both normal and plan-only physical plans reuse
that subquery through `ReusedSubqueryExec`. The raw plan hash misses this
equivalence, causing redundant previews and warnings. Could the key use a
canonical fingerprint, with a nested-subquery reuse regression? This is a
reporting issue; results and coverage values still match.
--
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]
