[VL]Use Velox's HashTableCache to cache the BHJ's HashTable (#12163)

JkSelf · web-flow · commit 0f8074c8fd0b · 2026-07-09T10:02:39.000+01:00
diff --git a/backends-velox/src/main/java/org/apache/gluten/vectorized/HashJoinBuilder.java b/backends-velox/src/main/java/org/apache/gluten/vectorized/HashJoinBuilder.java
@@ -35,12 +35,12 @@ public long rtHandle() {
     return runtime.getHandle();
   }
 
-  public static native void clearHashTable(long hashTableData);
+  public static native void clearHashTable(String cacheKey, long hashTableData);
 
-  public static native long cloneHashTable(long hashTableData);
+  public static native long cloneHashTable(String cacheKey, long hashTableData);
 
   public static native long deserializeHashTableDirect(
-      long address, int size, boolean ignoreNullKeys, boolean joinHasNullKeys);
+      String cacheKey, long address, int size, boolean ignoreNullKeys, boolean joinHasNullKeys);
 
   public static native boolean getHashTableIgnoreNullKeys(long hashTableHandle);
 
diff --git a/backends-velox/src/main/scala/org/apache/gluten/execution/HashJoinExecTransformer.scala b/backends-velox/src/main/scala/org/apache/gluten/execution/HashJoinExecTransformer.scala
@@ -106,7 +106,7 @@ case class BroadcastHashJoinExecTransformer(
     right,
     isNullAwareAntiJoin) {
 
-  // Unique ID for the build side.
+  // Unique ID for the build side
   lazy val buildBroadcastTableId: String = buildPlan.id.toString
 
   override protected lazy val substraitJoinType: JoinRel.JoinType = joinType match {
diff --git a/backends-velox/src/main/scala/org/apache/gluten/execution/SerializedBroadcastHashTable.scala b/backends-velox/src/main/scala/org/apache/gluten/execution/SerializedBroadcastHashTable.scala
@@ -73,8 +73,9 @@ class SerializedBroadcastHashTable(
    * @return
    *   Hash table builder handle
    */
-  def deserialize(): Long = {
+  def deserialize(cacheKey: String): Long = {
     HashJoinBuilder.deserializeHashTableDirect(
+      cacheKey,
       serializedData.address(),
       Math.toIntExact(serializedData.size()),
       ignoreNullKeys,
@@ -117,6 +118,7 @@ object SerializedBroadcastHashTable {
    */
   def fromHashTable(
       hashTableHandle: Long,
+      cacheKey: String,
       buildSideRelation: BuildSideRelation,
       droppedDuplicates: Boolean,
       numRows: Long): SerializedBroadcastHashTable = {
@@ -148,7 +150,7 @@ object SerializedBroadcastHashTable {
         buildSideRelation)
     } finally {
       synchronized {
-        HashJoinBuilder.clearHashTable(hashTableHandle)
+        HashJoinBuilder.clearHashTable(cacheKey, hashTableHandle)
       }
     }
   }
diff --git a/backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastBuildSideCache.scala b/backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastBuildSideCache.scala
@@ -161,6 +161,7 @@ object VeloxBroadcastBuildSideCache
           val result =
             SerializedBroadcastHashTable.fromHashTable(
               hashTableHandle,
+              broadcastId,
               relation,
               droppedDuplicates,
               numRows)
@@ -197,7 +198,7 @@ object VeloxBroadcastBuildSideCache
       (_: String) => {
         logInfo(s"Deserializing hash table on executor for broadcast ID: $broadcastHashTableId")
         val startTime = System.currentTimeMillis()
-        val hashTableHandle = serialized.deserialize()
+        val hashTableHandle = serialized.deserialize(broadcastHashTableId)
         val timeMs = System.currentTimeMillis() - startTime
         deserializeHashTableTimeMetric.foreach(_ += timeMs)
         BroadcastHashTable(
@@ -245,7 +246,7 @@ object VeloxBroadcastBuildSideCache
         }
       }
 
-      HashJoinBuilder.clearHashTable(value.pointer)
+      HashJoinBuilder.clearHashTable(key, value.pointer)
     }
   }
 }
diff --git a/backends-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarBuildSideRelation.scala b/backends-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarBuildSideRelation.scala
@@ -242,7 +242,9 @@ case class ColumnarBuildSideRelation(
         (hashTableData(droppedDuplicates), this, droppedDuplicates)
       } else {
         (
-          HashJoinBuilder.cloneHashTable(hashTableData(droppedDuplicates)),
+          HashJoinBuilder.cloneHashTable(
+            broadcastContext.buildHashTableId,
+            hashTableData(droppedDuplicates)),
           null,
           droppedDuplicates)
       }
@@ -332,7 +334,12 @@ case class ColumnarBuildSideRelation(
 
         (hashTableData(droppedDuplicates), this, droppedDuplicates)
       } else {
-        (HashJoinBuilder.cloneHashTable(hashTableData(droppedDuplicates)), null, droppedDuplicates)
+        (
+          HashJoinBuilder.cloneHashTable(
+            broadcastContext.buildHashTableId,
+            hashTableData(droppedDuplicates)),
+          null,
+          droppedDuplicates)
       }
     }
 
diff --git a/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala b/backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala
@@ -214,7 +214,9 @@ class UnsafeColumnarBuildSideRelation(
         (hashTableData(droppedDuplicates), this, droppedDuplicates)
       } else {
         (
-          HashJoinBuilder.cloneHashTable(hashTableData(droppedDuplicates)),
+          HashJoinBuilder.cloneHashTable(
+            broadcastContext.buildHashTableId,
+            hashTableData(droppedDuplicates)),
           null,
           droppedDuplicates)
       }
@@ -305,7 +307,12 @@ class UnsafeColumnarBuildSideRelation(
 
         (hashTableData(droppedDuplicates), this, droppedDuplicates)
       } else {
-        (HashJoinBuilder.cloneHashTable(hashTableData(droppedDuplicates)), null, droppedDuplicates)
+        (
+          HashJoinBuilder.cloneHashTable(
+            broadcastContext.buildHashTableId,
+            hashTableData(droppedDuplicates)),
+          null,
+          droppedDuplicates)
       }
     }
 
diff --git a/cpp/velox/jni/VeloxJniWrapper.cc b/cpp/velox/jni/VeloxJniWrapper.cc
@@ -46,6 +46,7 @@
 #include "velox/common/base/BloomFilter.h"
 #include "velox/common/file/FileSystems.h"
 #include "velox/exec/HashTable.h"
+#include "velox/exec/HashTableCache.h"
 
 #ifdef GLUTEN_ENABLE_GPU
 #include "cudf/CudfPlanValidator.h"
@@ -949,7 +950,7 @@ JNIEXPORT jobject JNICALL Java_org_apache_gluten_execution_IcebergWriteJniWrappe
 JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_nativeBuild( // NOLINT
     JNIEnv* env,
     jobject wrapper,
-    jstring /*tableId*/,
+    jstring tableId,
     jlongArray batchHandles,
     jobjectArray joinKeys,
     jobjectArray filterBuildColumns,
@@ -974,6 +975,8 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_native
       queryConf.get<uint32_t>(kAbandonDedupHashMapMinRows, kAbandonDedupHashMapMinRowsDefault);
   const auto abandonHashBuildDedupMinPct =
       queryConf.get<uint32_t>(kAbandonDedupHashMapMinPct, kAbandonDedupHashMapMinPctDefault);
+  const auto hashTableId = jStringToCString(env, tableId);
+
   // Convert Java String array to C++ vector<string>
   std::vector<std::string> hashJoinKeys;
   jsize joinKeysCount = env->GetArrayLength(joinKeys);
@@ -1052,6 +1055,12 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_native
         nullptr);
     builder->setHashTable(std::move(mainTable));
 
+    auto* cache = facebook::velox::exec::HashTableCache::instance();
+
+    if (!cache->exist(hashTableId)) {
+      cache->add(hashTableId, builder->hashTable(), builder->joinHasNullKeys(), defaultLeafVeloxMemoryPool());
+    }
+
     return gluten::getHashTableObjStore()->save(builder);
   }
 
@@ -1134,44 +1143,69 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_native
   }
 
   hashTableBuilders[0]->setHashTable(std::move(mainTable));
+
+  auto* cache = facebook::velox::exec::HashTableCache::instance();
+  if (!cache->exist(hashTableId)) {
+    cache->add(
+        hashTableId,
+        hashTableBuilders[0]->hashTable(),
+        hashTableBuilders[0]->joinHasNullKeys(),
+        defaultLeafVeloxMemoryPool());
+  }
+
   return gluten::getHashTableObjStore()->save(hashTableBuilders[0]);
   JNI_METHOD_END(kInvalidObjectHandle)
 }
 
 JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_cloneHashTable( // NOLINT
     JNIEnv* env,
     jclass,
+    jstring cacheKey,
     jlong tableHandler) {
   JNI_METHOD_START
+  auto cacheKeyStr = jStringToCString(env, cacheKey);
   auto hashTableHandler = ObjectStore::retrieve<gluten::HashTableBuilder>(tableHandler);
+  auto* cache = facebook::velox::exec::HashTableCache::instance();
+  if (!cache->exist(cacheKeyStr)) {
+    cache->add(
+        cacheKeyStr, hashTableHandler->hashTable(), hashTableHandler->joinHasNullKeys(), defaultLeafVeloxMemoryPool());
+  }
+
   return gluten::getHashTableObjStore()->save(hashTableHandler);
   JNI_METHOD_END(kInvalidObjectHandle)
 }
 
 JNIEXPORT void JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_clearHashTable( // NOLINT
     JNIEnv* env,
     jclass,
+    jstring cacheKey,
     jlong tableHandler) {
   JNI_METHOD_START
-  auto hashTableHandler = ObjectStore::retrieve<gluten::HashTableBuilder>(tableHandler);
-  hashTableHandler->hashTable()->clear(true);
+  auto cacheKeyStr = jStringToCString(env, cacheKey);
+  facebook::velox::exec::HashTableCache::instance()->drop(cacheKeyStr);
   ObjectStore::release(tableHandler);
   JNI_METHOD_END()
 }
 
 JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_deserializeHashTableDirect( // NOLINT
     JNIEnv* env,
     jclass,
+    jstring cacheKey,
     jlong address,
     jint size,
     jboolean ignoreNullKeys,
     jboolean joinHasNullKeys) {
   JNI_METHOD_START
+  auto cacheKeyStr = jStringToCString(env, cacheKey);
   auto builder = gluten::deserializeHashTable(
       reinterpret_cast<const uint8_t*>(address),
       static_cast<size_t>(size),
       static_cast<bool>(ignoreNullKeys),
       static_cast<bool>(joinHasNullKeys));
+  auto* cache = facebook::velox::exec::HashTableCache::instance();
+  if (!cache->exist(cacheKeyStr)) {
+    cache->add(cacheKeyStr, builder->hashTable(), builder->joinHasNullKeys(), defaultLeafVeloxMemoryPool());
+  }
   return gluten::getHashTableObjStore()->save(builder);
   JNI_METHOD_END(kInvalidObjectHandle)
 }
diff --git a/cpp/velox/substrait/SubstraitToVeloxPlan.cc b/cpp/velox/substrait/SubstraitToVeloxPlan.cc
@@ -458,26 +458,6 @@ core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
   } else if (
       sJoin.has_advanced_extension() &&
       SubstraitParser::configSetInOptimization(sJoin.advanced_extension(), "isBHJ=")) {
-    std::string hashTableId = sJoin.hashtableid();
-
-    std::shared_ptr<core::OpaqueHashTable> opaqueSharedHashTable = nullptr;
-    bool joinHasNullKeys = false;
-
-    try {
-      auto hashTableBuilder = ObjectStore::retrieve<gluten::HashTableBuilder>(getJoin(hashTableId));
-      joinHasNullKeys = hashTableBuilder->joinHasNullKeys();
-      auto originalShared = hashTableBuilder->hashTable();
-      opaqueSharedHashTable = std::shared_ptr<core::OpaqueHashTable>(
-          originalShared, reinterpret_cast<core::OpaqueHashTable*>(originalShared.get()));
-
-      LOG(INFO) << "Successfully retrieved and aliased HashTable for reuse. ID: " << hashTableId;
-    } catch (const std::exception& e) {
-      LOG(WARNING)
-          << "Error retrieving HashTable from ObjectStore: " << e.what()
-          << ". Falling back to building new table. To ensure correct results, please verify that spark.gluten.velox.buildHashTableOncePerExecutor.enabled is set to false.";
-      opaqueSharedHashTable = nullptr;
-    }
-
     // Create HashJoinNode node
     return std::make_shared<core::HashJoinNode>(
         nextPlanNodeId(),
@@ -489,10 +469,11 @@ core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
         leftNode,
         rightNode,
         getJoinOutputType(leftNode, rightNode, joinType),
+        true,
         false,
         false,
-        joinHasNullKeys,
-        opaqueSharedHashTable);
+        nullptr,
+        sJoin.hashtableid());
   } else {
     // Create HashJoinNode node
     return std::make_shared<core::HashJoinNode>(
diff --git a/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/execution/adaptive/velox/VeloxAdaptiveQueryExecSuite.scala b/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/execution/adaptive/velox/VeloxAdaptiveQueryExecSuite.scala
@@ -522,7 +522,9 @@ class VeloxAdaptiveQueryExecSuite extends AdaptiveQueryExecSuite with GlutenSQLT
     withSQLConf(
       SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
       SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "20000000",
-      SQLConf.SUBQUERY_REUSE_ENABLED.key -> "false") {
+      SQLConf.ADAPTIVE_OPTIMIZER_EXCLUDED_RULES.key -> AQEPropagateEmptyRelation.ruleName,
+      SQLConf.SUBQUERY_REUSE_ENABLED.key -> "false"
+    ) {
       val (plan, adaptivePlan) = runAdaptiveAndVerifyResult(
         "SELECT a FROM testData join testData2 ON key = a " +
           "where value >= (" +

Original file line number	Diff line number	Diff line change
`@@ -161,6 +161,7 @@ object VeloxBroadcastBuildSideCache`
`161`	`161`	`val result =`
`162`	`162`	`SerializedBroadcastHashTable.fromHashTable(`
`163`	`163`	`hashTableHandle,`
	`164`	`+ broadcastId,`
`164`	`165`	`relation,`
`165`	`166`	`droppedDuplicates,`
`166`	`167`	`numRows)`
`@@ -197,7 +198,7 @@ object VeloxBroadcastBuildSideCache`
`197`	`198`	`(_: String) => {`
`198`	`199`	`logInfo(s"Deserializing hash table on executor for broadcast ID: $broadcastHashTableId")`
`199`	`200`	`val startTime = System.currentTimeMillis()`
`200`		`- val hashTableHandle = serialized.deserialize()`
	`201`	`+ val hashTableHandle = serialized.deserialize(broadcastHashTableId)`
`201`	`202`	`val timeMs = System.currentTimeMillis() - startTime`
`202`	`203`	`deserializeHashTableTimeMetric.foreach(_ += timeMs)`
`203`	`204`	`BroadcastHashTable(`
`@@ -245,7 +246,7 @@ object VeloxBroadcastBuildSideCache`
`245`	`246`	`}`
`246`	`247`	`}`
`247`	`248`
`248`		`- HashJoinBuilder.clearHashTable(value.pointer)`
	`249`	`+ HashJoinBuilder.clearHashTable(key, value.pointer)`
`249`	`250`	`}`
`250`	`251`	`}`
`251`	`252`	`}`
Original file line number	Diff line number	Diff line change
`@@ -242,7 +242,9 @@ case class ColumnarBuildSideRelation(`
`242`	`242`	`(hashTableData(droppedDuplicates), this, droppedDuplicates)`
`243`	`243`	`} else {`
`244`	`244`	`(`
`245`		`- HashJoinBuilder.cloneHashTable(hashTableData(droppedDuplicates)),`
	`245`	`+ HashJoinBuilder.cloneHashTable(`
	`246`	`+ broadcastContext.buildHashTableId,`
	`247`	`+ hashTableData(droppedDuplicates)),`
`246`	`248`	`null,`
`247`	`249`	`droppedDuplicates)`
`248`	`250`	`}`
`@@ -332,7 +334,12 @@ case class ColumnarBuildSideRelation(`
`332`	`334`
`333`	`335`	`(hashTableData(droppedDuplicates), this, droppedDuplicates)`
`334`	`336`	`} else {`
`335`		`- (HashJoinBuilder.cloneHashTable(hashTableData(droppedDuplicates)), null, droppedDuplicates)`
	`337`	`+ (`
	`338`	`+ HashJoinBuilder.cloneHashTable(`
	`339`	`+ broadcastContext.buildHashTableId,`
	`340`	`+ hashTableData(droppedDuplicates)),`
	`341`	`+ null,`
	`342`	`+ droppedDuplicates)`
`336`	`343`	`}`
`337`	`344`	`}`
`338`	`345`
Original file line number	Diff line number	Diff line change
`@@ -214,7 +214,9 @@ class UnsafeColumnarBuildSideRelation(`
`214`	`214`	`(hashTableData(droppedDuplicates), this, droppedDuplicates)`
`215`	`215`	`} else {`
`216`	`216`	`(`
`217`		`- HashJoinBuilder.cloneHashTable(hashTableData(droppedDuplicates)),`
	`217`	`+ HashJoinBuilder.cloneHashTable(`
	`218`	`+ broadcastContext.buildHashTableId,`
	`219`	`+ hashTableData(droppedDuplicates)),`
`218`	`220`	`null,`
`219`	`221`	`droppedDuplicates)`
`220`	`222`	`}`
`@@ -305,7 +307,12 @@ class UnsafeColumnarBuildSideRelation(`
`305`	`307`
`306`	`308`	`(hashTableData(droppedDuplicates), this, droppedDuplicates)`
`307`	`309`	`} else {`
`308`		`- (HashJoinBuilder.cloneHashTable(hashTableData(droppedDuplicates)), null, droppedDuplicates)`
	`310`	`+ (`
	`311`	`+ HashJoinBuilder.cloneHashTable(`
	`312`	`+ broadcastContext.buildHashTableId,`
	`313`	`+ hashTableData(droppedDuplicates)),`
	`314`	`+ null,`
	`315`	`+ droppedDuplicates)`
`309`	`316`	`}`
`310`	`317`	`}`
`311`	`318`