|
46 | 46 | #include "velox/common/base/BloomFilter.h" |
47 | 47 | #include "velox/common/file/FileSystems.h" |
48 | 48 | #include "velox/exec/HashTable.h" |
| 49 | +#include "velox/exec/HashTableCache.h" |
49 | 50 |
|
50 | 51 | #ifdef GLUTEN_ENABLE_GPU |
51 | 52 | #include "cudf/CudfPlanValidator.h" |
@@ -1060,6 +1061,12 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_native |
1060 | 1061 | nullptr); |
1061 | 1062 | builder->setHashTable(std::move(mainTable)); |
1062 | 1063 |
|
| 1064 | + auto* cache = facebook::velox::exec::HashTableCache::instance(); |
| 1065 | + |
| 1066 | + if (!cache->hasTable(hashTableId)) { |
| 1067 | + cache->add(hashTableId, builder->hashTable(), builder->joinHasNullKeys(), defaultLeafVeloxMemoryPool()); |
| 1068 | + } |
| 1069 | + |
1063 | 1070 | return gluten::getHashTableObjStore()->save(builder); |
1064 | 1071 | } |
1065 | 1072 |
|
@@ -1138,30 +1145,50 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_native |
1138 | 1145 | } |
1139 | 1146 |
|
1140 | 1147 | hashTableBuilders[0]->setHashTable(std::move(mainTable)); |
| 1148 | + |
| 1149 | + auto* cache = facebook::velox::exec::HashTableCache::instance(); |
| 1150 | + if (!cache->hasTable(hashTableId)) { |
| 1151 | + cache->add( |
| 1152 | + hashTableId, |
| 1153 | + hashTableBuilders[0]->hashTable(), |
| 1154 | + hashTableBuilders[0]->joinHasNullKeys(), |
| 1155 | + defaultLeafVeloxMemoryPool()); |
| 1156 | + } |
| 1157 | + |
1141 | 1158 | return gluten::getHashTableObjStore()->save(hashTableBuilders[0]); |
1142 | 1159 | JNI_METHOD_END(kInvalidObjectHandle) |
1143 | 1160 | } |
1144 | 1161 |
|
1145 | 1162 | JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_cloneHashTable( // NOLINT |
1146 | 1163 | JNIEnv* env, |
1147 | 1164 | jclass, |
| 1165 | + jstring cacheKey, |
1148 | 1166 | jlong tableHandler) { |
1149 | 1167 | JNI_METHOD_START |
| 1168 | + auto cacheKeyStr = jStringToCString(env, cacheKey); |
1150 | 1169 | auto hashTableHandler = ObjectStore::retrieve<gluten::HashTableBuilder>(tableHandler); |
| 1170 | + auto* cache = facebook::velox::exec::HashTableCache::instance(); |
| 1171 | + if (!cache->hasTable(cacheKeyStr)) { |
| 1172 | + cache->add( |
| 1173 | + cacheKeyStr, hashTableHandler->hashTable(), hashTableHandler->joinHasNullKeys(), defaultLeafVeloxMemoryPool()); |
| 1174 | + } |
| 1175 | + |
1151 | 1176 | return gluten::getHashTableObjStore()->save(hashTableHandler); |
1152 | 1177 | JNI_METHOD_END(kInvalidObjectHandle) |
1153 | 1178 | } |
1154 | 1179 |
|
1155 | 1180 | JNIEXPORT void JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_clearHashTable( // NOLINT |
1156 | 1181 | JNIEnv* env, |
1157 | 1182 | jclass, |
| 1183 | + jstring cacheKey, |
1158 | 1184 | jlong tableHandler) { |
1159 | 1185 | JNI_METHOD_START |
1160 | | - auto hashTableHandler = ObjectStore::retrieve<gluten::HashTableBuilder>(tableHandler); |
1161 | | - hashTableHandler->hashTable()->clear(true); |
| 1186 | + auto cacheKeyStr = jStringToCString(env, cacheKey); |
| 1187 | + facebook::velox::exec::HashTableCache::instance()->drop(cacheKeyStr); |
1162 | 1188 | ObjectStore::release(tableHandler); |
1163 | 1189 | JNI_METHOD_END() |
1164 | 1190 | } |
| 1191 | + |
1165 | 1192 | #ifdef __cplusplus |
1166 | 1193 | } |
1167 | 1194 | #endif |
0 commit comments