Add hash table memory usage metric

Author: JkSelf

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);

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(

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 ||

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)
         }
       )

cpp/velox/jni/JniHashTable.cc

@@ -180,6 +180,11 @@ size_t serializedHashTableSize(std::shared_ptr<HashTableBuilder> builder) {
   return HashTableSerializer::serializedSize<true>(hashTableTrue);
 }
 
+int64_t hashTableMemoryUsage(std::shared_ptr<HashTableBuilder> builder) {
+  VELOX_CHECK_NOT_NULL(builder, "Hash table builder cannot be null");
+  return builder->hashTableMemoryUsage();
+}
+
 void serializeHashTableTo(std::shared_ptr<HashTableBuilder> 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");

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<HashTableBuilder> builder);
 
+// Return retained bytes for the hash table, including parallel-build subtables.
+int64_t hashTableMemoryUsage(std::shared_ptr<HashTableBuilder> builder);
+
 // Serialize hash table directly to a caller-provided buffer.
 void serializeHashTableTo(std::shared_ptr<HashTableBuilder> builder, uint8_t* data, size_t size);
 

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<gluten::HashTableBuilder>(hashTableHandle);
+  return static_cast<jlong>(gluten::hashTableMemoryUsage(builder));
+  JNI_METHOD_END(0L)
+}
+
 JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_serializedHashTableSizeDirect( // NOLINT
     JNIEnv* env,
     jclass,

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