[AURON #2351] Reduce native broadcast task serialization size

spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronCheckConvertBroadcastExchangeSuite.scala

@@ -16,13 +16,44 @@
  */
 package org.apache.auron
 
+import org.apache.spark.SparkEnv
 import org.apache.spark.sql.{AuronQueryTest, Row}
 import org.apache.spark.sql.execution.auron.plan.NativeBroadcastExchangeExec
 import org.apache.spark.sql.execution.exchange.BroadcastExchangeExec
 
 class AuronCheckConvertBroadcastExchangeSuite extends AuronQueryTest with BaseAuronSQLSuite {
   import testImplicits._
 
+  test("do not serialize the broadcast relation with Spark tasks") {
+    withSQLConf(
+      "spark.auron.enable.broadcastExchange" -> "true",
+      "spark.auron.enable.bhj" -> "false") {
+      val payload = "x" * 4096
+      (0 until 256)
+        .map(i => (i, s"$i$payload"))
+        .toDF("key", "payload")
+        .createOrReplaceTempView("broad_cast_table1")
+      Seq(0, 255).toDF("key").createOrReplaceTempView("broad_cast_table2")
+
+      val df = spark.sql(
+        "select /*+ broadcast(a)*/ b.key, a.payload from broad_cast_table1 a " +
+          "inner join broad_cast_table2 b on a.key = b.key")
+
+      checkAnswer(df, Seq(Row(0, s"0$payload"), Row(255, s"255$payload")))
+      val exchange = collectFirst(df.queryExecution.executedPlan) {
+        case broadcastExchangeExec: NativeBroadcastExchangeExec => broadcastExchangeExec
+      }.get
+      val broadcast = exchange.executeBroadcast[Any]()
+      try {
+        assert(broadcast eq exchange.executeBroadcast[Any]())
+        val serialized = SparkEnv.get.closureSerializer.newInstance().serialize(broadcast)
+        assert(serialized.remaining() < 64 * 1024)
+      } finally {
+        broadcast.destroy()
+      }
+    }
+  }
+
   test(
     "test bhj broadcastExchange to native where spark.auron.enable.broadcastExchange is true") {
     withSQLConf("spark.auron.enable.broadcastExchange" -> "true") {

spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeBroadcastExchangeBase.scala

@@ -27,7 +27,6 @@ import scala.collection.immutable.SortedMap
 import scala.concurrent.Promise
 import scala.jdk.CollectionConverters._
 
-import org.apache.commons.lang3.reflect.MethodUtils
 import org.apache.spark.OneToOneDependency
 import org.apache.spark.Partition
 import org.apache.spark.SparkException
@@ -125,7 +124,16 @@ abstract class NativeBroadcastExchangeBase(mode: BroadcastMode, override val chi
     relationFuture
   }
 
-  override def doExecuteBroadcast[T](): Broadcast[T] = {
+  // Mixed native/Spark execution needs a second JVM broadcast; otherwise the transformed relation
+  // would be serialized with every Spark task.
+  // Build it under a separate lock because relationFuture initializes lazy fields on another thread.
+  @transient
+  private lazy val sparkBroadcastLock = new Object
+
+  @transient
+  private var sparkBroadcast: Broadcast[Any] = _
+
+  private def createSparkBroadcast(): Broadcast[Any] = {
     val singlePartition = new Partition() {
       override def index: Int = 0
     }
@@ -145,18 +153,14 @@ abstract class NativeBroadcastExchangeBase(mode: BroadcastMode, override val chi
       .map(_.copy())
       .toArray
 
-    val broadcast = relationFuture.get // broadcast must be resolved
-    val v = mode.transform(dataRows)
-    val dummyBroadcasted = new Broadcast[Any](-1) {
-      override protected def getValue(): Any = v
-      override protected def doUnpersist(blocking: Boolean): Unit = {
-        MethodUtils.invokeMethod(broadcast, true, "doUnpersist", Array(blocking))
-      }
-      override protected def doDestroy(blocking: Boolean): Unit = {
-        MethodUtils.invokeMethod(broadcast, true, "doDestroy", Array(blocking))
-      }
+    sparkContext.broadcast(mode.transform(dataRows))
+  }
+
+  override def doExecuteBroadcast[T](): Broadcast[T] = sparkBroadcastLock.synchronized {
+    if (sparkBroadcast == null) {
+      sparkBroadcast = createSparkBroadcast()
     }
-    dummyBroadcasted.asInstanceOf[Broadcast[T]]
+    sparkBroadcast.asInstanceOf[Broadcast[T]]
   }
 
   def doExecuteBroadcastNative[T](): broadcast.Broadcast[T] = {