[VL]Optimize driver side build hash table performance

backends-velox/src/main/java/org/apache/gluten/vectorized/HashJoinBuilder.java

@@ -54,6 +54,15 @@ public static native long deserializeHashTableDirect(
 
   public static native void serializeHashTableDirect(long hashTableHandle, long address, long size);
 
+  /**
+   * Builds a Velox hash table from the given batches.
+   *
+   * @param serializeOnly when true, the table is built only to be handed to {@link
+   *     #serializeHashTableDirect} and is never probed in this process. The slot array is then
+   *     neither allocated nor populated, since it is not part of the serialized form and the
+   *     consuming side rebuilds it in {@code deserializeHashTableDirect}. Must be false for the
+   *     executor-side build, whose table is probed in place.
+   */
   public native long nativeBuild(
       String buildHashTableId,
       long[] batchHandlers,
@@ -66,5 +75,6 @@ public native long nativeBuild(
       byte[] namedStruct,
       boolean isNullAwareAntiJoin,
       long bloomFilterPushdownSize,
-      int broadcastHashTableBuildThreads);
+      int broadcastHashTableBuildThreads,
+      boolean serializeOnly);
 }

backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala

@@ -899,7 +899,11 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with Logging {
       math
         .ceil(dataSize.value.toDouble / VeloxConfig.get.veloxBroadcastHashTableBuildTargetBytes)
         .toInt
-    val buildThreadsValue = if (rawThreads < 1) 1 else rawThreads
+    // Each build thread constructs a full partial hash table that prepareJoinTable() later has to
+    // merge, and they all run on the shared Velox IO executor. Cap the fan-out so that a large
+    // build side cannot spawn an unbounded number of them.
+    val buildThreadsValue =
+      math.max(1, math.min(rawThreads, java.lang.Runtime.getRuntime.availableProcessors))
     buildThreads += buildThreadsValue
 
     // Create the base ColumnarBuildSideRelation first
@@ -927,7 +931,12 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with Logging {
     // Only do this for HashedRelationBroadcastMode and when offload is enabled
     val shouldBuildOnDriver = VeloxConfig.get.enableDriverSideBroadcastHashTableBuild &&
       mode.isInstanceOf[HashedRelationBroadcastMode] &&
-      offload
+      offload &&
+      // With cuDF the join builds its own GPU hash table from the build side value stream, so
+      // VeloxBroadcastBuildSideRDD asks the relation for its batches on the executor. A
+      // driver-built table cannot serve that: the raw build side is not part of the broadcast
+      // payload and the executor has no way to get at it.
+      !GlutenConfig.get.enableColumnarCudf
 
     if (shouldBuildOnDriver) {
       // Try to get broadcast join context from logical plan tag

backends-velox/src/main/scala/org/apache/gluten/execution/SerializedBroadcastHashTable.scala

@@ -27,6 +27,17 @@ import java.io.{Externalizable, ObjectInput, ObjectOutput}
 /**
  * Serialized broadcast hash table that can be efficiently broadcast to executors. This is built on
  * the driver and contains the serialized hash table data.
+ *
+ * @param broadcastId
+ *   Identity of the driver-side broadcast this table was built for. Several broadcast hash joins
+ *   can share one broadcast exchange (reused exchange), in which case they all see the same
+ *   `broadcastId` while having distinct per-join hash table ids. Executors key the (expensive)
+ *   deserialization on this id so that the table is materialized once and then shared.
+ * @param buildSideRelation
+ *   The raw build side relation. Driver-only: it backs [[BuildSideRelation.transform]] for DPP and
+ *   [[BuildSideRelation.deserialized]] for broadcast-mode conversion, both of which run on the
+ *   driver. It is deliberately excluded from the wire format, since shipping it would broadcast the
+ *   whole raw build side alongside the hash table.
  */
 class SerializedBroadcastHashTable(
     var serializedData: UnsafeByteArray,
@@ -36,10 +47,11 @@ class SerializedBroadcastHashTable(
     var droppedDuplicates: Boolean,
     var bloomFilterBlocksByteSize: Long,
     var hashProbeDynamicFiltersProduced: Long,
-    var buildSideRelation: BuildSideRelation)
+    var broadcastId: String,
+    @transient var buildSideRelation: BuildSideRelation)
   extends Externalizable {
 
-  def this() = this(null, 0, false, false, false, 0, 0, null) // Required for Externalizable
+  def this() = this(null, 0, false, false, false, 0, 0, null, null) // Required for Externalizable
 
   override def writeExternal(out: ObjectOutput): Unit = {
     out.writeLong(numRows)
@@ -48,8 +60,9 @@ class SerializedBroadcastHashTable(
     out.writeBoolean(droppedDuplicates)
     out.writeLong(bloomFilterBlocksByteSize)
     out.writeLong(hashProbeDynamicFiltersProduced)
+    out.writeUTF(if (broadcastId == null) "" else broadcastId)
     serializedData.writeExternal(out)
-    out.writeObject(buildSideRelation)
+    // 'buildSideRelation' is intentionally not written. See the class doc.
   }
 
   override def readExternal(in: ObjectInput): Unit = {
@@ -59,10 +72,14 @@ class SerializedBroadcastHashTable(
     droppedDuplicates = in.readBoolean()
     bloomFilterBlocksByteSize = in.readLong()
     hashProbeDynamicFiltersProduced = in.readLong()
+    broadcastId = in.readUTF() match {
+      case "" => null
+      case id => id
+    }
     val data = new UnsafeByteArray()
     data.readExternal(in)
     serializedData = data
-    buildSideRelation = in.readObject().asInstanceOf[BuildSideRelation]
+    buildSideRelation = null
   }
 
   /**
@@ -74,6 +91,12 @@ class SerializedBroadcastHashTable(
    *   Hash table builder handle
    */
   def deserialize(cacheKey: String): Long = {
+    if (serializedData == null || serializedData.isReleased) {
+      throw new IllegalStateException(
+        s"Serialized hash table bytes for broadcast $broadcastId have already been released. " +
+          "They are freed once the native table has been materialized from them, so this " +
+          "instance cannot be deserialized again.")
+    }
     HashJoinBuilder.deserializeHashTableDirect(
       cacheKey,
       serializedData.address(),
@@ -82,8 +105,20 @@ class SerializedBroadcastHashTable(
       joinHasNullKeys)
   }
 
+  /**
+   * Frees the off-heap buffer holding the serialized bytes. Only safe once the native hash table
+   * has been materialized from it, and only on an executor: on the driver the very same object is
+   * still owned by the broadcast variable and by
+   * [[VeloxBroadcastBuildSideCache.buildAndSerializeOnDriverInBroadcastExchange]]'s cache.
+   */
+  def releaseSerializedData(): Unit = {
+    if (serializedData != null) {
+      serializedData.release()
+    }
+  }
+
   /** Get the size of serialized data in bytes. */
-  def sizeInBytes: Long = serializedData.size()
+  def sizeInBytes: Long = if (serializedData == null) 0L else serializedData.size()
 }
 
 object SerializedBroadcastHashTable {
@@ -95,6 +130,7 @@ object SerializedBroadcastHashTable {
       droppedDuplicates: Boolean,
       bloomFilterBlocksByteSize: Long,
       hashProbeDynamicFiltersProduced: Long,
+      broadcastId: String,
       buildSideRelation: BuildSideRelation): SerializedBroadcastHashTable =
     new SerializedBroadcastHashTable(
       serializedData,
@@ -104,7 +140,9 @@ object SerializedBroadcastHashTable {
       droppedDuplicates,
       bloomFilterBlocksByteSize,
       hashProbeDynamicFiltersProduced,
-      buildSideRelation)
+      broadcastId,
+      buildSideRelation
+    )
 
   /**
    * Build and serialize a hash table on the driver.
@@ -147,7 +185,9 @@ object SerializedBroadcastHashTable {
         droppedDuplicates,
         bloomFilterBlocksByteSize,
         hashProbeDynamicFiltersProduced,
-        buildSideRelation)
+        cacheKey,
+        buildSideRelation
+      )
     } finally {
       synchronized {
         HashJoinBuilder.clearHashTable(cacheKey, hashTableHandle)

backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastBuildSideCache.scala

@@ -38,6 +38,12 @@ case class BroadcastHashTable(
     relation: BuildSideRelation,
     droppedDuplicates: Boolean)
 
+/**
+ * One natively deserialized hash table, shared by every broadcast hash join that reads the same
+ * driver-side broadcast.
+ */
+private case class SharedDeserializedHashTable(pointer: Long)
+
 /**
  * `VeloxBroadcastBuildSideCache` is used for controlling to build bhj hash table once.
  *
@@ -66,6 +72,26 @@ object VeloxBroadcastBuildSideCache
       .removalListener(this)
       .build[String, BroadcastHashTable]()
 
+  // Executor-side cache of natively deserialized hash tables, keyed by the driver-side broadcast
+  // id rather than by the per-join hash table id, so that joins sharing a reused broadcast
+  // exchange share one native table. Values are owned here; per-join entries in
+  // 'buildSideRelationCache' hold clones.
+  private val sharedDeserializedCache: Cache[String, SharedDeserializedHashTable] =
+    Caffeine.newBuilder
+      .expireAfterAccess(expiredTime, TimeUnit.SECONDS)
+      .removalListener(
+        new RemovalListener[String, SharedDeserializedHashTable] {
+          override def onRemoval(
+              key: String,
+              value: SharedDeserializedHashTable,
+              cause: RemovalCause): Unit = {
+            if (value != null) {
+              HashJoinBuilder.clearHashTable(key, value.pointer)
+            }
+          }
+        }
+      ).build[String, SharedDeserializedHashTable]()
+
   // Cache for driver-side serialized hash tables to avoid rebuilding for reuse exchange
   private val driverSerializedCache: Cache[String, SerializedBroadcastHashTable] =
     Caffeine.newBuilder
@@ -189,7 +215,34 @@ object VeloxBroadcastBuildSideCache
     }
   }
 
-  /** Deserialize hash table on executor from broadcast data. */
+  /**
+   * Returns the raw build side relation that the driver-side build of `broadcastId` was made from,
+   * if this JVM is the driver that built it.
+   *
+   * The relation is deliberately kept out of the broadcast payload, so a
+   * [[SerializedBroadcastHashTable]] read back from the payload has none. That includes the copy
+   * the driver itself gets: the broadcast is created with `serializedOnly = true`, so no
+   * deserialized copy is retained on the driver and even a driver-side `broadcast.value` goes
+   * through the wire format. Consumers that legitimately need the raw build side all run on the
+   * driver (DPP key extraction, broadcast mode conversion), where this lookup finds the original
+   * that [[buildAndSerializeOnDriverInBroadcastExchange]] cached. On an executor it finds nothing,
+   * which is the correct answer there.
+   */
+  def driverBuildSideRelation(broadcastId: String): Option[BuildSideRelation] =
+    Option(broadcastId)
+      .flatMap(id => Option(driverSerializedCache.getIfPresent(id)))
+      .flatMap(serialized => Option(serialized.buildSideRelation))
+
+  /**
+   * Deserialize hash table on executor from broadcast data.
+   *
+   * A reused broadcast exchange feeds several broadcast hash joins, each with its own
+   * `broadcastHashTableId` (the native side looks the table up by that id via [[get]]). Doing the
+   * deserialization per join id would materialize a full copy of the table per join. Instead the
+   * table is deserialized once per driver-side broadcast and every join id gets a cheap clone that
+   * shares the same native table, mirroring what the executor-side build path does via
+   * `cloneHashTable`.
+   */
   def deserializeOnExecutor(
       serialized: SerializedBroadcastHashTable,
       broadcastHashTableId: String,
@@ -199,11 +252,14 @@ object VeloxBroadcastBuildSideCache
     buildSideRelationCache.get(
       broadcastHashTableId,
       (_: String) => {
-        logInfo(s"Deserializing hash table on executor for broadcast ID: $broadcastHashTableId")
-        val startTime = System.currentTimeMillis()
-        val hashTableHandle = serialized.deserialize(broadcastHashTableId)
-        val timeMs = System.currentTimeMillis() - startTime
-        deserializeHashTableTimeMetric.foreach(_ += timeMs)
+        val shared = getOrDeserializeShared(
+          serialized,
+          broadcastHashTableId,
+          deserializeHashTableTimeMetric)
+        // Register the shared table under this join's id as well, so that the native probe side
+        // can resolve it. The clone holds its own reference to the same native table.
+        val hashTableHandle =
+          HashJoinBuilder.cloneHashTable(broadcastHashTableId, shared.pointer)
         BroadcastHashTable(
           hashTableHandle,
           serialized.buildSideRelation,
@@ -212,6 +268,50 @@ object VeloxBroadcastBuildSideCache
     )
   }
 
+  /**
+   * Returns a handle to the native hash table for `serialized`, deserializing it at most once per
+   * driver-side broadcast id. The returned handle is owned by [[sharedDeserializedCache]]; callers
+   * must clone it rather than releasing it.
+   */
+  private def getOrDeserializeShared(
+      serialized: SerializedBroadcastHashTable,
+      broadcastHashTableId: String,
+      deserializeHashTableTimeMetric: Option[org.apache.spark.sql.execution.metric.SQLMetric])
+      : SharedDeserializedHashTable = {
+    // Older payloads, and any path that did not go through the driver-side build, carry no
+    // broadcast id. Fall back to keying on the join id, which is what the previous behavior was.
+    val sharedKey =
+      if (serialized.broadcastId != null) serialized.broadcastId else broadcastHashTableId
+
+    sharedDeserializedCache.get(
+      sharedKey,
+      (key: String) => {
+        logInfo(s"Deserializing hash table on executor for broadcast ID: $key")
+        val startTime = System.currentTimeMillis()
+        val hashTableHandle = serialized.deserialize(key)
+        val timeMs = System.currentTimeMillis() - startTime
+        deserializeHashTableTimeMetric.foreach(_ += timeMs)
+
+        // The serialized bytes have been fully consumed into the native table. On an executor
+        // nothing else refers to them, so free the off-heap copy instead of waiting for the
+        // broadcast object to be collected. On the driver the same object is still owned by the
+        // broadcast variable and by driverSerializedCache, so it must be left alone.
+        if (!isDriver) {
+          serialized.releaseSerializedData()
+        }
+        SharedDeserializedHashTable(hashTableHandle)
+      }
+    )
+  }
+
+  // Mirrors the private[spark] SparkContext.DRIVER_IDENTIFIER.
+  private val driverExecutorId = "driver"
+
+  private def isDriver: Boolean = {
+    val env = SparkEnv.get
+    env == null || env.executorId == driverExecutorId
+  }
+
   /** This is called from c++ side. */
   def get(broadcastHashtableId: String): Long = {
     Option(buildSideRelationCache.getIfPresent(broadcastHashtableId))
@@ -232,6 +332,7 @@ object VeloxBroadcastBuildSideCache
 
   def cleanAll(): Unit = {
     buildSideRelationCache.invalidateAll()
+    sharedDeserializedCache.invalidateAll()
     driverSerializedCache.invalidateAll()
   }
 

backends-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarBuildSideRelation.scala

@@ -243,7 +243,9 @@ case class ColumnarBuildSideRelation(
               SubstraitUtil.toNameStruct(newOutput).toByteArray,
               broadcastContext.isNullAwareAntiJoin,
               broadcastContext.bloomFilterPushdownSize,
-              buildThreads
+              buildThreads,
+              // This table is probed in place on this executor.
+              false
             )
         } finally {
           jniWrapper.close(serializeHandle)
@@ -336,7 +338,10 @@ case class ColumnarBuildSideRelation(
               SubstraitUtil.toNameStruct(newOutput).toByteArray,
               broadcastContext.isNullAwareAntiJoin,
               broadcastContext.bloomFilterPushdownSize,
-              buildThreads
+              buildThreads,
+              // Driver-side build: the table is only serialized and broadcast, never probed here,
+              // so skip building the slot array that the executors rebuild anyway.
+              true
             )
         } finally {
           jniWrapper.close(serializeHandle)