more reuse cases

Author: peter-toth

sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/Exchange.scala

@@ -17,15 +17,17 @@
 
 package org.apache.spark.sql.execution.exchange
 
+import java.util.Objects
+
 import scala.collection.mutable
-import scala.collection.mutable.ArrayBuffer
 
 import org.apache.spark.broadcast
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, Expression, SortOrder}
 import org.apache.spark.sql.catalyst.plans.physical.Partitioning
 import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreeNodeRef
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.StructType
@@ -52,6 +54,13 @@ abstract class Exchange extends UnaryExecNode {
 case class ReusedExchangeExec(override val output: Seq[Attribute], child: Exchange)
   extends LeafExecNode {
 
+  override def equals(that: Any): Boolean = that match {
+    case ReusedExchangeExec(output, child) => this.child == output && this.child.eq(child)
+    case _ => false
+  }
+
+  override def hashCode: Int = Objects.hash(output, child)
+
   override def supportsColumnar: Boolean = child.supportsColumnar
 
   // Ignore this wrapper for canonicalizing.
@@ -113,27 +122,38 @@ case class ReuseExchange(conf: SQLConf) extends Rule[SparkPlan] {
     //   have the same schema
     val exchanges = mutable.Map[StructType, (Exchange, mutable.Map[SparkPlan, Exchange])]()
 
-    def reuse(plan: SparkPlan): SparkPlan = plan.transform {
-      case exchange: Exchange =>
-        val (firstSameSchemaExchange, sameResultExchanges) =
-          exchanges.getOrElseUpdate(exchange.schema, (exchange, mutable.Map()))
-        if (firstSameSchemaExchange.ne(exchange)) {
-          if (sameResultExchanges.isEmpty) {
-            sameResultExchanges += firstSameSchemaExchange.canonicalized -> firstSameSchemaExchange
-          }
-          val sameResultExchange =
-            sameResultExchanges.getOrElseUpdate(exchange.canonicalized, exchange)
-          if (sameResultExchange.ne(exchange)) {
-            ReusedExchangeExec(exchange.output, sameResultExchange)
+    def reuse(plan: SparkPlan): SparkPlan = {
+      // Track exchanges that are replaced to reused exchanges to be able to fix ReusedExchangeExec
+      // nodes referencing to them
+      val reuseExchanges = mutable.Map[TreeNodeRef, Exchange]()
+
+      plan.transform {
+        case exchange: Exchange =>
+          val (firstSameSchemaExchange, sameResultExchanges) =
+            exchanges.getOrElseUpdate(exchange.schema, (exchange, mutable.Map()))
+          if (firstSameSchemaExchange.ne(exchange)) {
+            if (sameResultExchanges.isEmpty) {
+              sameResultExchanges +=
+                firstSameSchemaExchange.canonicalized -> firstSameSchemaExchange
+            }
+            val sameResultExchange =
+              sameResultExchanges.getOrElseUpdate(exchange.canonicalized, exchange)
+            if (sameResultExchange.ne(exchange)) {
+              reuseExchanges += new TreeNodeRef(exchange) -> sameResultExchange
+              ReusedExchangeExec(exchange.output, sameResultExchange)
+            } else {
+              exchange
+            }
           } else {
             exchange
           }
-        } else {
-          exchange
+        case reuseExchange @ ReusedExchangeExec(output, child) =>
+          reuseExchanges.get(new TreeNodeRef(child)).map(ReusedExchangeExec(output, _))
+            .getOrElse(reuseExchange)
+        case other => other.transformExpressions {
+          case sub: ExecSubqueryExpression =>
+            sub.withNewPlan(reuse(sub.plan).asInstanceOf[BaseSubqueryExec])
         }
-      case other => other.transformExpressions {
-        case sub: ExecSubqueryExpression =>
-          sub.withNewPlan(reuse(sub.plan).asInstanceOf[BaseSubqueryExec])
       }
     }
 

sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala

@@ -1536,17 +1536,43 @@ class SubquerySuite extends QueryTest with SharedSparkSession {
 
         val plan = df.queryExecution.executedPlan
 
-        val countExchange = plan.collectInPlanAndSubqueries({ case _: Exchange => 1 }).sum
-        val countReusedExchange =
-          plan.collectInPlanAndSubqueries({ case _: ReusedExchangeExec => 1 }).sum
+        val exchangeIds = plan.collectInPlanAndSubqueries { case e: Exchange => e.id }
+        val reusedExchangeIds = plan.collectInPlanAndSubqueries {
+          case re: ReusedExchangeExec => re.child.id
+        }
+
+        if (reuse) {
+          assert(exchangeIds.size == 2, "Exchange reusing not working correctly")
+          assert(reusedExchangeIds.size == 3, "Exchange reusing not working correctly")
+          assert(reusedExchangeIds.forall(exchangeIds.contains(_)),
+            "ReusedExchangeExec should reuse an existing exchange")
+        } else {
+          assert(exchangeIds.size == 5, "expect 5 Exchange when not reusing")
+          assert(reusedExchangeIds.size == 0, "expect 0 ReusedExchangeExec when not reusing")
+        }
+
+        val df2 = sql(
+          """
+            SELECT
+              (SELECT min(a.key) FROM testData AS a JOIN testData AS b ON b.key = a.key),
+              (SELECT max(a.key) FROM testData AS a JOIN testData AS b ON b.key = a.key)
+          """.stripMargin)
+
+        val plan2 = df2.queryExecution.executedPlan
+
+        val exchangeIds2 = plan2.collectInPlanAndSubqueries { case e: Exchange => e.id }
+        val reusedExchangeIds2 = plan2.collectInPlanAndSubqueries {
+          case re: ReusedExchangeExec => re.child.id
+        }
 
         if (reuse) {
-          assert(countExchange == 2, "Exchange reusing not working correctly")
-          assert(countReusedExchange == 3, "Exchange reusing not working correctly")
+          assert(exchangeIds2.size == 3, "Exchange reusing not working correctly")
+          assert(reusedExchangeIds2.size == 3, "Exchange reusing not working correctly")
+          assert(reusedExchangeIds2.forall(exchangeIds2.contains(_)),
+            "ReusedExchangeExec should reuse an existing exchange")
         } else {
-          assert(countExchange == 5, "expect 5 Exchange when not reusing")
-          assert(countReusedExchange == 0,
-            "expect 0 ReusedExchangeExec when not reusing")
+          assert(exchangeIds2.size == 6, "expect 6 Exchange when not reusing")
+          assert(reusedExchangeIds2.size == 0, "expect 0 ReusedExchangeExec when not reusing")
         }
       }
     }