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 675db72a844..875dbb1cdbc 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 @@ -48,6 +48,8 @@ public static native long deserializeHashTableDirect( public static native long getHashTableBloomFilterBlocksByteSize(long hashTableHandle); + public static native long getHashTableMemoryUsage(long hashTableHandle); + public static native long serializedHashTableSizeDirect(long hashTableHandle); public static native void serializeHashTableDirect(long hashTableHandle, long address, long size); diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxMetricsApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxMetricsApi.scala index 272574fe79b..421e6c35959 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxMetricsApi.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxMetricsApi.scala @@ -727,7 +727,10 @@ class VeloxMetricsApi extends MetricsApi with Logging { "time to build hash table"), "deserializeHashTableTime" -> SQLMetrics.createTimingMetric( sparkContext, - "time to deserialize hash table") + "time to deserialize hash table"), + "hashTableMemorySize" -> SQLMetrics.createSizeMetric( + sparkContext, + "hash table memory size") ) override def genHashJoinTransformerMetricsUpdater( 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 58120aadc39..2c1c976e914 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 @@ -181,7 +181,8 @@ case class BroadcastHashJoinExecTransformer( metrics.get("buildHashTableTime"), metrics.get("serializeHashTableTime"), metrics.get("deserializeHashTableTime"), - metrics.get("serializedHashTableSize") + metrics.get("serializedHashTableSize"), + metrics.get("hashTableMemorySize") ) // Check the type of broadcast relation to determine the approach @@ -265,7 +266,8 @@ case class BroadcastHashJoinContext( buildHashTableTimeMetric: Option[SQLMetric] = None, serializeHashTableTimeMetric: Option[SQLMetric] = None, deserializeHashTableTimeMetric: Option[SQLMetric] = None, - serializedHashTableSizeMetric: Option[SQLMetric] = None) { + serializedHashTableSizeMetric: Option[SQLMetric] = None, + hashTableMemorySizeMetric: Option[SQLMetric] = None) { def droppedDuplicates: Boolean = { !hasMixedFiltCondition && ( substraitJoinType == JoinRel.JoinType.JOIN_TYPE_LEFT_SEMI || 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 2c68028e9d2..02b18b5d2be 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 @@ -98,6 +98,9 @@ object VeloxBroadcastBuildSideCache unsafe.buildHashTable(broadcastContext) } + broadcastContext.hashTableMemorySizeMetric.foreach( + _ += HashJoinBuilder.getHashTableMemoryUsage(pointer)) + BroadcastHashTable(pointer, relation, droppedDuplicates) } ) diff --git a/cpp/velox/jni/JniHashTable.cc b/cpp/velox/jni/JniHashTable.cc index d7259879c59..7db693e3d00 100644 --- a/cpp/velox/jni/JniHashTable.cc +++ b/cpp/velox/jni/JniHashTable.cc @@ -180,6 +180,11 @@ size_t serializedHashTableSize(std::shared_ptr builder) { return HashTableSerializer::serializedSize(hashTableTrue); } +int64_t hashTableMemoryUsage(std::shared_ptr builder) { + VELOX_CHECK_NOT_NULL(builder, "Hash table builder cannot be null"); + return builder->hashTableMemoryUsage(); +} + void serializeHashTableTo(std::shared_ptr builder, uint8_t* data, size_t size) { VELOX_CHECK_NOT_NULL(builder, "Hash table builder cannot be null"); VELOX_CHECK_NOT_NULL(data, "Serialized buffer cannot be null"); diff --git a/cpp/velox/jni/JniHashTable.h b/cpp/velox/jni/JniHashTable.h index d75d2d663bc..48bed2b4f00 100644 --- a/cpp/velox/jni/JniHashTable.h +++ b/cpp/velox/jni/JniHashTable.h @@ -94,6 +94,9 @@ long getJoin(const std::string& hashTableId); // Return the exact serialized hash table size for direct buffer allocation. size_t serializedHashTableSize(std::shared_ptr builder); +// Return retained bytes for the hash table, including parallel-build subtables. +int64_t hashTableMemoryUsage(std::shared_ptr builder); + // Serialize hash table directly to a caller-provided buffer. void serializeHashTableTo(std::shared_ptr builder, uint8_t* data, size_t size); diff --git a/cpp/velox/jni/VeloxJniWrapper.cc b/cpp/velox/jni/VeloxJniWrapper.cc index db189fdd14c..fa465755bf2 100644 --- a/cpp/velox/jni/VeloxJniWrapper.cc +++ b/cpp/velox/jni/VeloxJniWrapper.cc @@ -1054,6 +1054,7 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_native builder->dropDuplicates(), nullptr); builder->setHashTable(std::move(mainTable)); + builder->setHashTableMemoryUsage(builder->hashTable()->allocatedBytes()); auto* cache = facebook::velox::exec::HashTableCache::instance(); @@ -1122,6 +1123,10 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_native for (int i = 1; i < numThreads; ++i) { tables.push_back(std::move(otherTables[i])); } + int64_t otherTablesMemoryUsage = 0; + for (const auto& table : tables) { + otherTablesMemoryUsage += table->allocatedBytes(); + } // TODO: Get accurate signal if parallel join build is going to be applied // from hash table. Currently there is still a chance inside hash table that @@ -1143,6 +1148,8 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_native } hashTableBuilders[0]->setHashTable(std::move(mainTable)); + hashTableBuilders[0]->setHashTableMemoryUsage( + hashTableBuilders[0]->hashTable()->allocatedBytes() + otherTablesMemoryUsage); auto* cache = facebook::velox::exec::HashTableCache::instance(); if (!cache->exist(hashTableId)) { @@ -1260,6 +1267,16 @@ Java_org_apache_gluten_vectorized_HashJoinBuilder_getHashTableBloomFilterBlocksB JNI_METHOD_END(0L) } +JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_getHashTableMemoryUsage( // NOLINT + JNIEnv* env, + jclass, + jlong hashTableHandle) { + JNI_METHOD_START + auto builder = ObjectStore::retrieve(hashTableHandle); + return static_cast(gluten::hashTableMemoryUsage(builder)); + JNI_METHOD_END(0L) +} + JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_serializedHashTableSizeDirect( // NOLINT JNIEnv* env, jclass, diff --git a/cpp/velox/operators/hashjoin/HashTableBuilder.h b/cpp/velox/operators/hashjoin/HashTableBuilder.h index 98e459b0141..44769e7490f 100644 --- a/cpp/velox/operators/hashjoin/HashTableBuilder.h +++ b/cpp/velox/operators/hashjoin/HashTableBuilder.h @@ -75,6 +75,14 @@ class HashTableBuilder { return noMoreInput_; } + void setHashTableMemoryUsage(int64_t hashTableMemoryUsage) { + hashTableMemoryUsage_ = hashTableMemoryUsage; + } + + int64_t hashTableMemoryUsage() const { + return hashTableMemoryUsage_; + } + uint32_t joinBuildVectorHasherMaxNumDistinct() const { return joinBuildVectorHasherMaxNumDistinct_; } @@ -150,6 +158,7 @@ class HashTableBuilder { uint32_t abandonHashBuildDedupMinRows_{100'000}; uint32_t abandonHashBuildDedupMinPct_{0}; bool filterPropagatesNulls_{false}; + int64_t hashTableMemoryUsage_{0}; }; } // namespace gluten