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 index dbd8a3a5f70..675db72a844 100644 --- 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 index e77b2c61154..58120aadc39 100644 --- 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 index a59326daef8..323b88a4203 100644 --- 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 index 0d492817b14..2c68028e9d2 100644 --- 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 index 35587bde133..7374cf0b04d 100644 --- 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 index 2a6963bb727..d0129e887cf 100644 --- 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 index 36798f5cded..db189fdd14c 100644 --- 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(kAbandonDedupHashMapMinRows, kAbandonDedupHashMapMinRowsDefault); const auto abandonHashBuildDedupMinPct = queryConf.get(kAbandonDedupHashMapMinPct, kAbandonDedupHashMapMinPctDefault); + const auto hashTableId = jStringToCString(env, tableId); + // Convert Java String array to C++ vector std::vector 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,6 +1143,16 @@ 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) } @@ -1141,9 +1160,17 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_native 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(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) } @@ -1151,10 +1178,11 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_cloneH 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(tableHandler); - hashTableHandler->hashTable()->clear(true); + auto cacheKeyStr = jStringToCString(env, cacheKey); + facebook::velox::exec::HashTableCache::instance()->drop(cacheKeyStr); ObjectStore::release(tableHandler); JNI_METHOD_END() } @@ -1162,16 +1190,22 @@ JNIEXPORT void JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_clearHa 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(address), static_cast(size), static_cast(ignoreNullKeys), static_cast(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 index 5bd499582b8..bdc6d35c29c 100644 --- 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 opaqueSharedHashTable = nullptr; - bool joinHasNullKeys = false; - - try { - auto hashTableBuilder = ObjectStore::retrieve(getJoin(hashTableId)); - joinHasNullKeys = hashTableBuilder->joinHasNullKeys(); - auto originalShared = hashTableBuilder->hashTable(); - opaqueSharedHashTable = std::shared_ptr( - originalShared, reinterpret_cast(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( 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( 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 index 6ccddac0c62..6f1272b7f77 100644 --- 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 >= (" +