[VL][TEST]Use Velox's HashTableCache to cache the BHJ's HashTable

backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHIteratorApi.scala

@@ -259,6 +259,7 @@ class CHIteratorApi extends IteratorApi with Logging with LogLevelUtil {
         GlutenPartition(
           index,
           planByteArray,
+          wsCtx.queryId,
           splitInfos.toArray
         )
     }
@@ -315,6 +316,7 @@ class CHIteratorApi extends IteratorApi with Logging with LogLevelUtil {
       inputIterators: Seq[Iterator[ColumnarBatch]],
       sparkConf: SparkConf,
       rootNode: PlanNode,
+      queryId: String,
       pipelineTime: SQLMetric,
       updateNativeMetrics: IMetrics => Unit,
       partitionIndex: Int,

backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHSparkPlanExecApi.scala

@@ -481,7 +481,8 @@ class CHSparkPlanExecApi extends SparkPlanExecApi with Logging {
       child: SparkPlan,
       numOutputRows: SQLMetric,
       dataSize: SQLMetric,
-      buildThreads: SQLMetric): BuildSideRelation = {
+      buildThreads: SQLMetric,
+      queryId: String): BuildSideRelation = {
 
     val (buildKeys, isNullAware) = mode match {
       case mode1: HashedRelationBroadcastMode =>

backends-velox/src-delta33/main/scala/org/apache/spark/sql/delta/stats/GlutenDeltaJobStatsTracker.scala

@@ -194,7 +194,8 @@ object GlutenDeltaJobStatsTracker extends Logging {
           null,
           null,
           0,
-          BackendsApiManager.getSparkPlanExecApiInstance.rewriteSpillPath(spillDirPath)
+          BackendsApiManager.getSparkPlanExecApiInstance.rewriteSpillPath(spillDirPath),
+          s"Gluten_Delta_Stats_${UUID.randomUUID()}"
         )
       nativeOutItr
     }

backends-velox/src-delta40/main/scala/org/apache/spark/sql/delta/stats/GlutenDeltaJobStatsTracker.scala

@@ -198,7 +198,8 @@ object GlutenDeltaJobStatsTracker extends Logging {
           null,
           null,
           0,
-          BackendsApiManager.getSparkPlanExecApiInstance.rewriteSpillPath(spillDirPath)
+          BackendsApiManager.getSparkPlanExecApiInstance.rewriteSpillPath(spillDirPath),
+          s"Gluten_Delta_Stats_${UUID.randomUUID()}"
         )
       nativeOutItr
     }

backends-velox/src/main/java/org/apache/gluten/vectorized/HashJoinBuilder.java

@@ -35,9 +35,9 @@ 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 native long nativeBuild(
       String buildHashTableId,
@@ -51,5 +51,6 @@ public native long nativeBuild(
       byte[] namedStruct,
       boolean isNullAwareAntiJoin,
       long bloomFilterPushdownSize,
-      int broadcastHashTableBuildThreads);
+      int broadcastHashTableBuildThreads,
+      String queryId);
 }

backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxIteratorApi.scala

@@ -130,6 +130,7 @@ class VeloxIteratorApi extends IteratorApi with Logging {
         GlutenPartition(
           index,
           planByteArray,
+          wsCtx.queryId,
           splitInfos.toArray
         )
     }
@@ -216,7 +217,8 @@ class VeloxIteratorApi extends IteratorApi with Logging {
         if (splitInfoByteArray.nonEmpty) splitInfoByteArray else null,
         if (columnarNativeIterators.nonEmpty) columnarNativeIterators.toArray else null,
         partitionIndex,
-        BackendsApiManager.getSparkPlanExecApiInstance.rewriteSpillPath(spillDirPath)
+        BackendsApiManager.getSparkPlanExecApiInstance.rewriteSpillPath(spillDirPath),
+        inputPartition.queryId
       )
     resIter.noMoreSplits()
     val itrMetrics = IteratorMetricsJniWrapper.create()
@@ -243,6 +245,7 @@ class VeloxIteratorApi extends IteratorApi with Logging {
       inputIterators: Seq[Iterator[ColumnarBatch]],
       sparkConf: SparkConf,
       rootNode: PlanNode,
+      queryId: String,
       pipelineTime: SQLMetric,
       updateNativeMetrics: IMetrics => Unit,
       partitionIndex: Int,
@@ -270,7 +273,8 @@ class VeloxIteratorApi extends IteratorApi with Logging {
         null,
         if (columnarNativeIterator.nonEmpty) columnarNativeIterator.toArray else null,
         partitionIndex,
-        BackendsApiManager.getSparkPlanExecApiInstance.rewriteSpillPath(spillDirPath)
+        BackendsApiManager.getSparkPlanExecApiInstance.rewriteSpillPath(spillDirPath),
+        queryId
       )
     nativeResultIterator.noMoreSplits()
     val itrMetrics = IteratorMetricsJniWrapper.create()

backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala

@@ -715,7 +715,8 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with Logging {
       child: SparkPlan,
       numOutputRows: SQLMetric,
       dataSize: SQLMetric,
-      buildThreads: SQLMetric): BuildSideRelation = {
+      buildThreads: SQLMetric,
+      queryId: String): BuildSideRelation = {
 
     val buildKeys = mode match {
       case mode1: HashedRelationBroadcastMode =>
@@ -875,7 +876,8 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with Logging {
           mode,
           newBuildKeys,
           offload,
-          buildThreadsValue)
+          buildThreadsValue,
+          queryId)
       }
     } else {
       ColumnarBuildSideRelation(
@@ -884,7 +886,8 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with Logging {
         mode,
         newBuildKeys,
         offload,
-        buildThreadsValue)
+        buildThreadsValue,
+        queryId)
     }
   }
 

backends-velox/src/main/scala/org/apache/gluten/execution/HashJoinExecTransformer.scala

@@ -26,8 +26,10 @@ import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.optimizer.{BuildRight, BuildSide}
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.execution.{SparkPlan, SQLExecution}
+import org.apache.spark.sql.execution.ColumnarBuildSideRelation
 import org.apache.spark.sql.execution.joins.BuildSideRelation
 import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.execution.unsafe.UnsafeColumnarBuildSideRelation
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
 import io.substrait.proto.JoinRel
@@ -106,7 +108,7 @@ case class BroadcastHashJoinExecTransformer(
     isNullAwareAntiJoin) {
 
   // Unique ID for built table
-  lazy val buildBroadcastTableId: String = buildPlan.id.toString
+  lazy val buildBroadcastTableId: String = canonicalBuildHashTableId(buildPlan)
 
   override protected lazy val substraitJoinType: JoinRel.JoinType = joinType match {
     case _: InnerLike =>
@@ -134,17 +136,18 @@ case class BroadcastHashJoinExecTransformer(
 
   override def columnarInputRDDs: Seq[RDD[ColumnarBatch]] = {
     val streamedRDD = getColumnarInputRDDs(streamedPlan)
+    val broadcast = buildPlan.executeBroadcast[BuildSideRelation]()
+    val queryId = extractQueryId(broadcast.value)
     val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
+    val cacheKey = s"$queryId:$buildHashTableId"
     if (executionId != null) {
-      GlutenDriverEndpoint.collectResources(executionId, buildBroadcastTableId)
+      GlutenDriverEndpoint.collectResources(executionId, cacheKey)
     } else {
       logWarning(
-        s"Can not trace broadcast table data $buildBroadcastTableId" +
+        s"Can not trace broadcast table data $cacheKey" +
           s" because execution id is null." +
           s" Will clean up until expire time.")
     }
-
-    val broadcast = buildPlan.executeBroadcast[BuildSideRelation]()
     val bloomFilterPushdownSize = if (VeloxConfig.get.hashProbeDynamicFilterPushdownEnabled) {
       VeloxConfig.get.hashProbeBloomFilterPushdownMaxSize
     } else {
@@ -174,7 +177,7 @@ case class BroadcastHashJoinExecTransformer(
         buildPlan.output,
         filterBuildColumns,
         filterPropagatesNulls,
-        buildBroadcastTableId,
+        cacheKey,
         isNullAwareAntiJoin,
         bloomFilterPushdownSize,
         metrics.get("buildHashTableTime")
@@ -183,6 +186,15 @@ case class BroadcastHashJoinExecTransformer(
     // FIXME: Do we have to make build side a RDD?
     streamedRDD :+ broadcastRDD
   }
+
+  private def extractQueryId(relation: BuildSideRelation): String = relation match {
+    case columnar: ColumnarBuildSideRelation if columnar.queryId.nonEmpty =>
+      columnar.queryId
+    case unsafe: UnsafeColumnarBuildSideRelation if unsafe.queryIdValue.nonEmpty =>
+      unsafe.queryIdValue
+    case _ =>
+      GlutenQueryContext.getRequiredQueryId(buildPlan)
+  }
 }
 
 case class BroadcastHashJoinContext(

backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastBuildSideCache.scala

@@ -103,7 +103,7 @@ object VeloxBroadcastBuildSideCache
         }
       }
 
-      HashJoinBuilder.clearHashTable(value.pointer)
+      HashJoinBuilder.clearHashTable(key, value.pointer)
     }
   }
 }

backends-velox/src/main/scala/org/apache/spark/sql/execution/ColumnarBuildSideRelation.scala

@@ -52,7 +52,8 @@ object ColumnarBuildSideRelation {
       mode: BroadcastMode,
       newBuildKeys: Seq[Expression] = Seq.empty,
       offload: Boolean = false,
-      buildThreads: Int = 1): ColumnarBuildSideRelation = {
+      buildThreads: Int = 1,
+      queryId: String = ""): ColumnarBuildSideRelation = {
     val boundMode = mode match {
       case HashedRelationBroadcastMode(keys, isNullAware) =>
         // Bind each key to the build-side output so simple cols become BoundReference
@@ -68,7 +69,8 @@ object ColumnarBuildSideRelation {
       BroadcastModeUtils.toSafe(boundMode),
       newBuildKeys,
       offload,
-      buildThreads)
+      buildThreads,
+      queryId)
   }
 }
 
@@ -78,7 +80,8 @@ case class ColumnarBuildSideRelation(
     safeBroadcastMode: SafeBroadcastMode,
     newBuildKeys: Seq[Expression],
     offload: Boolean,
-    buildThreads: Int)
+    buildThreads: Int,
+    queryId: String)
   extends BuildSideRelation
   with Logging
   with KnownSizeEstimation {
@@ -223,7 +226,8 @@ case class ColumnarBuildSideRelation(
             SubstraitUtil.toNameStruct(newOutput).toByteArray,
             broadcastContext.isNullAwareAntiJoin,
             broadcastContext.bloomFilterPushdownSize,
-            buildThreads
+            buildThreads,
+            queryId
           )
 
         jniWrapper.close(serializeHandle)
@@ -234,7 +238,7 @@ case class ColumnarBuildSideRelation(
 
         (hashTableData, this)
       } else {
-        (HashJoinBuilder.cloneHashTable(hashTableData), null)
+        (HashJoinBuilder.cloneHashTable(broadcastContext.buildHashTableId, hashTableData), null)
       }
     }
 

backends-velox/src/main/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelation.scala

@@ -57,7 +57,8 @@ object UnsafeColumnarBuildSideRelation {
       mode: BroadcastMode,
       newBuildKeys: Seq[Expression] = Seq.empty,
       offload: Boolean = false,
-      buildThreads: Int = 1): UnsafeColumnarBuildSideRelation = {
+      buildThreads: Int = 1,
+      queryId: String = ""): UnsafeColumnarBuildSideRelation = {
     val boundMode = mode match {
       case HashedRelationBroadcastMode(keys, isNullAware) =>
         // Bind each key to the build-side output so simple cols become BoundReference
@@ -73,7 +74,8 @@ object UnsafeColumnarBuildSideRelation {
       BroadcastModeUtils.toSafe(boundMode),
       newBuildKeys,
       offload,
-      buildThreads)
+      buildThreads,
+      queryId)
   }
 }
 
@@ -94,7 +96,8 @@ class UnsafeColumnarBuildSideRelation(
     private var safeBroadcastMode: SafeBroadcastMode,
     private var newBuildKeys: Seq[Expression],
     private var offload: Boolean,
-    private var buildThreads: Int)
+    private var buildThreads: Int,
+    private var queryId: String)
   extends BuildSideRelation
   with Externalizable
   with Logging
@@ -114,9 +117,11 @@ class UnsafeColumnarBuildSideRelation(
 
   def isOffload: Boolean = offload
 
+  def queryIdValue: String = queryId
+
   /** needed for serialization. */
   def this() = {
-    this(null, null, null, Seq.empty, false, 1)
+    this(null, null, null, Seq.empty, false, 1, "")
   }
 
   private[unsafe] def getBatches(): Seq[UnsafeByteArray] = {
@@ -193,7 +198,8 @@ class UnsafeColumnarBuildSideRelation(
             SubstraitUtil.toNameStruct(newOutput).toByteArray,
             broadcastContext.isNullAwareAntiJoin,
             broadcastContext.bloomFilterPushdownSize,
-            buildThreads
+            buildThreads,
+            queryId
           )
 
         jniWrapper.close(serializeHandle)
@@ -204,7 +210,7 @@ class UnsafeColumnarBuildSideRelation(
 
         (hashTableData, this)
       } else {
-        (HashJoinBuilder.cloneHashTable(hashTableData), null)
+        (HashJoinBuilder.cloneHashTable(broadcastContext.buildHashTableId, hashTableData), null)
       }
     }
 
@@ -219,6 +225,7 @@ class UnsafeColumnarBuildSideRelation(
     out.writeObject(newBuildKeys)
     out.writeBoolean(offload)
     out.writeInt(buildThreads)
+    out.writeUTF(queryId)
   }
 
   override def write(kryo: Kryo, out: Output): Unit = Utils.tryOrIOException {
@@ -228,6 +235,7 @@ class UnsafeColumnarBuildSideRelation(
     kryo.writeClassAndObject(out, newBuildKeys)
     out.writeBoolean(offload)
     out.writeInt(buildThreads)
+    out.writeString(queryId)
   }
 
   override def readExternal(in: ObjectInput): Unit = Utils.tryOrIOException {
@@ -237,6 +245,7 @@ class UnsafeColumnarBuildSideRelation(
     newBuildKeys = in.readObject().asInstanceOf[Seq[Expression]]
     offload = in.readBoolean()
     buildThreads = in.readInt()
+    queryId = in.readUTF()
   }
 
   override def read(kryo: Kryo, in: Input): Unit = Utils.tryOrIOException {
@@ -246,6 +255,7 @@ class UnsafeColumnarBuildSideRelation(
     newBuildKeys = kryo.readClassAndObject(in).asInstanceOf[Seq[Expression]]
     offload = in.readBoolean()
     buildThreads = in.readInt()
+    queryId = in.readString()
   }
 
   private def transformProjection: UnsafeProjection = safeBroadcastMode match {

cpp/core/compute/Runtime.h

@@ -42,12 +42,15 @@ struct SparkTaskInfo {
   int32_t partitionId{0};
   // Same as TID.
   int64_t taskId{0};
+  // Same as Spark SQL execution id. -1 means unavailable.
+  int64_t executionId{-1};
+  std::string queryId{};
   // virtual id for each backend internal use
   int32_t vId{0};
 
   std::string toString() const {
-    return "[Stage: " + std::to_string(stageId) + " TID: " + std::to_string(taskId) + " VID: " + std::to_string(vId) +
-        "]";
+    return "[Stage: " + std::to_string(stageId) + " Execution: " + std::to_string(executionId) +
+        " TID: " + std::to_string(taskId) + " VID: " + std::to_string(vId) + "]";
   }
 
   friend std::ostream& operator<<(std::ostream& os, const SparkTaskInfo& taskInfo) {

cpp/core/jni/JniWrapper.cc

@@ -468,13 +468,15 @@ Java_org_apache_gluten_vectorized_PlanEvaluatorJniWrapper_nativeCreateKernelWith
     jint stageId,
     jint partitionId,
     jlong taskId,
+    jlong executionId,
+    jstring queryId,
     jboolean enableDumping,
     jstring spillDir) {
   JNI_METHOD_START
 
   auto ctx = getRuntime(env, wrapper);
 
-  ctx->setSparkTaskInfo({stageId, partitionId, taskId});
+  ctx->setSparkTaskInfo({stageId, partitionId, taskId, executionId, jStringToCString(env, queryId)});
 
   if (enableDumping) {
     ctx->enableDumping();