fix: Check reused broadcast plan in non-AQE and make setNumPartitions thread safe (#2398)

* fix: Check broadcast plan of ReusedExchangeExec

* thread-safety

Author: wForget

spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastExchangeExec.scala

@@ -276,7 +276,7 @@ object CometBroadcastExchangeExec {
  */
 class CometBatchRDD(
     sc: SparkContext,
-    numPartitions: Int,
+    @volatile var numPartitions: Int,
     value: broadcast.Broadcast[Array[ChunkedByteBuffer]])
     extends RDD[ColumnarBatch](sc, Nil) {
 
@@ -289,6 +289,12 @@ class CometBatchRDD(
     partition.value.value.toIterator
       .flatMap(Utils.decodeBatches(_, this.getClass.getSimpleName))
   }
+
+  def withNumPartitions(numPartitions: Int): CometBatchRDD = {
+    this.numPartitions = numPartitions
+    this
+  }
+
 }
 
 class CometBatchPartition(

spark/src/main/scala/org/apache/spark/sql/comet/operators.scala

@@ -238,13 +238,15 @@ abstract class CometNativeExec extends CometExec {
           case (_: CometBroadcastExchangeExec, _) => false
           case (BroadcastQueryStageExec(_, _: CometBroadcastExchangeExec, _), _) => false
           case (BroadcastQueryStageExec(_, _: ReusedExchangeExec, _), _) => false
+          case (ReusedExchangeExec(_, _: CometBroadcastExchangeExec), _) => false
           case _ => true
         }
 
         val containsBroadcastInput = sparkPlans.exists {
           case _: CometBroadcastExchangeExec => true
           case BroadcastQueryStageExec(_, _: CometBroadcastExchangeExec, _) => true
           case BroadcastQueryStageExec(_, _: ReusedExchangeExec, _) => true
+          case ReusedExchangeExec(_, _: CometBroadcastExchangeExec) => true
           case _ => false
         }
 
@@ -272,16 +274,28 @@ abstract class CometNativeExec extends CometExec {
         sparkPlans.zipWithIndex.foreach { case (plan, idx) =>
           plan match {
             case c: CometBroadcastExchangeExec =>
-              inputs += c.setNumPartitions(firstNonBroadcastPlanNumPartitions).executeColumnar()
+              inputs += c
+                .executeColumnar()
+                .asInstanceOf[CometBatchRDD]
+                .withNumPartitions(firstNonBroadcastPlanNumPartitions)
             case BroadcastQueryStageExec(_, c: CometBroadcastExchangeExec, _) =>
-              inputs += c.setNumPartitions(firstNonBroadcastPlanNumPartitions).executeColumnar()
+              inputs += c
+                .executeColumnar()
+                .asInstanceOf[CometBatchRDD]
+                .withNumPartitions(firstNonBroadcastPlanNumPartitions)
             case ReusedExchangeExec(_, c: CometBroadcastExchangeExec) =>
-              inputs += c.setNumPartitions(firstNonBroadcastPlanNumPartitions).executeColumnar()
+              inputs += c
+                .executeColumnar()
+                .asInstanceOf[CometBatchRDD]
+                .withNumPartitions(firstNonBroadcastPlanNumPartitions)
             case BroadcastQueryStageExec(
                   _,
                   ReusedExchangeExec(_, c: CometBroadcastExchangeExec),
                   _) =>
-              inputs += c.setNumPartitions(firstNonBroadcastPlanNumPartitions).executeColumnar()
+              inputs += c
+                .executeColumnar()
+                .asInstanceOf[CometBatchRDD]
+                .withNumPartitions(firstNonBroadcastPlanNumPartitions)
             case _: CometNativeExec =>
             // no-op
             case _ if idx == firstNonBroadcastPlan.get._2 =>