fix: potential native broadcast failure in scenarios with ReusedExhange

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

@@ -277,7 +277,7 @@ object CometBroadcastExchangeExec {
 class CometBatchRDD(
     sc: SparkContext,
     numPartitions: Int,
-    value: broadcast.Broadcast[Array[ChunkedByteBuffer]])
+    private[comet] val value: broadcast.Broadcast[Array[ChunkedByteBuffer]])
     extends RDD[ColumnarBatch](sc, Nil) {
 
   override def getPartitions: Array[Partition] = (0 until numPartitions).toArray.map { i =>

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

@@ -233,13 +233,19 @@ abstract class CometNativeExec extends CometExec {
 
         foreachUntilCometInput(this)(sparkPlans += _)
 
-        // Find the first non broadcast plan
-        val firstNonBroadcastPlan = sparkPlans.zipWithIndex.find {
+        val nonBroadcastPredicate: (SparkPlan, Int) => Boolean = {
           case (_: CometBroadcastExchangeExec, _) => false
           case (BroadcastQueryStageExec(_, _: CometBroadcastExchangeExec, _), _) => false
           case (BroadcastQueryStageExec(_, _: ReusedExchangeExec, _), _) => false
           case _ => true
         }
+        // Find the first non broadcast plan which is not a reused exchange if possible
+        val firstNonBroadcastPlan = sparkPlans.zipWithIndex
+          .find {
+            case (ReusedExchangeExec(_, _: CometBroadcastExchangeExec), _) => false
+            case (p, idx) => nonBroadcastPredicate(p, idx)
+          }
+          .orElse(sparkPlans.zipWithIndex.find(_._1.isInstanceOf[ReusedExchangeExec]))
 
         val containsBroadcastInput = sparkPlans.exists {
           case _: CometBroadcastExchangeExec => true
@@ -303,7 +309,16 @@ abstract class CometNativeExec extends CometExec {
         }
 
         if (inputs.nonEmpty) {
-          ZippedPartitionsRDD(sparkContext, inputs.toSeq)(createCometExecIter)
+          val fixedInputs = inputs.map {
+            case cometBatchRDD: CometBatchRDD
+                if cometBatchRDD.getNumPartitions != firstNonBroadcastPlanNumPartitions =>
+              new CometBatchRDD(
+                sparkContext,
+                firstNonBroadcastPlanNumPartitions,
+                cometBatchRDD.value)
+            case other => other
+          }
+          ZippedPartitionsRDD(sparkContext, fixedInputs.toSeq)(createCometExecIter)
         } else {
           val partitionNum = firstNonBroadcastPlanNumPartitions
           CometExecRDD(sparkContext, partitionNum)(createCometExecIter)

spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala

@@ -2107,6 +2107,48 @@ class CometExecSuite extends CometTestBase {
     }
   }
 
+  test("ReusedExchange broadcast with incompatible partitions number does not fail") {
+    withTable("tbl1", "tbl2", "tbl3") {
+      // enforce different number of partitions for future broadcasts/exchanges
+      spark
+        .range(50)
+        .withColumnRenamed("id", "x")
+        .repartition(2)
+        .writeTo("tbl1")
+        .using("parquet")
+        .create()
+      spark
+        .range(50)
+        .withColumnRenamed("id", "y")
+        .repartition(3)
+        .writeTo("tbl2")
+        .using("parquet")
+        .create()
+      spark
+        .range(50)
+        .withColumnRenamed("id", "z")
+        .repartition(4)
+        .writeTo("tbl3")
+        .using("parquet")
+        .create()
+      val df1 = spark.table("tbl1")
+      val df2 = spark.table("tbl2")
+      val df3 = spark.table("tbl3")
+      Seq("true", "false").foreach(aqeEnabled =>
+        withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled) {
+          val dfWithReusedExchange = df1
+            .join(df3.hint("broadcast").join(df1, $"x" === $"z"), "x", "right")
+            .join(
+              df3
+                .hint("broadcast")
+                .join(df2, $"y" === $"z", "right")
+                .withColumnRenamed("z", "z1"),
+              $"x" === $"y")
+          checkSparkAnswerAndOperator(dfWithReusedExchange, classOf[ReusedExchangeExec])
+        })
+    }
+  }
+
   test("SparkToColumnar override node name for columnar input") {
     withSQLConf(
       SQLConf.USE_V1_SOURCE_LIST.key -> "",