ulysses-you commented on code in PR #37612:
URL: https://github.com/apache/spark/pull/37612#discussion_r960155139


##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/LogicalQueryStage.scala:
##########
@@ -40,7 +40,15 @@ case class LogicalQueryStage(
   override def output: Seq[Attribute] = logicalPlan.output
   override val isStreaming: Boolean = logicalPlan.isStreaming
   override val outputOrdering: Seq[SortOrder] = physicalPlan.outputOrdering
-  override protected val nodePatterns: Seq[TreePattern] = 
Seq(LOGICAL_QUERY_STAGE)
+  override protected val nodePatterns: Seq[TreePattern] = {
+    // Repartition is a special node that it represents a shuffle exchange,
+    // then in AQE the repartition will be always wrapped into 
`LogicalQueryStage`
+    val repartitionPattern = logicalPlan match {
+      case _: RepartitionOperation => Some(REPARTITION_OPERATION)
+      case _ => None
+    }
+    Seq(LOGICAL_QUERY_STAGE) ++ repartitionPattern

Review Comment:
   this is for AQE side to tag `TreePattern`. cc @cloud-fan 



-- 
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]

Reply via email to