[SPARK-41214][SQL] - SQL Metrics are missing from Spark UI when AQE for Cached DataFrame is enabled

[erenavsarogullari]

What changes were proposed in this pull request?

JIRA: https://issues.apache.org/jira/browse/SPARK-41214
spark.sql.optimizer.canChangeCachedPlanOutputPartitioning enables AQE optimizations under InMemoryRelation(IMR) nodes by creating sub-SparkPlans for AQE optimizations. However, when spark.sql.optimizer.canChangeCachedPlanOutputPartitioning = true, Spark UI does not show correct final DAG due to lack of final sub-plans (under IMR) submissions (into UI).

Why are the changes needed?

spark.sql.optimizer.canChangeCachedPlanOutputPartitioning enables AQE optimizations under InMemoryRelation(IMR) nodes. Sample query in Jira has IMR node on both BroadcastHashJoin legs. However,
when spark.sql.optimizer.canChangeCachedPlanOutputPartitioning = true, following datas are missed due to lack of final sub-plans (under IMR) submissions (into UI).

• Physical operators' SQLMetrics (before AdaptiveSparkPlan) are missed such as Exchange and HashAggregate on left BHJ leg and HashAggregate on right BHJ leg,
• WSCG blocks are missed on left BHJ leg,
• AQEShuffleRead node is missed on left BHJ leg.

DAG before fix:

DAG after fix:

Does this PR introduce any user-facing change?

Currently, Spark UI does not show final DAG when spark.sql.optimizer.canChangeCachedPlanOutputPartitioning = true so this causes some physical operators metrics as missed. This PR aims to fix this problem.

How was this patch tested?

Specific UT Coverage will be added (In Progress).

[erenavsarogullari]

cc @cloud-fan @Ngone51

[cloud-fan]

is this still a problem after #38558 ?

[ulysses-you]

The reason is, the AdaptiveSparkPlanExec inside InMemoryTableScanExec does not register a new execution id in SQLExecution. Then when it materializes, AdaptiveSparkPlanExec.getExecutionId can not find the corresponding QueryExecution.

spark/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala

Lines 221 to 226 in 0ae82d9

	private def getExecutionId: Option[Long] = {
	// If the `QueryExecution` does not match the current execution ID, it means the execution ID
	// belongs to another (parent) query, and we should not call update UI in this query.
	Option(context.session.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY))
	.map(_.toLong).filter(SQLExecution.getQueryExecution(_) eq context.qe)
	}

A simple example like:

          ...
           |
  AdaptiveSparkPlanExec (queryexecution 0, no execution id)
           |
  InMemoryTableScanExec
           |
          ...
           |
  AdaptiveSparkPlanExec (queryexecution 1, execution id 0)

[ulysses-you]

i s this still a problem after #38558 ?

@cloud-fan oh just see it. it not related

[cloud-fan]

after #38558 , I think only the first cache access will miss the metrics?

[ulysses-you]

There is another issue.. it seems the existed framework does not support to get metrics across multi sql execution. This issue is not about AQE, that means even we disable AQE for cache plan, we still can not get valid metrics after first access.

[github-actions]

We're closing this PR because it hasn't been updated in a while. This isn't a judgement on the merit of the PR in any way. It's just a way of keeping the PR queue manageable.
If you'd like to revive this PR, please reopen it and ask a committer to remove the Stale tag!