This is an automated email from the ASF dual-hosted git repository.
JkSelf pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new be08b0ffdb [VL] Follow up optimizing BHJ by building hash table in
driver side like spark (#11993)
be08b0ffdb is described below
commit be08b0ffdb35bdc08e7194c4d2aa2d001ee35c9b
Author: JiaKe <[email protected]>
AuthorDate: Fri Jul 3 08:51:28 2026 +0100
[VL] Follow up optimizing BHJ by building hash table in driver side like
spark (#11993)
---
.../clickhouse/CHSparkPlanExecApi.scala | 5 +-
.../apache/gluten/vectorized/HashJoinBuilder.java | 13 ++
.../gluten/backendsapi/velox/VeloxMetricsApi.scala | 16 +-
.../backendsapi/velox/VeloxSparkPlanExecApi.scala | 227 +++++++++++++++++++--
.../org/apache/gluten/config/VeloxConfig.scala | 12 ++
.../gluten/execution/HashJoinExecTransformer.scala | 76 ++++++-
.../execution/SerializedBroadcastHashTable.scala | 155 ++++++++++++++
.../execution/VeloxBroadcastBuildSideCache.scala | 138 ++++++++++++-
.../execution/VeloxSerializedBroadcastRDD.scala | 66 ++++++
.../apache/gluten/metrics/JoinMetricsUpdater.scala | 10 +-
.../sql/execution/ColumnarBuildSideRelation.scala | 125 ++++++++++--
.../SerializedHashTableBroadcastRelation.scala | 137 +++++++++++++
.../unsafe/UnsafeColumnarBuildSideRelation.scala | 128 ++++++++++--
.../execution/DynamicOffHeapSizingSuite.scala | 1 +
.../gluten/execution/VeloxHashJoinSuite.scala | 144 ++++++++++++-
.../UnsafeColumnarBuildSideRelationTest.scala | 3 +
cpp/velox/CMakeLists.txt | 1 +
cpp/velox/jni/JniHashTable.cc | 83 ++++++++
cpp/velox/jni/JniHashTable.h | 10 +
cpp/velox/jni/VeloxJniWrapper.cc | 108 +++++++++-
cpp/velox/operators/hashjoin/HashTableBuilder.cc | 5 +-
.../operators/hashjoin/HashTableSerializer.cc | 62 ++++++
cpp/velox/operators/hashjoin/HashTableSerializer.h | 67 ++++++
docs/velox-configuration.md | 1 +
.../apache/gluten/memory/NativeMemoryManager.scala | 1 +
.../scala/org/apache/gluten/runtime/Runtime.scala | 1 +
.../gluten/extension/BroadcastJoinContextTag.scala | 46 +++++
.../gluten/extension/GlutenJoinKeysCapture.scala | 139 +++++++++++--
.../gluten/backendsapi/SparkPlanExecApi.scala | 5 +-
.../gluten/execution/JoinExecTransformer.scala | 3 +
.../apache/gluten/substrait/SubstraitContext.scala | 3 +
.../execution/ColumnarBroadcastExchangeExec.scala | 6 +-
32 files changed, 1712 insertions(+), 85 deletions(-)
diff --git
a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHSparkPlanExecApi.scala
b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHSparkPlanExecApi.scala
index 36c9813fbb..c00738434d 100644
---
a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHSparkPlanExecApi.scala
+++
b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHSparkPlanExecApi.scala
@@ -497,7 +497,10 @@ class CHSparkPlanExecApi extends SparkPlanExecApi with
Logging {
child: SparkPlan,
numOutputRows: SQLMetric,
dataSize: SQLMetric,
- buildThreads: SQLMetric): BuildSideRelation = {
+ buildThreads: SQLMetric,
+ buildHashTableTimeMetric: SQLMetric,
+ serializeHashTableTimeMetric: SQLMetric,
+ serializedHashTableSizeMetric: SQLMetric): BuildSideRelation = {
val (buildKeys, isNullAware) = mode match {
case mode1: HashedRelationBroadcastMode =>
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 106d8b9817..dbd8a3a5f7 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
@@ -39,6 +39,19 @@ public class HashJoinBuilder implements RuntimeAware {
public static native long cloneHashTable(long hashTableData);
+ public static native long deserializeHashTableDirect(
+ long address, int size, boolean ignoreNullKeys, boolean joinHasNullKeys);
+
+ public static native boolean getHashTableIgnoreNullKeys(long
hashTableHandle);
+
+ public static native boolean getHashTableJoinHasNullKeys(long
hashTableHandle);
+
+ public static native long getHashTableBloomFilterBlocksByteSize(long
hashTableHandle);
+
+ public static native long serializedHashTableSizeDirect(long
hashTableHandle);
+
+ public static native void serializeHashTableDirect(long hashTableHandle,
long address, long size);
+
public native long nativeBuild(
String buildHashTableId,
long[] batchHandlers,
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 83e83555fa..272574fe79 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
@@ -599,7 +599,16 @@ class VeloxMetricsApi extends MetricsApi with Logging {
"numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of
output rows"),
"collectTime" -> SQLMetrics.createTimingMetric(sparkContext, "time to
collect"),
"broadcastTime" -> SQLMetrics.createTimingMetric(sparkContext, "time to
broadcast"),
- "buildThreads" -> SQLMetrics.createMetric(sparkContext, "build threads")
+ "buildThreads" -> SQLMetrics.createMetric(sparkContext, "build threads"),
+ "driverBuildHashTableTime" -> SQLMetrics.createTimingMetric(
+ sparkContext,
+ "time to build hash table on driver"),
+ "driverSerializeHashTableTime" -> SQLMetrics.createTimingMetric(
+ sparkContext,
+ "time to serialize hash table on driver"),
+ "serializedHashTableSize" -> SQLMetrics.createSizeMetric(
+ sparkContext,
+ "serialized hash table size")
)
override def genColumnarSubqueryBroadcastMetrics(
@@ -715,7 +724,10 @@ class VeloxMetricsApi extends MetricsApi with Logging {
"time of loading lazy vectors"),
"buildHashTableTime" -> SQLMetrics.createTimingMetric(
sparkContext,
- "time to build hash table")
+ "time to build hash table"),
+ "deserializeHashTableTime" -> SQLMetrics.createTimingMetric(
+ sparkContext,
+ "time to deserialize hash table")
)
override def genHashJoinTransformerMetricsUpdater(
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala
b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala
index 2448e04e0c..d320fa0179 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala
+++
b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala
@@ -22,7 +22,7 @@ import org.apache.gluten.exception.{GlutenExceptionUtil,
GlutenNotSupportExcepti
import org.apache.gluten.execution._
import org.apache.gluten.expression._
import org.apache.gluten.expression.aggregate.{HLLAdapter,
VeloxBloomFilterAggregate, VeloxCollectList, VeloxCollectSet}
-import org.apache.gluten.extension.JoinKeysTag
+import org.apache.gluten.extension.{BroadcastJoinContextInfo,
BroadcastJoinContextTag}
import org.apache.gluten.extension.columnar.FallbackTags
import org.apache.gluten.shuffle.NeedCustomColumnarBatchSerializer
import org.apache.gluten.sql.shims.SparkShimLoader
@@ -44,7 +44,7 @@ import org.apache.spark.sql.catalyst.expressions._
import
org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression,
CollectList, CollectSet}
import org.apache.spark.sql.catalyst.expressions.objects.{AssertNotNull,
StaticInvoke}
import org.apache.spark.sql.catalyst.optimizer.BuildSide
-import org.apache.spark.sql.catalyst.plans.JoinType
+import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, FullOuter,
InnerLike, JoinType, LeftAnti, LeftOuter, LeftSemi, RightOuter}
import org.apache.spark.sql.catalyst.plans.physical._
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.adaptive.AQEShuffleReadExec
@@ -63,6 +63,7 @@ import org.apache.spark.sql.types._
import org.apache.spark.sql.vectorized.ColumnarBatch
import org.apache.spark.task.TaskResources
+import io.substrait.proto.JoinRel
import org.apache.commons.lang3.ClassUtils
import javax.ws.rs.core.UriBuilder
@@ -715,7 +716,23 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with
Logging {
child: SparkPlan,
numOutputRows: SQLMetric,
dataSize: SQLMetric,
- buildThreads: SQLMetric): BuildSideRelation = {
+ buildThreads: SQLMetric,
+ buildHashTableTimeMetric: SQLMetric,
+ serializeHashTableTimeMetric: SQLMetric,
+ serializedHashTableSizeMetric: SQLMetric): BuildSideRelation = {
+
+ @scala.annotation.tailrec
+ def findLogicalLink(
+ plan: SparkPlan):
Option[org.apache.spark.sql.catalyst.plans.logical.LogicalPlan] = {
+ plan.logicalLink match {
+ case some @ Some(_) => some
+ case None =>
+ plan.children match {
+ case Seq(child) => findLogicalLink(child)
+ case _ => None
+ }
+ }
+ }
val buildKeys = mode match {
case mode1: HashedRelationBroadcastMode =>
@@ -729,22 +746,29 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with
Logging {
if (VeloxConfig.get.enableBroadcastBuildOncePerExecutor) {
// Try to lookup from TreeNodeTag using child's logical plan
- // Need to recursively find logicalLink in case of AQE or other
wrappers
- @scala.annotation.tailrec
- def findLogicalLink(
- plan: SparkPlan):
Option[org.apache.spark.sql.catalyst.plans.logical.LogicalPlan] = {
- plan.logicalLink match {
- case some @ Some(_) => some
- case None =>
- plan.children match {
- case Seq(child) => findLogicalLink(child)
- case _ => None
- }
- }
- }
-
val newBuildKeys = findLogicalLink(child)
- .flatMap(_.getTagValue(JoinKeysTag.ORIGINAL_JOIN_KEYS))
+ .flatMap {
+ logicalPlan =>
+ logicalPlan
+ .getTagValue(BroadcastJoinContextTag.BROADCAST_JOIN_CONTEXT)
+ .flatMap {
+ contexts =>
+ val childOutputSet = child.outputSet
+ contexts.find {
+ ctx =>
+ val buildOutputSet = ctx.buildOutputSet
+ childOutputSet.subsetOf(buildOutputSet) &&
buildOutputSet.subsetOf(
+ childOutputSet)
+ }
+ }.map {
+ ctx =>
+ if (ctx.buildRight) {
+ ctx.originalRightKeys
+ } else {
+ ctx.originalLeftKeys
+ }
+ }
+ }
.getOrElse {
if (SparkHashJoinUtils.canRewriteAsLongType(buildKeys) &&
buildKeys.nonEmpty) {
SparkHashJoinUtils.getOriginalKeysFromPacked(buildKeys.head)
@@ -850,6 +874,7 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with
Logging {
.mapPartitions(itr => Iterator(BroadcastUtils.serializeStream(itr)))
.filter(_.numRows != 0)
.collect
+ val buildSideRowCount = serialized.map(_.numRows).sum
val rawSize = serialized.map(_.sizeInBytes()).sum
if (rawSize >= GlutenConfig.get.maxBroadcastTableSize) {
throw new SparkException(
@@ -857,7 +882,7 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with
Logging {
s"${SparkMemoryUtil.bytesToString(GlutenConfig.get.maxBroadcastTableSize)}: " +
s"${SparkMemoryUtil.bytesToString(rawSize)}")
}
- numOutputRows += serialized.map(_.numRows).sum
+ numOutputRows += buildSideRowCount
dataSize += rawSize
val rawThreads =
@@ -867,7 +892,8 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with
Logging {
val buildThreadsValue = if (rawThreads < 1) 1 else rawThreads
buildThreads += buildThreadsValue
- if (useOffheapBroadcastBuildRelation) {
+ // Create the base ColumnarBuildSideRelation first
+ val columnarRelation = if (useOffheapBroadcastBuildRelation) {
TaskResources.runUnsafe {
UnsafeColumnarBuildSideRelation(
newOutput,
@@ -886,6 +912,167 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with
Logging {
offload,
buildThreadsValue)
}
+
+ // Check if we should build hash table on driver (Spark-native approach)
+ // Only do this for HashedRelationBroadcastMode and when offload is enabled
+ val shouldBuildOnDriver =
VeloxConfig.get.enableDriverSideBroadcastHashTableBuild &&
+ mode.isInstanceOf[HashedRelationBroadcastMode] &&
+ offload
+
+ if (shouldBuildOnDriver) {
+ // Try to get broadcast join context from logical plan tag
+ // In multi-join scenarios, there may be multiple contexts. Find the one
that matches
+ // the current broadcast child's output.
+ val joinContextOpt: Option[BroadcastJoinContextInfo] =
+ findLogicalLink(child).flatMap {
+ logicalPlan =>
+ logicalPlan.getTagValue(
+ BroadcastJoinContextTag.BROADCAST_JOIN_CONTEXT
+ ).flatMap {
+ contexts =>
+ val childOutputSet = AttributeSet(newOutput)
+ // Find the context whose build output matches the child's
output
+ contexts.find {
+ ctx =>
+ val buildOutputMatches =
childOutputSet.subsetOf(ctx.buildOutputSet) &&
+ ctx.buildOutputSet.subsetOf(childOutputSet)
+ buildOutputMatches
+ }
+ }
+ }
+
+ joinContextOpt match {
+ case Some(joinContext) =>
+ // We have join context information - build hash table on driver
+ logInfo(
+ s"Building hash table on driver in BroadcastExchangeExec " +
+ s"with join context: $joinContext")
+
+ // Create a broadcast ID for this hash table
+ val broadcastId =
s"broadcast_exchange_${child.id}_${System.identityHashCode(mode)}"
+
+ // Convert Spark JoinType to Substrait JoinType
+ val substraitJoinType = joinContext.joinType match {
+ case _: InnerLike =>
+ JoinRel.JoinType.JOIN_TYPE_INNER
+ case FullOuter =>
+ JoinRel.JoinType.JOIN_TYPE_OUTER
+ case LeftOuter |
+ RightOuter =>
+ JoinRel.JoinType.JOIN_TYPE_LEFT
+ case LeftSemi |
+ ExistenceJoin(_) =>
+ JoinRel.JoinType.JOIN_TYPE_LEFT_SEMI
+ case LeftAnti =>
+ JoinRel.JoinType.JOIN_TYPE_LEFT_ANTI
+ case _ =>
+ JoinRel.JoinType.UNRECOGNIZED
+ }
+
+ // Extract filter information from join condition
+ val (filterBuildColumns, filterPropagatesNulls,
hasMixedFiltCondition) =
+ joinContext.condition match {
+ case Some(cond) =>
+ val buildAttrs = joinContext.buildOutputSet
+ val cols: Array[String] = cond.references.toSeq.collect {
+ case a: Attribute if buildAttrs.contains(a) =>
+ ConverterUtils.genColumnNameWithExprId(a)
+ }.toArray
+ val propagatesNulls =
SparkShimLoader.getSparkShims.isNullIntolerant(cond)
+ (cols, propagatesNulls, true)
+ case None =>
+ (Array.empty[String], false, false)
+ }
+
+ // Calculate bloom filter pushdown size if enabled
+ val bloomFilterPushdownSize = if
(VeloxConfig.get.hashProbeDynamicFilterPushdownEnabled) {
+ VeloxConfig.get.hashProbeBloomFilterPushdownMaxSize
+ } else {
+ -1
+ }
+
+ // Use the join keys from the matched context
+ // Since we already matched the context by comparing outputs,
+ // we know this is the correct one
+ val joinKeys = if (joinContext.buildRight) {
+ joinContext.originalRightKeys
+ } else {
+ joinContext.originalLeftKeys
+ }
+ val buildContext = BroadcastHashJoinContext(
+ buildSideJoinKeys = if (newBuildKeys.nonEmpty) newBuildKeys else
joinKeys,
+ substraitJoinType = substraitJoinType,
+ buildRight = joinContext.buildRight,
+ hasMixedFiltCondition = hasMixedFiltCondition,
+ isExistenceJoin = joinContext.joinType
+ .isInstanceOf[ExistenceJoin],
+ buildSideStructure = newOutput,
+ filterBuildColumns = filterBuildColumns,
+ filterPropagatesNulls = filterPropagatesNulls,
+ buildHashTableId = broadcastId,
+ isNullAwareAntiJoin = joinContext.isNullAwareAntiJoin,
+ bloomFilterPushdownSize = bloomFilterPushdownSize,
+ buildHashTableTimeMetric = Option(buildHashTableTimeMetric),
+ serializeHashTableTimeMetric =
Option(serializeHashTableTimeMetric),
+ serializedHashTableSizeMetric =
Option(serializedHashTableSizeMetric)
+ )
+
+ try {
+ // Build and serialize hash table on driver
+ val (serializedHashTable, safeMode) = columnarRelation match {
+ case rel: ColumnarBuildSideRelation =>
+ (
+ VeloxBroadcastBuildSideCache
+ .buildAndSerializeOnDriverInBroadcastExchange(
+ rel,
+ buildContext,
+ buildSideRowCount),
+ rel.safeBroadcastMode
+ )
+ case rel: UnsafeColumnarBuildSideRelation =>
+ (
+ VeloxBroadcastBuildSideCache
+ .buildAndSerializeOnDriverInBroadcastExchange(
+ rel,
+ buildContext,
+ buildSideRowCount),
+ rel.getSafeBroadcastMode
+ )
+ }
+
+ logInfo(
+ s"Successfully built hash table on driver: " +
+ s"size=${serializedHashTable.sizeInBytes} bytes, " +
+ s"rows=${serializedHashTable.numRows}, " +
+ s"joinType=${joinContext.joinType}, " +
+ s"broadcastId=$broadcastId")
+
+ // Return SerializedHashTableBroadcastRelation
+ SerializedHashTableBroadcastRelation(
+ serializedHashTable,
+ safeMode,
+ newOutput,
+ 0L, // buildTimeMs - tracked inside SerializedBroadcastHashTable
+ 0L // serializeTimeMs - tracked inside
SerializedBroadcastHashTable
+ )
+ } catch {
+ case e: Exception =>
+ logWarning(
+ s"Failed to build hash table on driver for
broadcastId=$broadcastId, " +
+ s"falling back to executor-side build: ${e.getMessage}",
+ e)
+ columnarRelation
+ }
+
+ case None =>
+ // No join context available - fall back to executor-side build
+ logInfo(s"No broadcast join context found in logical plan, using
executor-side build")
+ columnarRelation
+ }
+ } else {
+ // Return ColumnarBuildSideRelation for executor-side build (legacy
approach)
+ columnarRelation
+ }
}
override def doCanonicalizeForBroadcastMode(mode: BroadcastMode):
BroadcastMode = {
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala
b/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala
index dcb79a462f..dacb835b8c 100644
--- a/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala
+++ b/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala
@@ -113,6 +113,9 @@ class VeloxConfig(conf: SQLConf) extends GlutenConfig(conf)
{
getConf(VALUE_STREAM_DYNAMIC_FILTER_ENABLED)
def enableTimestampNtzValidation: Boolean =
getConf(ENABLE_TIMESTAMP_NTZ_VALIDATION)
+
+ def enableDriverSideBroadcastHashTableBuild: Boolean =
+ getConf(VELOX_DRIVER_SIDE_BROADCAST_HASH_TABLE_BUILD)
}
object VeloxConfig extends ConfigRegistry {
@@ -682,6 +685,15 @@ object VeloxConfig extends ConfigRegistry {
.booleanConf
.createWithDefault(true)
+ val VELOX_DRIVER_SIDE_BROADCAST_HASH_TABLE_BUILD =
+
buildConf("spark.gluten.sql.columnar.backend.velox.driverSideBroadcastHashTableBuild")
+ .doc(
+ "Enable driver-side broadcast hash table build. When enabled, the hash
table is " +
+ "built and serialized on the driver, then broadcast to executors.
When disabled, " +
+ "each executor builds its own hash table from the broadcast data.")
+ .booleanConf
+ .createWithDefault(false)
+
val QUERY_TRACE_ENABLED =
buildConf("spark.gluten.sql.columnar.backend.velox.queryTraceEnabled")
.doc("Enable query tracing flag.")
.booleanConf
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 06f59acf03..e77b2c6115 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
@@ -25,9 +25,10 @@ import org.apache.spark.rpc.GlutenDriverEndpoint
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,
SerializedHashTableBroadcastRelation, SparkPlan, SQLExecution}
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
@@ -177,9 +178,73 @@ case class BroadcastHashJoinExecTransformer(
buildBroadcastTableId,
isNullAwareAntiJoin,
bloomFilterPushdownSize,
- metrics.get("buildHashTableTime")
+ metrics.get("buildHashTableTime"),
+ metrics.get("serializeHashTableTime"),
+ metrics.get("deserializeHashTableTime"),
+ metrics.get("serializedHashTableSize")
)
- val broadcastRDD = VeloxBroadcastBuildSideRDD(sparkContext, broadcast,
context)
+
+ // Check the type of broadcast relation to determine the approach
+ val broadcastRDD = broadcast.value match {
+ case serializedRelation: SerializedHashTableBroadcastRelation =>
+ joinParamsForMetrics.foreach(_.usesDriverSideSerializedHashTable =
true)
+ // Hash table was already built and serialized in
BroadcastExchangeExec.
+ // Reuse the existing broadcast variable (no re-broadcast).
+ logInfo(
+ s"Using pre-built serialized hash table from BroadcastExchangeExec "
+
+ s"for $buildBroadcastTableId")
+
+ val rdd = VeloxSerializedBroadcastRDD(sparkContext, broadcast, context)
+
+ // Update bloom filter metrics
+ val (bloomFilterSize, dynamicFiltersProduced) =
rdd.getBloomFilterMetrics
+
metrics.get("bloomFilterBlocksByteSize").foreach(_.set(bloomFilterSize))
+
metrics.get("hashProbeDynamicFiltersProduced").foreach(_.set(dynamicFiltersProduced))
+
+ // Update size metric from the pre-built hash table
+ val (_, sizeInBytes, _, _) = serializedRelation.getMetrics
+ metrics.get("serializedHashTableSize").foreach(_.set(sizeInBytes))
+
+ rdd
+
+ case columnar: ColumnarBuildSideRelation =>
+ joinParamsForMetrics.foreach(_.usesDriverSideSerializedHashTable =
false)
+ // Legacy path: ColumnarBuildSideRelation from BroadcastExchangeExec.
+ // Hash table is built on each executor from broadcast data.
+ val canOffload = columnar.offload
+
+ if (!canOffload) {
+ logWarning(
+ s"Build side cannot be offloaded for $buildBroadcastTableId, " +
+ "falling back to executor-side build")
+ } else {
+ logInfo(s"Using executor-side broadcast hash table build for
$buildBroadcastTableId")
+ }
+ VeloxBroadcastBuildSideRDD(sparkContext, broadcast, context)
+
+ case unsafe: UnsafeColumnarBuildSideRelation =>
+ joinParamsForMetrics.foreach(_.usesDriverSideSerializedHashTable =
false)
+ // Similar to ColumnarBuildSideRelation
+ val canOffload = unsafe.isOffload
+
+ if (!canOffload) {
+ logWarning(
+ s"Build side cannot be offloaded for $buildBroadcastTableId, " +
+ "falling back to executor-side build")
+ } else {
+ logInfo(s"Using executor-side broadcast hash table build for
$buildBroadcastTableId")
+ }
+ VeloxBroadcastBuildSideRDD(sparkContext, broadcast, context)
+
+ case other =>
+ joinParamsForMetrics.foreach(_.usesDriverSideSerializedHashTable =
false)
+ // Fallback for unknown types
+ logWarning(
+ s"Unknown broadcast relation type: ${other.getClass.getName}, " +
+ "using executor-side build")
+ VeloxBroadcastBuildSideRDD(sparkContext, broadcast, context)
+ }
+
// FIXME: Do we have to make build side a RDD?
streamedRDD :+ broadcastRDD
}
@@ -197,7 +262,10 @@ case class BroadcastHashJoinContext(
buildHashTableId: String,
isNullAwareAntiJoin: Boolean = false,
bloomFilterPushdownSize: Long,
- buildHashTableTimeMetric: Option[SQLMetric] = None) {
+ buildHashTableTimeMetric: Option[SQLMetric] = None,
+ serializeHashTableTimeMetric: Option[SQLMetric] = None,
+ deserializeHashTableTimeMetric: Option[SQLMetric] = None,
+ serializedHashTableSizeMetric: 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/SerializedBroadcastHashTable.scala
b/backends-velox/src/main/scala/org/apache/gluten/execution/SerializedBroadcastHashTable.scala
new file mode 100644
index 0000000000..a59326daef
--- /dev/null
+++
b/backends-velox/src/main/scala/org/apache/gluten/execution/SerializedBroadcastHashTable.scala
@@ -0,0 +1,155 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.execution
+
+import org.apache.gluten.vectorized.HashJoinBuilder
+
+import org.apache.spark.sql.execution.joins.BuildSideRelation
+import org.apache.spark.sql.execution.unsafe.JniUnsafeByteBuffer
+import org.apache.spark.sql.execution.unsafe.UnsafeByteArray
+
+import java.io.{Externalizable, ObjectInput, ObjectOutput}
+
+/**
+ * Serialized broadcast hash table that can be efficiently broadcast to
executors. This is built on
+ * the driver and contains the serialized hash table data.
+ */
+class SerializedBroadcastHashTable(
+ var serializedData: UnsafeByteArray,
+ var numRows: Long,
+ var ignoreNullKeys: Boolean,
+ var joinHasNullKeys: Boolean,
+ var droppedDuplicates: Boolean,
+ var bloomFilterBlocksByteSize: Long,
+ var hashProbeDynamicFiltersProduced: Long,
+ var buildSideRelation: BuildSideRelation)
+ extends Externalizable {
+
+ def this() = this(null, 0, false, false, false, 0, 0, null) // Required for
Externalizable
+
+ override def writeExternal(out: ObjectOutput): Unit = {
+ out.writeLong(numRows)
+ out.writeBoolean(ignoreNullKeys)
+ out.writeBoolean(joinHasNullKeys)
+ out.writeBoolean(droppedDuplicates)
+ out.writeLong(bloomFilterBlocksByteSize)
+ out.writeLong(hashProbeDynamicFiltersProduced)
+ serializedData.writeExternal(out)
+ out.writeObject(buildSideRelation)
+ }
+
+ override def readExternal(in: ObjectInput): Unit = {
+ numRows = in.readLong()
+ ignoreNullKeys = in.readBoolean()
+ joinHasNullKeys = in.readBoolean()
+ droppedDuplicates = in.readBoolean()
+ bloomFilterBlocksByteSize = in.readLong()
+ hashProbeDynamicFiltersProduced = in.readLong()
+ val data = new UnsafeByteArray()
+ data.readExternal(in)
+ serializedData = data
+ buildSideRelation = in.readObject().asInstanceOf[BuildSideRelation]
+ }
+
+ /**
+ * Deserialize the hash table on executor side. The serialized Velox hash
table is already in a
+ * prepared, probe-ready form, so executor side only needs deserialization
without re-running
+ * prepareJoinTable.
+ *
+ * @return
+ * Hash table builder handle
+ */
+ def deserialize(): Long = {
+ HashJoinBuilder.deserializeHashTableDirect(
+ serializedData.address(),
+ Math.toIntExact(serializedData.size()),
+ ignoreNullKeys,
+ joinHasNullKeys)
+ }
+
+ /** Get the size of serialized data in bytes. */
+ def sizeInBytes: Long = serializedData.size()
+}
+
+object SerializedBroadcastHashTable {
+ def apply(
+ serializedData: UnsafeByteArray,
+ numRows: Long,
+ ignoreNullKeys: Boolean,
+ joinHasNullKeys: Boolean,
+ droppedDuplicates: Boolean,
+ bloomFilterBlocksByteSize: Long,
+ hashProbeDynamicFiltersProduced: Long,
+ buildSideRelation: BuildSideRelation): SerializedBroadcastHashTable =
+ new SerializedBroadcastHashTable(
+ serializedData,
+ numRows,
+ ignoreNullKeys,
+ joinHasNullKeys,
+ droppedDuplicates,
+ bloomFilterBlocksByteSize,
+ hashProbeDynamicFiltersProduced,
+ buildSideRelation)
+
+ /**
+ * Build and serialize a hash table on the driver.
+ *
+ * @param hashTableHandle
+ * Handle to the built hash table
+ * @param buildSideRelation
+ * The build side relation for metadata
+ * @return
+ * Serialized broadcast hash table
+ */
+ def fromHashTable(
+ hashTableHandle: Long,
+ buildSideRelation: BuildSideRelation,
+ droppedDuplicates: Boolean,
+ numRows: Long): SerializedBroadcastHashTable = {
+ try {
+ val serializedSize =
HashJoinBuilder.serializedHashTableSizeDirect(hashTableHandle)
+ val byteBuffer = JniUnsafeByteBuffer.allocate(serializedSize)
+ HashJoinBuilder.serializeHashTableDirect(
+ hashTableHandle,
+ byteBuffer.address(),
+ byteBuffer.size())
+ val serializedData = byteBuffer.toUnsafeByteArray()
+ val ignoreNullKeys = HashJoinBuilder
+ .getHashTableIgnoreNullKeys(hashTableHandle)
+ val joinHasNullKeys = HashJoinBuilder
+ .getHashTableJoinHasNullKeys(hashTableHandle)
+
+ val bloomFilterBlocksByteSize = HashJoinBuilder
+ .getHashTableBloomFilterBlocksByteSize(hashTableHandle)
+ val hashProbeDynamicFiltersProduced = if (bloomFilterBlocksByteSize > 0)
1L else 0L
+
+ SerializedBroadcastHashTable(
+ serializedData,
+ numRows,
+ ignoreNullKeys,
+ joinHasNullKeys,
+ droppedDuplicates,
+ bloomFilterBlocksByteSize,
+ hashProbeDynamicFiltersProduced,
+ buildSideRelation)
+ } finally {
+ synchronized {
+ HashJoinBuilder.clearHashTable(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 451c633be4..0d492817b1 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
@@ -16,7 +16,9 @@
*/
package org.apache.gluten.execution
+import org.apache.gluten.backendsapi.BackendsApiManager
import org.apache.gluten.backendsapi.velox.VeloxBackendSettings
+import org.apache.gluten.runtime.Runtimes
import org.apache.gluten.vectorized.HashJoinBuilder
import org.apache.spark.SparkEnv
@@ -25,6 +27,7 @@ import org.apache.spark.internal.Logging
import org.apache.spark.sql.execution.ColumnarBuildSideRelation
import org.apache.spark.sql.execution.joins.BuildSideRelation
import org.apache.spark.sql.execution.unsafe.UnsafeColumnarBuildSideRelation
+import org.apache.spark.task.TaskResources
import com.github.benmanes.caffeine.cache.{Cache, Caffeine, RemovalCause,
RemovalListener}
@@ -40,6 +43,11 @@ case class BroadcastHashTable(
*
* The complicated part is due to reuse exchange, where multiple BHJ IDs
correspond to a
* `BuildSideRelation`.
+ *
+ * This implementation supports two modes:
+ * 1. Driver-side build (new): Hash table is built and serialized on driver,
then broadcast to
+ * executors.
+ * 2. Executor-side build (legacy): Each executor builds its own hash table
from broadcast data
*/
object VeloxBroadcastBuildSideCache
extends Logging
@@ -58,6 +66,23 @@ object VeloxBroadcastBuildSideCache
.removalListener(this)
.build[String, BroadcastHashTable]()
+ // Cache for driver-side serialized hash tables to avoid rebuilding for
reuse exchange
+ private val driverSerializedCache: Cache[String,
SerializedBroadcastHashTable] =
+ Caffeine.newBuilder
+ .expireAfterAccess(expiredTime, TimeUnit.SECONDS)
+ .removalListener(
+ new RemovalListener[String, SerializedBroadcastHashTable] {
+ override def onRemoval(
+ key: String,
+ value: SerializedBroadcastHashTable,
+ cause: RemovalCause): Unit = {
+ if (value != null && value.serializedData != null) {
+ value.serializedData.release()
+ }
+ }
+ }
+ ).build[String, SerializedBroadcastHashTable]()
+
def getOrBuildBroadcastHashTable(
broadcast: Broadcast[BuildSideRelation],
broadcastContext: BroadcastHashJoinContext): BroadcastHashTable = {
@@ -78,6 +103,111 @@ object VeloxBroadcastBuildSideCache
)
}
+ /**
+ * Build hash table on driver and serialize for broadcasting. This version
is called from
+ * BroadcastExchangeExec and doesn't need a broadcast variable.
+ *
+ * This is the Spark-native approach where hash table is built in
BroadcastExchangeExec.
+ */
+ def buildAndSerializeOnDriverInBroadcastExchange(
+ relation: BuildSideRelation,
+ broadcastContext: BroadcastHashJoinContext,
+ numRows: Long): SerializedBroadcastHashTable = {
+
+ val broadcastId = broadcastContext.buildHashTableId
+
+ val cached = driverSerializedCache.getIfPresent(broadcastId)
+ if (cached != null) {
+ logInfo(s"Reusing cached serialized hash table for broadcast ID:
$broadcastId")
+ return cached
+ }
+
+ def resetRelation(droppedDuplicates: Boolean): Unit = relation match {
+ case r: ColumnarBuildSideRelation => r.reset(droppedDuplicates)
+ case r: UnsafeColumnarBuildSideRelation => r.reset(droppedDuplicates)
+ case _ =>
+ }
+
+ relation.synchronized {
+ val cachedAfterLock = driverSerializedCache.getIfPresent(broadcastId)
+ if (cachedAfterLock != null) {
+ logInfo(s"Reusing cached serialized hash table for broadcast ID:
$broadcastId (after lock)")
+ return cachedAfterLock
+ }
+
+ logInfo(
+ s"Building hash table on driver in BroadcastExchangeExec " +
+ s"for broadcast ID: $broadcastId")
+
+ val backendName = BackendsApiManager.getBackendName
+ TaskResources.runUnsafe {
+ val runtime = Runtimes.contextInstance(
+ backendName,
+ "DriverBroadcastHashTableBuild"
+ )
+
+ resetRelation(broadcastContext.droppedDuplicates)
+ val (hashTableHandle, _, droppedDuplicates) = relation match {
+ case r: ColumnarBuildSideRelation =>
+ r.buildHashTableWithRuntime(broadcastContext, runtime)
+ case r: UnsafeColumnarBuildSideRelation =>
+ r.buildHashTableWithRuntime(broadcastContext, runtime)
+ case other =>
+ throw new IllegalArgumentException(
+ s"Unsupported relation type for driver-side build:
${other.getClass.getName}")
+ }
+ try {
+ val startSerializeTime = System.currentTimeMillis()
+ val result =
+ SerializedBroadcastHashTable.fromHashTable(
+ hashTableHandle,
+ relation,
+ droppedDuplicates,
+ numRows)
+ val serializeTimeMs = System.currentTimeMillis() - startSerializeTime
+
+ logInfo(
+ s"Built and serialized hash table on driver: " +
+ s"size=${result.sizeInBytes} bytes, " +
+ s"rows=${result.numRows}, " +
+ s"serializeTime=${serializeTimeMs}ms " +
+ s"for broadcast ID: $broadcastId")
+
+ broadcastContext.serializeHashTableTimeMetric.foreach(_ +=
serializeTimeMs)
+ broadcastContext.serializedHashTableSizeMetric.foreach(_ +=
result.sizeInBytes)
+
+ driverSerializedCache.put(broadcastId, result)
+ result
+ } finally {
+ resetRelation(droppedDuplicates)
+ }
+ }
+ }
+ }
+
+ /** Deserialize hash table on executor from broadcast data. */
+ def deserializeOnExecutor(
+ serialized: SerializedBroadcastHashTable,
+ broadcastHashTableId: String,
+ deserializeHashTableTimeMetric:
Option[org.apache.spark.sql.execution.metric.SQLMetric] =
+ None): BroadcastHashTable = {
+
+ buildSideRelationCache.get(
+ broadcastHashTableId,
+ (_: String) => {
+ logInfo(s"Deserializing hash table on executor for broadcast ID:
$broadcastHashTableId")
+ val startTime = System.currentTimeMillis()
+ val hashTableHandle = serialized.deserialize()
+ val timeMs = System.currentTimeMillis() - startTime
+ deserializeHashTableTimeMetric.foreach(_ += timeMs)
+ BroadcastHashTable(
+ hashTableHandle,
+ serialized.buildSideRelation,
+ serialized.droppedDuplicates)
+ }
+ )
+ }
+
/** This is called from c++ side. */
def get(broadcastHashtableId: String): Long = {
Option(buildSideRelationCache.getIfPresent(broadcastHashtableId))
@@ -93,7 +223,13 @@ object VeloxBroadcastBuildSideCache
/** Only used in UT. */
def size(): Long = buildSideRelationCache.estimatedSize()
- def cleanAll(): Unit = buildSideRelationCache.invalidateAll()
+ /** Only used in UT. */
+ def driverSerializedCacheSize(): Long = driverSerializedCache.estimatedSize()
+
+ def cleanAll(): Unit = {
+ buildSideRelationCache.invalidateAll()
+ driverSerializedCache.invalidateAll()
+ }
override def onRemoval(
key: String,
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/execution/VeloxSerializedBroadcastRDD.scala
b/backends-velox/src/main/scala/org/apache/gluten/execution/VeloxSerializedBroadcastRDD.scala
new file mode 100644
index 0000000000..761d000097
--- /dev/null
+++
b/backends-velox/src/main/scala/org/apache/gluten/execution/VeloxSerializedBroadcastRDD.scala
@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.execution
+
+import org.apache.spark.{broadcast, SparkContext}
+import org.apache.spark.sql.execution.SerializedHashTableBroadcastRelation
+import org.apache.spark.sql.execution.joins.BuildSideRelation
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+/**
+ * RDD for handling serialized broadcast hash tables built on the driver. This
RDD deserializes the
+ * hash table on each executor.
+ */
+case class VeloxSerializedBroadcastRDD(
+ @transient private val sc: SparkContext,
+ broadcasted: broadcast.Broadcast[BuildSideRelation],
+ broadcastContext: BroadcastHashJoinContext)
+ extends BroadcastBuildSideRDD(sc, broadcasted) {
+
+ override def genBroadcastBuildSideIterator(): Iterator[ColumnarBatch] = {
+ val serialized = broadcasted.value match {
+ case relation: SerializedHashTableBroadcastRelation =>
relation.getSerializedHashTable
+ case other =>
+ throw new IllegalStateException(
+ s"VeloxSerializedBroadcastRDD expects
SerializedHashTableBroadcastRelation, " +
+ s"but got: ${other.getClass.getName}")
+ }
+ VeloxBroadcastBuildSideCache.deserializeOnExecutor(
+ serialized,
+ broadcastContext.buildHashTableId,
+ broadcastContext.deserializeHashTableTimeMetric
+ )
+
+ // Return empty iterator as hash table is already built
+ Iterator.empty
+ }
+
+ /**
+ * Get bloom filter metrics from the serialized hash table. This is called
from the driver to get
+ * metrics that were computed during hash table build.
+ */
+ def getBloomFilterMetrics: (Long, Long) = {
+ val serialized = broadcasted.value match {
+ case relation: SerializedHashTableBroadcastRelation =>
relation.getSerializedHashTable
+ case other =>
+ throw new IllegalStateException(
+ s"VeloxSerializedBroadcastRDD expects
SerializedHashTableBroadcastRelation, " +
+ s"but got: ${other.getClass.getName}")
+ }
+ (serialized.bloomFilterBlocksByteSize,
serialized.hashProbeDynamicFiltersProduced)
+ }
+}
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/metrics/JoinMetricsUpdater.scala
b/backends-velox/src/main/scala/org/apache/gluten/metrics/JoinMetricsUpdater.scala
index b056cd36a8..b7d32f94bd 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/metrics/JoinMetricsUpdater.scala
+++
b/backends-velox/src/main/scala/org/apache/gluten/metrics/JoinMetricsUpdater.scala
@@ -129,8 +129,14 @@ class HashJoinMetricsUpdater(override val metrics:
Map[String, SQLMetric])
hashProbeSpilledPartitions += hashProbeMetrics.spilledPartitions
hashProbeSpilledFiles += hashProbeMetrics.spilledFiles
hashProbeReplacedWithDynamicFilterRows +=
hashProbeMetrics.numReplacedWithDynamicFilterRows
- hashProbeDynamicFiltersProduced +=
hashProbeMetrics.numDynamicFiltersProduced
- bloomFilterBlocksByteSize += hashProbeMetrics.bloomFilterBlocksByteSize
+
+ // Only skip these metrics when this join actually reuses a pre-built
serialized
+ // hash table from driver-side build. Fallbacks still build on executors
and
+ // must accumulate native probe metrics here.
+ if (!joinParams.usesDriverSideSerializedHashTable) {
+ hashProbeDynamicFiltersProduced +=
hashProbeMetrics.numDynamicFiltersProduced
+ bloomFilterBlocksByteSize += hashProbeMetrics.bloomFilterBlocksByteSize
+ }
idx += 1
// HashBuild
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 285cbe7ef6..35587bde13 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
@@ -214,24 +214,117 @@ case class ColumnarBuildSideRelation(
val hashJoinBuilder = HashJoinBuilder.create(runtime)
- // Build the hash table
- hashTableData(droppedDuplicates) = hashJoinBuilder
- .nativeBuild(
- broadcastContext.buildHashTableId,
- batchArray.toArray,
- joinKeys,
- broadcastContext.filterBuildColumns,
- broadcastContext.filterPropagatesNulls,
- broadcastContext.substraitJoinType.ordinal(),
- broadcastContext.hasMixedFiltCondition,
- broadcastContext.isExistenceJoin,
- SubstraitUtil.toNameStruct(newOutput).toByteArray,
- broadcastContext.isNullAwareAntiJoin,
- broadcastContext.bloomFilterPushdownSize,
- buildThreads
+ try {
+ // Build the hash table
+ hashTableData(droppedDuplicates) = hashJoinBuilder
+ .nativeBuild(
+ broadcastContext.buildHashTableId,
+ batchArray.toArray,
+ joinKeys,
+ broadcastContext.filterBuildColumns,
+ broadcastContext.filterPropagatesNulls,
+ broadcastContext.substraitJoinType.ordinal(),
+ broadcastContext.hasMixedFiltCondition,
+ broadcastContext.isExistenceJoin,
+ SubstraitUtil.toNameStruct(newOutput).toByteArray,
+ broadcastContext.isNullAwareAntiJoin,
+ broadcastContext.bloomFilterPushdownSize,
+ buildThreads
+ )
+ } finally {
+ jniWrapper.close(serializeHandle)
+ }
+
+ // Update build hash table time metric
+ val elapsedTime = System.nanoTime() - startTime
+ broadcastContext.buildHashTableTimeMetric.foreach(_ += elapsedTime /
1000000)
+
+ (hashTableData(droppedDuplicates), this, droppedDuplicates)
+ } else {
+ (
+ HashJoinBuilder.cloneHashTable(hashTableData(droppedDuplicates)),
+ null,
+ droppedDuplicates)
+ }
+ }
+
+ /**
+ * Build hash table with provided runtime (for driver-side build). This
version doesn't rely on
+ * TaskContext and can be called from the driver.
+ */
+ def buildHashTableWithRuntime(
+ broadcastContext: BroadcastHashJoinContext,
+ runtime: org.apache.gluten.runtime.Runtime): (Long,
ColumnarBuildSideRelation, Boolean) =
+ synchronized {
+ val droppedDuplicates = broadcastContext.droppedDuplicates
+ if (!hashTableData.contains(droppedDuplicates)) {
+ val startTime = System.nanoTime()
+ val jniWrapper = ColumnarBatchSerializerJniWrapper.create(runtime)
+ val serializeHandle: Long = {
+ val allocator = ArrowBufferAllocators.globalInstance()
+ val cSchema = ArrowSchema.allocateNew(allocator)
+ val arrowSchema = SparkArrowUtil.toArrowSchema(
+ SparkShimLoader.getSparkShims.structFromAttributes(output),
+ SQLConf.get.sessionLocalTimeZone)
+ ArrowAbiUtil.exportSchema(allocator, arrowSchema, cSchema)
+ val handle = jniWrapper
+ .init(cSchema.memoryAddress())
+ cSchema.close()
+ handle
+ }
+
+ val batchArray = new ArrayBuffer[Long]
+
+ var batchId = 0
+ while (batchId < batches.size) {
+ batchArray.append(jniWrapper.deserialize(serializeHandle,
batches(batchId)))
+ batchId += 1
+ }
+
+ logDebug(
+ s"BHJ value size: " +
+ s"${broadcastContext.buildHashTableId} = ${batches.length}")
+
+ val (keys, newOutput) = if (newBuildKeys.isEmpty) {
+ (
+ broadcastContext.buildSideJoinKeys.asJava,
+ broadcastContext.buildSideStructure.asJava
)
+ } else {
+ (
+ newBuildKeys.asJava,
+ output.asJava
+ )
+ }
- jniWrapper.close(serializeHandle)
+ val joinKeys = keys.asScala.map {
+ key =>
+ val attr = ConverterUtils.getAttrFromExpr(key)
+ ConverterUtils.genColumnNameWithExprId(attr)
+ }.toArray
+
+ val hashJoinBuilder = HashJoinBuilder.create(runtime)
+
+ try {
+ // Build the hash table
+ hashTableData(droppedDuplicates) = hashJoinBuilder
+ .nativeBuild(
+ broadcastContext.buildHashTableId,
+ batchArray.toArray,
+ joinKeys,
+ broadcastContext.filterBuildColumns,
+ broadcastContext.filterPropagatesNulls,
+ broadcastContext.substraitJoinType.ordinal(),
+ broadcastContext.hasMixedFiltCondition,
+ broadcastContext.isExistenceJoin,
+ SubstraitUtil.toNameStruct(newOutput).toByteArray,
+ broadcastContext.isNullAwareAntiJoin,
+ broadcastContext.bloomFilterPushdownSize,
+ buildThreads
+ )
+ } finally {
+ jniWrapper.close(serializeHandle)
+ }
// Update build hash table time metric
val elapsedTime = System.nanoTime() - startTime
diff --git
a/backends-velox/src/main/scala/org/apache/spark/sql/execution/SerializedHashTableBroadcastRelation.scala
b/backends-velox/src/main/scala/org/apache/spark/sql/execution/SerializedHashTableBroadcastRelation.scala
new file mode 100644
index 0000000000..19949d4107
--- /dev/null
+++
b/backends-velox/src/main/scala/org/apache/spark/sql/execution/SerializedHashTableBroadcastRelation.scala
@@ -0,0 +1,137 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.execution
+
+import org.apache.gluten.execution.SerializedBroadcastHashTable
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression}
+import org.apache.spark.sql.catalyst.plans.physical.BroadcastMode
+import org.apache.spark.sql.execution.joins.BuildSideRelation
+import org.apache.spark.sql.vectorized.ColumnarBatch
+import org.apache.spark.util.KnownSizeEstimation
+
+/**
+ * Broadcast relation that contains a pre-built and serialized hash table.
This is similar to
+ * Spark's native HashedRelation broadcast approach where the hash table is
built once on the driver
+ * and broadcast to executors.
+ *
+ * Unlike ColumnarBuildSideRelation which broadcasts raw data and builds hash
table on each
+ * executor, this class broadcasts the serialized hash table directly, saving
CPU time on executors.
+ *
+ * @param serializedHashTable
+ * The serialized hash table built on driver
+ * @param mode
+ * The broadcast mode (HashedRelationBroadcastMode or IdentityBroadcastMode)
+ * @param output
+ * The output attributes
+ * @param buildTimeMs
+ * Time spent building hash table on driver (milliseconds)
+ * @param serializeTimeMs
+ * Time spent serializing hash table on driver (milliseconds)
+ */
+case class SerializedHashTableBroadcastRelation(
+ serializedHashTable: SerializedBroadcastHashTable,
+ safeBroadcastMode: SafeBroadcastMode,
+ output: Seq[Attribute],
+ buildTimeMs: Long,
+ serializeTimeMs: Long)
+ extends BuildSideRelation
+ with KnownSizeEstimation {
+
+ // Rebuild the real BroadcastMode on demand; never serialize it.
+ @transient override lazy val mode: BroadcastMode =
+ BroadcastModeUtils.fromSafe(safeBroadcastMode, output)
+
+ /**
+ * Returns an iterator of deserialized columnar batches. Note: This is not
the primary use case
+ * for this class. The main purpose is to provide the serialized hash table
directly to the join
+ * operator.
+ */
+ override def deserialized: Iterator[ColumnarBatch] = {
+ serializedHashTable.buildSideRelation match {
+ case _: SerializedHashTableBroadcastRelation =>
+ throw new IllegalStateException(
+ "Unexpected nested SerializedHashTableBroadcastRelation in
SerializedBroadcastHashTable")
+ case other =>
+ other.deserialized
+ }
+ }
+
+ override def asReadOnlyCopy(): SerializedHashTableBroadcastRelation = this
+
+ /**
+ * Get the serialized hash table for use in join operations. This is the
primary interface for
+ * consuming this broadcast relation.
+ */
+ def getSerializedHashTable: SerializedBroadcastHashTable =
serializedHashTable
+
+ /**
+ * Transform is used for DPP (Dynamic Partition Pruning) to extract keys. We
delegate to the
+ * underlying buildSideRelation in the serialized hash table.
+ */
+ override def transform(key: Expression): Array[InternalRow] = {
+ serializedHashTable.buildSideRelation.transform(key)
+ }
+
+ override def estimatedSize: Long = {
+ serializedHashTable.sizeInBytes
+ }
+
+ /**
+ * Get metrics for monitoring.
+ */
+ def getMetrics: (Long, Long, Long, Long) = {
+ (
+ serializedHashTable.numRows,
+ serializedHashTable.sizeInBytes,
+ buildTimeMs,
+ serializeTimeMs
+ )
+ }
+}
+
+object SerializedHashTableBroadcastRelation {
+
+ /**
+ * Create SerializedHashTableBroadcastRelation from
ColumnarBuildSideRelation by building and
+ * serializing the hash table on the driver.
+ */
+ def fromColumnarRelation(
+ columnarRelation: ColumnarBuildSideRelation,
+ broadcastContext: org.apache.gluten.execution.BroadcastHashJoinContext,
+ numRows: Long)
+ : SerializedHashTableBroadcastRelation = {
+
+ val startBuildTime = System.currentTimeMillis()
+
+ // Build and serialize hash table on driver
+ val serializedHashTable =
org.apache.gluten.execution.VeloxBroadcastBuildSideCache
+ .buildAndSerializeOnDriverInBroadcastExchange(columnarRelation,
broadcastContext, numRows)
+
+ val buildTimeMs = System.currentTimeMillis() - startBuildTime
+ val serializeTimeMs = 0L // This is tracked inside
SerializedBroadcastHashTable
+
+ SerializedHashTableBroadcastRelation(
+ serializedHashTable,
+ columnarRelation.safeBroadcastMode,
+ columnarRelation.output,
+ buildTimeMs,
+ serializeTimeMs
+ )
+ }
+}
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 57c53540ab..2a6963bb72 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
@@ -105,6 +105,8 @@ class UnsafeColumnarBuildSideRelation(
@transient override lazy val mode: BroadcastMode =
BroadcastModeUtils.fromSafe(safeBroadcastMode, output)
+ def getSafeBroadcastMode: SafeBroadcastMode = safeBroadcastMode
+
// If we stored expression bytes, deserialize once and cache locally (not
serialized).
@transient private lazy val exprKeysFromBytes: Option[Seq[Expression]] =
safeBroadcastMode match {
case HashExprSafeBroadcastMode(bytes, _) =>
@@ -184,24 +186,118 @@ class UnsafeColumnarBuildSideRelation(
val hashJoinBuilder = HashJoinBuilder.create(runtime)
- // Build the hash table
- hashTableData(droppedDuplicates) = hashJoinBuilder
- .nativeBuild(
- broadcastContext.buildHashTableId,
- batchArray.toArray,
- joinKeys,
- broadcastContext.filterBuildColumns,
- broadcastContext.filterPropagatesNulls,
- broadcastContext.substraitJoinType.ordinal(),
- broadcastContext.hasMixedFiltCondition,
- broadcastContext.isExistenceJoin,
- SubstraitUtil.toNameStruct(newOutput).toByteArray,
- broadcastContext.isNullAwareAntiJoin,
- broadcastContext.bloomFilterPushdownSize,
- buildThreads
+ try {
+ // Build the hash table
+ hashTableData(droppedDuplicates) = hashJoinBuilder
+ .nativeBuild(
+ broadcastContext.buildHashTableId,
+ batchArray.toArray,
+ joinKeys,
+ broadcastContext.filterBuildColumns,
+ broadcastContext.filterPropagatesNulls,
+ broadcastContext.substraitJoinType.ordinal(),
+ broadcastContext.hasMixedFiltCondition,
+ broadcastContext.isExistenceJoin,
+ SubstraitUtil.toNameStruct(newOutput).toByteArray,
+ broadcastContext.isNullAwareAntiJoin,
+ broadcastContext.bloomFilterPushdownSize,
+ buildThreads
+ )
+ } finally {
+ jniWrapper.close(serializeHandle)
+ }
+
+ // Update build hash table time metric
+ val elapsedTime = System.nanoTime() - startTime
+ broadcastContext.buildHashTableTimeMetric.foreach(_ += elapsedTime /
1000000)
+
+ (hashTableData(droppedDuplicates), this, droppedDuplicates)
+ } else {
+ (
+ HashJoinBuilder.cloneHashTable(hashTableData(droppedDuplicates)),
+ null,
+ droppedDuplicates)
+ }
+ }
+
+ /**
+ * Build hash table with provided runtime (for driver-side build). This
version doesn't rely on
+ * TaskContext and can be called from the driver.
+ */
+ def buildHashTableWithRuntime(
+ broadcastContext: BroadcastHashJoinContext,
+ runtime: org.apache.gluten.runtime.Runtime): (Long, BuildSideRelation,
Boolean) =
+ synchronized {
+ val droppedDuplicates = broadcastContext.droppedDuplicates
+ if (!hashTableData.contains(droppedDuplicates)) {
+ val startTime = System.nanoTime()
+ val jniWrapper = ColumnarBatchSerializerJniWrapper.create(runtime)
+ val serializeHandle: Long = {
+ val allocator = ArrowBufferAllocators.globalInstance()
+ val cSchema = ArrowSchema.allocateNew(allocator)
+ val arrowSchema = SparkArrowUtil.toArrowSchema(
+ SparkShimLoader.getSparkShims.structFromAttributes(output),
+ SQLConf.get.sessionLocalTimeZone)
+ ArrowAbiUtil.exportSchema(allocator, arrowSchema, cSchema)
+ val handle = jniWrapper
+ .init(cSchema.memoryAddress())
+ cSchema.close()
+ handle
+ }
+
+ val batchArray = new ArrayBuffer[Long]
+
+ var batchId = 0
+ while (batchId < batches.size) {
+ val (offset, length) = (batches(batchId).address(),
batches(batchId).size())
+ batchArray.append(jniWrapper.deserializeDirect(serializeHandle,
offset, length.toInt))
+ batchId += 1
+ }
+
+ logDebug(
+ s"BHJ value size: " +
+ s"${broadcastContext.buildHashTableId} = ${batches.size}")
+
+ val (keys, newOutput) = if (newBuildKeys.isEmpty) {
+ (
+ broadcastContext.buildSideJoinKeys.asJava,
+ broadcastContext.buildSideStructure.asJava
+ )
+ } else {
+ (
+ newBuildKeys.asJava,
+ output.asJava
)
+ }
- jniWrapper.close(serializeHandle)
+ val joinKeys = keys.asScala.map {
+ key =>
+ val attr = ConverterUtils.getAttrFromExpr(key)
+ ConverterUtils.genColumnNameWithExprId(attr)
+ }.toArray
+
+ val hashJoinBuilder = HashJoinBuilder.create(runtime)
+
+ try {
+ // Build the hash table
+ hashTableData(droppedDuplicates) = hashJoinBuilder
+ .nativeBuild(
+ broadcastContext.buildHashTableId,
+ batchArray.toArray,
+ joinKeys,
+ broadcastContext.filterBuildColumns,
+ broadcastContext.filterPropagatesNulls,
+ broadcastContext.substraitJoinType.ordinal(),
+ broadcastContext.hasMixedFiltCondition,
+ broadcastContext.isExistenceJoin,
+ SubstraitUtil.toNameStruct(newOutput).toByteArray,
+ broadcastContext.isNullAwareAntiJoin,
+ broadcastContext.bloomFilterPushdownSize,
+ buildThreads
+ )
+ } finally {
+ jniWrapper.close(serializeHandle)
+ }
// Update build hash table time metric
val elapsedTime = System.nanoTime() - startTime
diff --git
a/backends-velox/src/test/scala/org/apache/gluten/execution/DynamicOffHeapSizingSuite.scala
b/backends-velox/src/test/scala/org/apache/gluten/execution/DynamicOffHeapSizingSuite.scala
index ddd76f917d..523aae65c5 100644
---
a/backends-velox/src/test/scala/org/apache/gluten/execution/DynamicOffHeapSizingSuite.scala
+++
b/backends-velox/src/test/scala/org/apache/gluten/execution/DynamicOffHeapSizingSuite.scala
@@ -35,6 +35,7 @@ class DynamicOffHeapSizingSuite extends
VeloxWholeStageTransformerSuite {
.set("spark.shuffle.manager",
"org.apache.spark.shuffle.sort.ColumnarShuffleManager")
.set("spark.executor.memory", "2GB")
.set("spark.memory.offHeap.enabled", "false")
+
.set("spark.gluten.sql.columnar.backend.velox.driverSideBroadcastHashTableBuild",
"false")
.set(
"spark.gluten.velox.buildHashTableOncePerExecutor.enabled",
"false"
diff --git
a/backends-velox/src/test/scala/org/apache/gluten/execution/VeloxHashJoinSuite.scala
b/backends-velox/src/test/scala/org/apache/gluten/execution/VeloxHashJoinSuite.scala
index 518aad3427..c5c8d23123 100644
---
a/backends-velox/src/test/scala/org/apache/gluten/execution/VeloxHashJoinSuite.scala
+++
b/backends-velox/src/test/scala/org/apache/gluten/execution/VeloxHashJoinSuite.scala
@@ -22,8 +22,9 @@ import org.apache.gluten.sql.shims.SparkShimLoader
import org.apache.spark.SparkConf
import org.apache.spark.sql.Row
import org.apache.spark.sql.catalyst.expressions.AttributeReference
-import org.apache.spark.sql.execution.{ColumnarBroadcastExchangeExec,
ColumnarSubqueryBroadcastExec, InputIteratorTransformer}
+import org.apache.spark.sql.execution.{ColumnarBroadcastExchangeExec,
ColumnarSubqueryBroadcastExec, InputIteratorTransformer,
SerializedHashTableBroadcastRelation}
import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
+import org.apache.spark.sql.execution.joins.BuildSideRelation
import org.apache.spark.sql.execution.joins.HashedRelationBroadcastMode
class VeloxHashJoinSuite extends VeloxWholeStageTransformerSuite {
@@ -39,6 +40,27 @@ class VeloxHashJoinSuite extends
VeloxWholeStageTransformerSuite {
override protected def sparkConf: SparkConf = super.sparkConf
.set("spark.unsafe.exceptionOnMemoryLeak", "true")
+ override protected def beforeEach(): Unit = {
+ super.beforeEach()
+ VeloxBroadcastBuildSideCache.cleanAll()
+ }
+
+ override protected def afterEach(): Unit = {
+ try {
+ VeloxBroadcastBuildSideCache.cleanAll()
+ } finally {
+ super.afterEach()
+ }
+ }
+
+ private def collectBroadcastRelations(df: org.apache.spark.sql.DataFrame)
+ : Seq[BuildSideRelation] = {
+ collectWithSubqueries(df.queryExecution.executedPlan) {
+ case exchange: ColumnarBroadcastExchangeExec =>
+ exchange.executeBroadcast[BuildSideRelation]().value
+ }
+ }
+
test("generate hash join plan - v1") {
withSQLConf(
("spark.sql.autoBroadcastJoinThreshold", "-1"),
@@ -313,6 +335,125 @@ class VeloxHashJoinSuite extends
VeloxWholeStageTransformerSuite {
}
}
+ test(
+ "driver-side broadcast hash table build uses serialized relation and
preserves join results") {
+ withSQLConf(
+ ("spark.sql.autoBroadcastJoinThreshold", "10MB"),
+ ("spark.sql.adaptive.enabled", "false"),
+ (VeloxConfig.VELOX_DRIVER_SIDE_BROADCAST_HASH_TABLE_BUILD.key, "true"),
+ (VeloxConfig.VELOX_BROADCAST_BUILD_RELATION_USE_OFFHEAP.key, "true")
+ ) {
+ withTable("driver_build_fact", "driver_build_dim") {
+ spark.range(
+ 0,
+ 200).selectExpr("id as k", "id % 11 as
v").write.saveAsTable("driver_build_fact")
+ spark.range(0, 50).selectExpr("id as k", "concat('dim_', cast(id as
string)) as name").write
+ .saveAsTable("driver_build_dim")
+
+ val query =
+ """
+ |SELECT /*+ BROADCAST(driver_build_dim) */
+ | f.k, f.v, d.name
+ |FROM driver_build_fact f
+ |JOIN driver_build_dim d
+ |ON f.k = d.k
+ |ORDER BY f.k
+ |""".stripMargin
+
+ runQueryAndCompare(query) {
+ df =>
+ val plan = df.queryExecution.executedPlan
+ val broadcastJoins = plan.collect { case bhj:
BroadcastHashJoinExecTransformer => bhj }
+ assert(broadcastJoins.nonEmpty, "Should use broadcast hash join")
+
+ val relations = collectBroadcastRelations(df)
+ assert(
+ relations.size == 1,
+ s"Expected a single broadcast relation, got ${relations.size}")
+ assert(
+
relations.head.isInstanceOf[SerializedHashTableBroadcastRelation],
+ s"Expected SerializedHashTableBroadcastRelation," +
+ s" got ${relations.head.getClass.getName}"
+ )
+
+ val serializedRelation =
+ relations.head.asInstanceOf[SerializedHashTableBroadcastRelation]
+ assert(serializedRelation.getSerializedHashTable.sizeInBytes > 0)
+ assert(
+ VeloxBroadcastBuildSideCache.driverSerializedCacheSize() >= 1,
+ s"Expected driver serialized cache to contain entries, got " +
+ s"${VeloxBroadcastBuildSideCache.driverSerializedCacheSize()}"
+ )
+ }
+
+ VeloxBroadcastBuildSideCache.cleanAll()
+ assert(VeloxBroadcastBuildSideCache.driverSerializedCacheSize() == 0)
+ }
+ }
+ }
+
+ test("driver-side broadcast hash table build reuses exchange and cache
cleanup works") {
+ withSQLConf(
+ ("spark.sql.autoBroadcastJoinThreshold", "10MB"),
+ ("spark.sql.adaptive.enabled", "false"),
+ ("spark.sql.exchange.reuse", "true"),
+ (VeloxConfig.VELOX_DRIVER_SIDE_BROADCAST_HASH_TABLE_BUILD.key, "true"),
+ (VeloxConfig.VELOX_BROADCAST_BUILD_RELATION_USE_OFFHEAP.key, "true")
+ ) {
+ withTable("reuse_fact", "reuse_dim") {
+ spark.range(0, 400).selectExpr("id as id", "id % 40 as
k").write.saveAsTable("reuse_fact")
+ spark.range(0, 40).selectExpr("id as k", "id * 3 as
v").write.saveAsTable("reuse_dim")
+
+ val query =
+ """
+ |WITH shared_dim AS (
+ | SELECT * FROM reuse_dim WHERE k < 20
+ |)
+ |SELECT /*+ BROADCAST(d1), BROADCAST(d2) */ count(*)
+ |FROM reuse_fact f
+ |JOIN shared_dim d1
+ | ON f.k = d1.k
+ |JOIN shared_dim d2
+ | ON f.k = d2.k
+ |""".stripMargin
+
+ runQueryAndCompare(query) {
+ df =>
+ val plan = df.queryExecution.executedPlan
+ val broadcastJoins = plan.collect { case bhj:
BroadcastHashJoinExecTransformer => bhj }
+ assert(
+ broadcastJoins.size == 2,
+ s"Expected two broadcast hash joins, got ${broadcastJoins.size}")
+
+ val reusedExchanges = collectWithSubqueries(plan) {
+ case reused: ReusedExchangeExec => reused
+ }
+ assert(
+ reusedExchanges.nonEmpty,
+ "Expected reused broadcast exchange for shared build side")
+
+ val relations = collectBroadcastRelations(df)
+ assert(
+ relations.size == 1,
+ s"Expected one materialized broadcast relation, got
${relations.size}")
+ assert(
+
relations.head.isInstanceOf[SerializedHashTableBroadcastRelation],
+ s"Expected SerializedHashTableBroadcastRelation," +
+ s" got ${relations.head.getClass.getName}"
+ )
+ assert(
+ VeloxBroadcastBuildSideCache.driverSerializedCacheSize() >= 1,
+ s"Expected driver serialized cache to contain entries, got " +
+ s"${VeloxBroadcastBuildSideCache.driverSerializedCacheSize()}"
+ )
+ }
+
+ VeloxBroadcastBuildSideCache.cleanAll()
+ assert(VeloxBroadcastBuildSideCache.driverSerializedCacheSize() == 0)
+ }
+ }
+ }
+
test("Broadcast join with multiple cast expressions in join keys") {
withSQLConf(
("spark.sql.autoBroadcastJoinThreshold", "10MB"),
@@ -513,4 +654,5 @@ class VeloxHashJoinSuite extends
VeloxWholeStageTransformerSuite {
}
})
}
+
}
diff --git
a/backends-velox/src/test/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelationTest.scala
b/backends-velox/src/test/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelationTest.scala
index c881d77ed1..89471119e3 100644
---
a/backends-velox/src/test/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelationTest.scala
+++
b/backends-velox/src/test/scala/org/apache/spark/sql/execution/unsafe/UnsafeColumnarBuildSideRelationTest.scala
@@ -16,6 +16,7 @@
*/
package org.apache.spark.sql.execution.unsafe
+import org.apache.gluten.execution.VeloxBroadcastBuildSideCache
import org.apache.gluten.memory.arrow.alloc.ArrowBufferAllocators
import
org.apache.gluten.memory.memtarget.ThrowOnOomMemoryTarget.OutOfMemoryException
@@ -77,6 +78,8 @@ class UnsafeColumnarBuildSideRelationTest extends
SharedSparkSession {
// be released after a full-GC.
unsafeRelWithIdentityMode = null
unsafeRelWithHashMode = null
+ // Clean up the broadcast build side cache to release any cached
serialized hash tables
+ VeloxBroadcastBuildSideCache.cleanAll()
System.gc()
Thread.sleep(1000)
// Since we trigger GC in beforeAll() to clean up residual memory from
previous test suites,
diff --git a/cpp/velox/CMakeLists.txt b/cpp/velox/CMakeLists.txt
index bbb20a8461..d935a12d04 100644
--- a/cpp/velox/CMakeLists.txt
+++ b/cpp/velox/CMakeLists.txt
@@ -180,6 +180,7 @@ set(VELOX_SRCS
operators/functions/SparkExprToSubfieldFilterParser.cc
operators/plannodes/RowVectorStream.cc
operators/hashjoin/HashTableBuilder.cc
+ operators/hashjoin/HashTableSerializer.cc
operators/reader/FileReaderIterator.cc
operators/reader/ParquetReaderIterator.cc
operators/serializer/VeloxColumnarBatchSerializer.cc
diff --git a/cpp/velox/jni/JniHashTable.cc b/cpp/velox/jni/JniHashTable.cc
index 1187347157..d7259879c5 100644
--- a/cpp/velox/jni/JniHashTable.cc
+++ b/cpp/velox/jni/JniHashTable.cc
@@ -23,6 +23,7 @@
#include "folly/String.h"
#include "memory/ColumnarBatch.h"
#include "memory/VeloxColumnarBatch.h"
+#include "operators/hashjoin/HashTableSerializer.h"
#include "substrait/algebra.pb.h"
#include "substrait/type.pb.h"
#include "velox/core/PlanNode.h"
@@ -163,4 +164,86 @@ long getJoin(const std::string& hashTableId) {
return JniHashTableContext::getInstance().callJavaGet(hashTableId);
}
+size_t serializedHashTableSize(std::shared_ptr<HashTableBuilder> builder) {
+ VELOX_CHECK_NOT_NULL(builder, "Hash table builder cannot be null");
+
+ auto hashTable = builder->hashTable();
+ VELOX_CHECK_NOT_NULL(hashTable, "Hash table cannot be null");
+
+ auto* hashTableFalse =
dynamic_cast<facebook::velox::exec::HashTable<false>*>(hashTable.get());
+ if (hashTableFalse != nullptr) {
+ return HashTableSerializer::serializedSize<false>(hashTableFalse);
+ }
+
+ auto* hashTableTrue =
dynamic_cast<facebook::velox::exec::HashTable<true>*>(hashTable.get());
+ VELOX_CHECK_NOT_NULL(hashTableTrue, "Hash table must be either
HashTable<false> or HashTable<true>");
+ return HashTableSerializer::serializedSize<true>(hashTableTrue);
+}
+
+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");
+
+ auto hashTable = builder->hashTable();
+ VELOX_CHECK_NOT_NULL(hashTable, "Hash table cannot be null");
+
+ auto* hashTableFalse =
dynamic_cast<facebook::velox::exec::HashTable<false>*>(hashTable.get());
+ if (hashTableFalse != nullptr) {
+ HashTableSerializer::serializeTo<false>(hashTableFalse, data, size);
+ return;
+ }
+
+ auto* hashTableTrue =
dynamic_cast<facebook::velox::exec::HashTable<true>*>(hashTable.get());
+ VELOX_CHECK_NOT_NULL(hashTableTrue, "Hash table must be either
HashTable<false> or HashTable<true>");
+ HashTableSerializer::serializeTo<true>(hashTableTrue, data, size);
+}
+
+std::shared_ptr<HashTableBuilder>
+deserializeHashTable(const uint8_t* data, size_t size, bool ignoreNullKeys,
bool joinHasNullKeys) {
+ VELOX_CHECK_NOT_NULL(data, "Serialized data cannot be null");
+ VELOX_CHECK_GT(size, 0, "Invalid data size");
+
+ auto pool = defaultLeafVeloxMemoryPool();
+ auto* poolPtr = pool.get();
+
+ std::unique_ptr<facebook::velox::exec::BaseHashTable> hashTable;
+ if (ignoreNullKeys) {
+ auto derived = HashTableSerializer::deserialize<true>(data, size, poolPtr);
+ hashTable = std::move(derived);
+ } else {
+ auto derived = HashTableSerializer::deserialize<false>(data, size,
poolPtr);
+ hashTable = std::move(derived);
+ }
+
+ std::vector<std::shared_ptr<const
facebook::velox::core::FieldAccessTypedExpr>> emptyKeys;
+ std::vector<uint32_t> emptyChannels;
+
+ auto keyTypes = hashTable->rows()->keyTypes();
+ std::vector<std::string> names;
+ for (size_t i = 0; i < keyTypes.size(); ++i) {
+ names.push_back("key" + std::to_string(i));
+ }
+ auto rowType = facebook::velox::ROW(std::move(names), std::move(keyTypes));
+
+ auto builder = std::make_shared<HashTableBuilder>(
+ facebook::velox::core::JoinType::kInner,
+ false,
+ false,
+ -1,
+ emptyKeys,
+ emptyChannels,
+ false,
+ rowType,
+ poolPtr,
+ 1000,
+ 1000000,
+ 100000,
+ 0);
+
+ builder->setHashTable(std::move(hashTable));
+ // Restore the joinHasNullKeys flag
+ builder->setJoinHasNullKeys(joinHasNullKeys);
+ return builder;
+}
+
} // namespace gluten
diff --git a/cpp/velox/jni/JniHashTable.h b/cpp/velox/jni/JniHashTable.h
index 47f89d1799..d75d2d663b 100644
--- a/cpp/velox/jni/JniHashTable.h
+++ b/cpp/velox/jni/JniHashTable.h
@@ -91,6 +91,16 @@ std::shared_ptr<HashTableBuilder> nativeHashTableBuild(
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);
+
+// Serialize hash table directly to a caller-provided buffer.
+void serializeHashTableTo(std::shared_ptr<HashTableBuilder> builder, uint8_t*
data, size_t size);
+
+// Deserialize hash table from broadcast data with explicit ignoreNullKeys
parameter
+std::shared_ptr<HashTableBuilder>
+deserializeHashTable(const uint8_t* data, size_t size, bool ignoreNullKeys,
bool joinHasNullKeys = false);
+
// Initialize the JNI hash table context
inline void initVeloxJniHashTable(JNIEnv* env, JavaVM* javaVm) {
JniHashTableContext::getInstance().initialize(env, javaVm);
diff --git a/cpp/velox/jni/VeloxJniWrapper.cc b/cpp/velox/jni/VeloxJniWrapper.cc
index 369d6716ed..36798f5cde 100644
--- a/cpp/velox/jni/VeloxJniWrapper.cc
+++ b/cpp/velox/jni/VeloxJniWrapper.cc
@@ -125,7 +125,6 @@ void JNI_OnUnload(JavaVM* vm, void*) {
env->DeleteGlobalRef(blockStripesClass);
env->DeleteGlobalRef(infoCls);
-
finalizeVeloxJniUDF(env);
finalizeVeloxJniFileSystem(env);
finalizeVeloxJniHashTable(env);
@@ -1021,6 +1020,11 @@ JNIEXPORT jlong JNICALL
Java_org_apache_gluten_vectorized_HashJoinBuilder_native
}
if (numThreads == 1) {
+ // Use default global pool for driver-side build
+ // The hash table will be serialized and broadcast, so it doesn't need
runtime's pool
+ // Using runtime pool causes lifecycle management issues
+ auto memoryPool = defaultLeafVeloxMemoryPool();
+
auto builder = nativeHashTableBuild(
hashJoinKeys,
filterColumns,
@@ -1037,7 +1041,7 @@ JNIEXPORT jlong JNICALL
Java_org_apache_gluten_vectorized_HashJoinBuilder_native
abandonHashBuildDedupMinRows,
abandonHashBuildDedupMinPct,
cb,
- defaultLeafVeloxMemoryPool());
+ memoryPool);
auto mainTable = builder->uniqueTable();
mainTable->prepareJoinTable(
@@ -1072,6 +1076,10 @@ JNIEXPORT jlong JNICALL
Java_org_apache_gluten_vectorized_HashJoinBuilder_native
threadBatches.push_back(cb[i]);
}
+ // Use default global pool for driver-side build
+ // The hash table will be serialized and broadcast, so it doesn't need
runtime's pool
+ auto threadMemoryPool = defaultLeafVeloxMemoryPool();
+
auto builder = nativeHashTableBuild(
hashJoinKeys,
filterColumns,
@@ -1088,7 +1096,7 @@ JNIEXPORT jlong JNICALL
Java_org_apache_gluten_vectorized_HashJoinBuilder_native
abandonHashBuildDedupMinRows,
abandonHashBuildDedupMinPct,
threadBatches,
- defaultLeafVeloxMemoryPool());
+ threadMemoryPool);
hashTableBuilders[t] = std::move(builder);
otherTables[t] = std::move(hashTableBuilders[t]->uniqueTable());
@@ -1150,6 +1158,100 @@ JNIEXPORT void JNICALL
Java_org_apache_gluten_vectorized_HashJoinBuilder_clearHa
ObjectStore::release(tableHandler);
JNI_METHOD_END()
}
+
+JNIEXPORT jlong JNICALL
Java_org_apache_gluten_vectorized_HashJoinBuilder_deserializeHashTableDirect(
// NOLINT
+ JNIEnv* env,
+ jclass,
+ jlong address,
+ jint size,
+ jboolean ignoreNullKeys,
+ jboolean joinHasNullKeys) {
+ JNI_METHOD_START
+ auto builder = gluten::deserializeHashTable(
+ reinterpret_cast<const uint8_t*>(address),
+ static_cast<size_t>(size),
+ static_cast<bool>(ignoreNullKeys),
+ static_cast<bool>(joinHasNullKeys));
+ return gluten::getHashTableObjStore()->save(builder);
+ JNI_METHOD_END(kInvalidObjectHandle)
+}
+
+JNIEXPORT jboolean JNICALL
Java_org_apache_gluten_vectorized_HashJoinBuilder_getHashTableIgnoreNullKeys(
// NOLINT
+ JNIEnv* env,
+ jclass,
+ jlong hashTableHandle) {
+ JNI_METHOD_START
+ auto builder =
ObjectStore::retrieve<gluten::HashTableBuilder>(hashTableHandle);
+ auto hashTable = builder->hashTable();
+ VELOX_CHECK_NOT_NULL(hashTable, "Hash table cannot be null");
+ return
static_cast<jboolean>(dynamic_cast<facebook::velox::exec::HashTable<true>*>(hashTable.get())
!= nullptr);
+ JNI_METHOD_END(false)
+}
+
+JNIEXPORT jboolean JNICALL
Java_org_apache_gluten_vectorized_HashJoinBuilder_getHashTableJoinHasNullKeys(
// NOLINT
+ JNIEnv* env,
+ jclass,
+ jlong hashTableHandle) {
+ JNI_METHOD_START
+ auto builder =
ObjectStore::retrieve<gluten::HashTableBuilder>(hashTableHandle);
+ return static_cast<jboolean>(builder->joinHasNullKeys());
+ JNI_METHOD_END(false)
+}
+
+JNIEXPORT jlong JNICALL
+Java_org_apache_gluten_vectorized_HashJoinBuilder_getHashTableBloomFilterBlocksByteSize(
// NOLINT
+ JNIEnv* env,
+ jclass,
+ jlong hashTableHandle) {
+ JNI_METHOD_START
+ auto builder =
ObjectStore::retrieve<gluten::HashTableBuilder>(hashTableHandle);
+ auto hashTable = builder->hashTable();
+ VELOX_CHECK_NOT_NULL(hashTable, "Hash table cannot be null");
+
+ auto* baseTable =
dynamic_cast<facebook::velox::exec::BaseHashTable*>(hashTable.get());
+ VELOX_CHECK_NOT_NULL(baseTable, "Hash table must derive from BaseHashTable");
+
+ int64_t bloomFilterBlocksByteSize = 0;
+ for (const auto& hasher : baseTable->hashers()) {
+ const auto& bloomFilter = hasher->getBloomFilter();
+ if (bloomFilter == nullptr) {
+ continue;
+ }
+ auto* bfFilter =
dynamic_cast<facebook::velox::common::BigintValuesUsingBloomFilter*>(bloomFilter.get());
+ if (bfFilter != nullptr) {
+ bloomFilterBlocksByteSize += bfFilter->blocksByteSize();
+ }
+ }
+ return static_cast<jlong>(bloomFilterBlocksByteSize);
+ JNI_METHOD_END(0L)
+}
+
+JNIEXPORT jlong JNICALL
Java_org_apache_gluten_vectorized_HashJoinBuilder_serializedHashTableSizeDirect(
// NOLINT
+ JNIEnv* env,
+ jclass,
+ jlong hashTableHandle) {
+ JNI_METHOD_START
+ auto builder =
ObjectStore::retrieve<gluten::HashTableBuilder>(hashTableHandle);
+ return static_cast<jlong>(gluten::serializedHashTableSize(builder));
+ JNI_METHOD_END(0L)
+}
+
+JNIEXPORT void JNICALL
Java_org_apache_gluten_vectorized_HashJoinBuilder_serializeHashTableDirect( //
NOLINT
+ JNIEnv* env,
+ jclass,
+ jlong hashTableHandle,
+ jlong address,
+ jlong size) {
+ JNI_METHOD_START
+ auto builder =
ObjectStore::retrieve<gluten::HashTableBuilder>(hashTableHandle);
+ VELOX_CHECK_GT(address, 0, "Serialized hash table buffer address must be
positive");
+ VELOX_CHECK_GE(size, 0, "Serialized hash table buffer size must be
non-negative");
+ const auto serializedSize = gluten::serializedHashTableSize(builder);
+ VELOX_CHECK_EQ(static_cast<size_t>(size), serializedSize, "Hash table buffer
size mismatch");
+ gluten::serializeHashTableTo(builder, reinterpret_cast<uint8_t*>(address),
serializedSize);
+ JNI_METHOD_END()
+}
+
#ifdef __cplusplus
}
#endif
diff --git a/cpp/velox/operators/hashjoin/HashTableBuilder.cc
b/cpp/velox/operators/hashjoin/HashTableBuilder.cc
index 363edea4ec..4c3129fccf 100644
--- a/cpp/velox/operators/hashjoin/HashTableBuilder.cc
+++ b/cpp/velox/operators/hashjoin/HashTableBuilder.cc
@@ -181,7 +181,10 @@ bool HashTableBuilder::abandonHashBuildDedupEarly(int64_t
numDistinct) const {
void HashTableBuilder::abandonHashBuildDedup() {
abandonHashBuildDedup_ = true;
- uniqueTable_->setAllowDuplicates(true);
+
+ if (!dropDuplicates_) {
+ uniqueTable_->setAllowDuplicates(true);
+ }
lookup_.reset();
}
diff --git a/cpp/velox/operators/hashjoin/HashTableSerializer.cc
b/cpp/velox/operators/hashjoin/HashTableSerializer.cc
new file mode 100644
index 0000000000..76365d76a4
--- /dev/null
+++ b/cpp/velox/operators/hashjoin/HashTableSerializer.cc
@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "operators/hashjoin/HashTableSerializer.h"
+#include "velox/common/base/Exceptions.h"
+
+namespace gluten {
+
+template <bool ignoreNullKeys>
+size_t HashTableSerializer::serializedSize(const
facebook::velox::exec::HashTable<ignoreNullKeys>* hashTable) {
+ VELOX_CHECK_NOT_NULL(hashTable, "Hash table cannot be null");
+ return hashTable->serializedSize();
+}
+
+template <bool ignoreNullKeys>
+void HashTableSerializer::serializeTo(
+ const facebook::velox::exec::HashTable<ignoreNullKeys>* hashTable,
+ uint8_t* data,
+ size_t size) {
+ VELOX_CHECK_NOT_NULL(hashTable, "Hash table cannot be null");
+ VELOX_CHECK_NOT_NULL(data, "Serialized buffer cannot be null");
+ hashTable->serializeTo(data, size);
+}
+
+template <bool ignoreNullKeys>
+std::unique_ptr<facebook::velox::exec::HashTable<ignoreNullKeys>>
+HashTableSerializer::deserialize(const uint8_t* data, size_t size,
facebook::velox::memory::MemoryPool* pool) {
+ VELOX_CHECK_NOT_NULL(data, "Serialized data cannot be null");
+ VELOX_CHECK_GT(size, 0, "Invalid serialized data size");
+ VELOX_CHECK_NOT_NULL(pool, "Memory pool cannot be null");
+ return
facebook::velox::exec::HashTable<ignoreNullKeys>::deserializeFrom(data, size,
pool);
+}
+
+template size_t HashTableSerializer::serializedSize<true>(const
facebook::velox::exec::HashTable<true>*);
+
+template size_t HashTableSerializer::serializedSize<false>(const
facebook::velox::exec::HashTable<false>*);
+
+template void HashTableSerializer::serializeTo<true>(const
facebook::velox::exec::HashTable<true>*, uint8_t*, size_t);
+
+template void HashTableSerializer::serializeTo<false>(const
facebook::velox::exec::HashTable<false>*, uint8_t*, size_t);
+
+template std::unique_ptr<facebook::velox::exec::HashTable<true>>
+HashTableSerializer::deserialize<true>(const uint8_t*, size_t,
facebook::velox::memory::MemoryPool*);
+
+template std::unique_ptr<facebook::velox::exec::HashTable<false>>
+HashTableSerializer::deserialize<false>(const uint8_t*, size_t,
facebook::velox::memory::MemoryPool*);
+
+} // namespace gluten
diff --git a/cpp/velox/operators/hashjoin/HashTableSerializer.h
b/cpp/velox/operators/hashjoin/HashTableSerializer.h
new file mode 100644
index 0000000000..5dbcdf272e
--- /dev/null
+++ b/cpp/velox/operators/hashjoin/HashTableSerializer.h
@@ -0,0 +1,67 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include "velox/exec/HashTable.h"
+
+namespace gluten {
+
+/**
+ * HashTableSerializer provides serialization and deserialization for Velox
hash tables.
+ * This is a thin wrapper around the HashTable's native serialize/deserialize
methods
+ * from IBM Velox's verified implementation.
+ */
+class HashTableSerializer {
+ public:
+ /**
+ * Returns the exact serialized size of a hash table.
+ *
+ * @param hashTable The hash table to serialize
+ * @return Exact size in bytes
+ */
+ template <bool ignoreNullKeys>
+ static size_t serializedSize(const
facebook::velox::exec::HashTable<ignoreNullKeys>* hashTable);
+
+ /**
+ * Serialize a hash table directly into a caller-provided buffer.
+ *
+ * @param hashTable The hash table to serialize
+ * @param data Destination buffer
+ * @param size Exact destination size in bytes
+ */
+ template <bool ignoreNullKeys>
+ static void
+ serializeTo(const facebook::velox::exec::HashTable<ignoreNullKeys>*
hashTable, uint8_t* data, size_t size);
+
+ /**
+ * Deserialize a hash table from a memory buffer.
+ * Directly uses HashTable's deserialize() method from IBM Velox.
+ *
+ * @param data Pointer to serialized data
+ * @param size Size of serialized data
+ * @param pool Memory pool for allocations
+ * @return Deserialized hash table
+ */
+ template <bool ignoreNullKeys>
+ static std::unique_ptr<facebook::velox::exec::HashTable<ignoreNullKeys>>
+ deserialize(const uint8_t* data, size_t size,
facebook::velox::memory::MemoryPool* pool);
+};
+
+} // namespace gluten
diff --git a/docs/velox-configuration.md b/docs/velox-configuration.md
index 9a50e5ec8a..9772080887 100644
--- a/docs/velox-configuration.md
+++ b/docs/velox-configuration.md
@@ -27,6 +27,7 @@ nav_order: 16
| spark.gluten.sql.columnar.backend.velox.cudf.memoryResource
| ⚓ Static | async | GPU RMM memory resource.
[...]
| spark.gluten.sql.columnar.backend.velox.cudf.shuffleMaxPrefetchBytes
| 🔄 Dynamic | 1028MB | Maximum bytes to prefetch in CPU
memory during GPU shuffle read while waiting for GPU available.
[...]
| spark.gluten.sql.columnar.backend.velox.directorySizeGuess
| ⚓ Static | 32KB | Deprecated, rename to
spark.gluten.sql.columnar.backend.velox.footerEstimatedSize
[...]
+| spark.gluten.sql.columnar.backend.velox.driverSideBroadcastHashTableBuild
| 🔄 Dynamic | false | Enable driver-side broadcast hash
table build. When enabled, the hash table is built and serialized on the
driver, then broadcast to executors. When disabled, each executor builds its
own hash table from the broadcast data.
[...]
| spark.gluten.sql.columnar.backend.velox.enableTimestampNtzValidation
| 🔄 Dynamic | false | Enable validation fallback for
TimestampNTZ type. When true, any plan containing TimestampNTZ will fall back
to Spark execution. When false, allows native execution for TimestampNTZ scan.
[...]
| spark.gluten.sql.columnar.backend.velox.fileHandleCacheEnabled
| ⚓ Static | false | Disables caching if false. File
handle cache should be disabled if files are mutable, i.e. file content may
change while file path stays the same.
[...]
| spark.gluten.sql.columnar.backend.velox.filePreloadThreshold
| ⚓ Static | 1MB | Set the file preload threshold for
velox file scan, refer to Velox's file-preload-threshold
[...]
diff --git
a/gluten-arrow/src/main/scala/org/apache/gluten/memory/NativeMemoryManager.scala
b/gluten-arrow/src/main/scala/org/apache/gluten/memory/NativeMemoryManager.scala
index 3308f357a7..b1a9920cce 100644
---
a/gluten-arrow/src/main/scala/org/apache/gluten/memory/NativeMemoryManager.scala
+++
b/gluten-arrow/src/main/scala/org/apache/gluten/memory/NativeMemoryManager.scala
@@ -115,4 +115,5 @@ object NativeMemoryManager {
def apply(backendName: String, name: String): NativeMemoryManager = {
new Impl(backendName, name)
}
+
}
diff --git
a/gluten-arrow/src/main/scala/org/apache/gluten/runtime/Runtime.scala
b/gluten-arrow/src/main/scala/org/apache/gluten/runtime/Runtime.scala
index 885dd831e8..3a2572563a 100644
--- a/gluten-arrow/src/main/scala/org/apache/gluten/runtime/Runtime.scala
+++ b/gluten-arrow/src/main/scala/org/apache/gluten/runtime/Runtime.scala
@@ -89,4 +89,5 @@ object Runtime {
override def resourceName(): String = s"runtime"
}
+
}
diff --git
a/gluten-core/src/main/scala/org/apache/gluten/extension/BroadcastJoinContextTag.scala
b/gluten-core/src/main/scala/org/apache/gluten/extension/BroadcastJoinContextTag.scala
new file mode 100644
index 0000000000..16d2c55150
--- /dev/null
+++
b/gluten-core/src/main/scala/org/apache/gluten/extension/BroadcastJoinContextTag.scala
@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.extension
+
+import org.apache.spark.sql.catalyst.expressions.{AttributeSet, Expression}
+import org.apache.spark.sql.catalyst.plans.JoinType
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+
+/**
+ * Information needed to build a broadcast hash table. This is attached to
logical plan nodes to
+ * pass information from join operators to BroadcastExchangeExec.
+ */
+case class BroadcastJoinContextInfo(
+ joinType: JoinType,
+ buildRight: Boolean,
+ isNullAwareAntiJoin: Boolean,
+ condition: Option[Expression],
+ originalLeftKeys: Seq[Expression],
+ originalRightKeys: Seq[Expression],
+ // Filter-related information
+ buildOutputSet: AttributeSet)
+
+/** TreeNodeTag for storing broadcast join context information. */
+object BroadcastJoinContextTag {
+
+ /**
+ * Tag to store broadcast join context on logical plan nodes. This allows
BroadcastExchangeExec to
+ * access join information when building hash tables.
+ */
+ val BROADCAST_JOIN_CONTEXT: TreeNodeTag[Seq[BroadcastJoinContextInfo]] =
+ TreeNodeTag[Seq[BroadcastJoinContextInfo]]("gluten.broadcastJoinContext")
+}
diff --git
a/gluten-core/src/main/scala/org/apache/gluten/extension/GlutenJoinKeysCapture.scala
b/gluten-core/src/main/scala/org/apache/gluten/extension/GlutenJoinKeysCapture.scala
index 5d1cb8d90a..50d4096dbe 100644
---
a/gluten-core/src/main/scala/org/apache/gluten/extension/GlutenJoinKeysCapture.scala
+++
b/gluten-core/src/main/scala/org/apache/gluten/extension/GlutenJoinKeysCapture.scala
@@ -17,13 +17,14 @@
package org.apache.gluten.extension
import org.apache.spark.sql.catalyst.planning.{ExtractEquiJoinKeys,
ExtractSingleColumnNullAwareAntiJoin}
+import org.apache.spark.sql.catalyst.plans.logical.{HintInfo, JoinHint}
import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan}
import org.apache.spark.sql.execution.{SparkPlan, SparkStrategy}
/**
- * Strategy to capture join keys from logical plan before Spark's
JoinSelection transforms them.
- * This strategy runs early in the planning phase to preserve the original
join keys before any
- * transformations like rewriteKeyExpr.
+ * Strategy to capture join keys and context from logical plan before Spark's
JoinSelection
+ * transforms them. This strategy runs early in the planning phase to preserve
the original join
+ * keys and join context information before any transformations like
rewriteKeyExpr.
*/
case class GlutenJoinKeysCapture() extends SparkStrategy {
@@ -35,23 +36,72 @@ case class GlutenJoinKeysCapture() extends SparkStrategy {
plan match {
- case ExtractEquiJoinKeys(_, leftKeys, rightKeys, _, _, left, right, _) =>
- if (leftKeys.nonEmpty) {
- left.setTagValue(JoinKeysTag.ORIGINAL_JOIN_KEYS, leftKeys)
- }
- if (rightKeys.nonEmpty) {
- right.setTagValue(JoinKeysTag.ORIGINAL_JOIN_KEYS, rightKeys)
- }
+ case ExtractEquiJoinKeys(
+ joinType,
+ leftKeys,
+ rightKeys,
+ condition,
+ _,
+ left,
+ right,
+ hint) =>
+ // Set broadcast join context for the build side
+ // This information will be used by BroadcastExchangeExec to build
hash table
+ val buildRight = chooseBuildRight(joinType, left, right, hint)
+ val buildPlan = if (buildRight) right else left
+
+ val contextInfo = BroadcastJoinContextInfo(
+ joinType = joinType,
+ buildRight = buildRight,
+ isNullAwareAntiJoin = false,
+ condition = condition,
+ originalLeftKeys = leftKeys,
+ originalRightKeys = rightKeys,
+ buildOutputSet = buildPlan.outputSet
+ )
+
+ // Append context to both sides - support multiple joins using the
same table
+ val leftContexts =
left.getTagValue(BroadcastJoinContextTag.BROADCAST_JOIN_CONTEXT)
+ .getOrElse(Seq.empty)
+ left.setTagValue(
+ BroadcastJoinContextTag.BROADCAST_JOIN_CONTEXT,
+ leftContexts :+ contextInfo)
+
+ val rightContexts =
right.getTagValue(BroadcastJoinContextTag.BROADCAST_JOIN_CONTEXT)
+ .getOrElse(Seq.empty)
+ right.setTagValue(
+ BroadcastJoinContextTag.BROADCAST_JOIN_CONTEXT,
+ rightContexts :+ contextInfo)
Nil
case j @ ExtractSingleColumnNullAwareAntiJoin(leftKeys, rightKeys) =>
- if (leftKeys.nonEmpty) {
- j.left.setTagValue(JoinKeysTag.ORIGINAL_JOIN_KEYS, leftKeys)
- }
- if (rightKeys.nonEmpty) {
- j.right.setTagValue(JoinKeysTag.ORIGINAL_JOIN_KEYS, rightKeys)
- }
+ // Set broadcast join context for null-aware anti join
+ val buildRight = true
+ val buildPlan = j.right
+
+ val contextInfo = BroadcastJoinContextInfo(
+ joinType = j.joinType,
+ buildRight = buildRight,
+ isNullAwareAntiJoin = true,
+ condition = j.condition,
+ originalLeftKeys = leftKeys,
+ originalRightKeys = rightKeys,
+ buildOutputSet = buildPlan.outputSet
+ )
+
+ // Append context to both sides - support multiple joins using the
same table
+ val leftContexts =
j.left.getTagValue(BroadcastJoinContextTag.BROADCAST_JOIN_CONTEXT)
+ .getOrElse(Seq.empty)
+ j.left.setTagValue(
+ BroadcastJoinContextTag.BROADCAST_JOIN_CONTEXT,
+ leftContexts :+ contextInfo)
+
+ val rightContexts =
j.right.getTagValue(BroadcastJoinContextTag.BROADCAST_JOIN_CONTEXT)
+ .getOrElse(Seq.empty)
+ j.right.setTagValue(
+ BroadcastJoinContextTag.BROADCAST_JOIN_CONTEXT,
+ rightContexts :+ contextInfo)
Nil
@@ -59,4 +109,61 @@ case class GlutenJoinKeysCapture() extends SparkStrategy {
case _ => Nil
}
}
+
+ /**
+ * Determine if we can build on the right side for this join type. This is a
simplified version -
+ * actual logic may be more complex.
+ */
+ private def chooseBuildRight(
+ joinType: org.apache.spark.sql.catalyst.plans.JoinType,
+ left: LogicalPlan,
+ right: LogicalPlan,
+ hint: JoinHint): Boolean = {
+ import org.apache.spark.sql.catalyst.plans._
+ val hintBuildRight = hintedBuildRight(hint)
+ val canBuildLeft = joinType match {
+ case _: InnerLike => true
+ case RightOuter => true
+ case _ => false
+ }
+ val canBuildRight = joinType match {
+ case _: InnerLike => true
+ case LeftOuter => true
+ case LeftSemi => true
+ case LeftAnti => true
+ case ExistenceJoin(_) => true
+ case _ => false
+ }
+
+ if (hintBuildRight.contains(true) && canBuildRight) {
+ true
+ } else if (hintBuildRight.contains(false) && canBuildLeft) {
+ false
+ } else if (canBuildRight && !canBuildLeft) {
+ true
+ } else if (canBuildLeft && !canBuildRight) {
+ false
+ } else if (canBuildLeft && canBuildRight) {
+ right.stats.sizeInBytes <= left.stats.sizeInBytes
+ } else {
+ true
+ }
+ }
+
+ private def hintedBuildRight(hint: JoinHint): Option[Boolean] = {
+ def isBroadcast(hintInfo: HintInfo): Boolean = {
+ hintInfo.strategy.exists(_.toString.equalsIgnoreCase("BROADCAST"))
+ }
+
+ val broadcastLeft = hint.leftHint.exists(isBroadcast)
+ val broadcastRight = hint.rightHint.exists(isBroadcast)
+
+ if (broadcastLeft && !broadcastRight) {
+ Some(false)
+ } else if (broadcastRight && !broadcastLeft) {
+ Some(true)
+ } else {
+ None
+ }
+ }
}
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/SparkPlanExecApi.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/SparkPlanExecApi.scala
index 79f3d67c0e..b0efc00fdf 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/SparkPlanExecApi.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/SparkPlanExecApi.scala
@@ -437,7 +437,10 @@ trait SparkPlanExecApi {
child: SparkPlan,
numOutputRows: SQLMetric,
dataSize: SQLMetric,
- buildThreads: SQLMetric = null): BuildSideRelation
+ buildThreads: SQLMetric = null,
+ buildHashTableTimeMetric: SQLMetric = null,
+ serializeHashTableTimeMetric: SQLMetric = null,
+ serializedHashTableSizeMetric: SQLMetric = null): BuildSideRelation
def doCanonicalizeForBroadcastMode(mode: BroadcastMode): BroadcastMode = {
mode.canonicalized
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/JoinExecTransformer.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/JoinExecTransformer.scala
index eac9e6ab81..ba0b55bcbf 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/execution/JoinExecTransformer.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/execution/JoinExecTransformer.scala
@@ -153,6 +153,8 @@ trait HashJoinLikeExecTransformer extends BaseJoinExec with
TransformSupport {
override def metricsUpdater(): MetricsUpdater =
BackendsApiManager.getMetricsApiInstance.genHashJoinTransformerMetricsUpdater(metrics)
+ protected var joinParamsForMetrics: Option[JoinParams] = None
+
override def outputPartitioning: Partitioning = joinBuildSide match {
case BuildLeft =>
joinType match {
@@ -238,6 +240,7 @@ trait HashJoinLikeExecTransformer extends BaseJoinExec with
TransformSupport {
val operatorId = context.nextOperatorId(this.nodeName)
val joinParams = new JoinParams
+ joinParamsForMetrics = Some(joinParams)
if (JoinUtils.preProjectionNeeded(streamedKeyExprs)) {
joinParams.streamPreProjectionNeeded = true
}
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/substrait/SubstraitContext.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/substrait/SubstraitContext.scala
index 1ceb2d4155..8e448ade83 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/substrait/SubstraitContext.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/substrait/SubstraitContext.scala
@@ -35,6 +35,9 @@ case class JoinParams() {
// Whether the join is with condition
var isWithCondition = false
+
+ // Whether this join actually uses a pre-built serialized hash table from
driver-side build.
+ var usesDriverSideSerializedHashTable = false
}
case class AggregationParams() {
diff --git
a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarBroadcastExchangeExec.scala
b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarBroadcastExchangeExec.scala
index 2ead6bf07d..f8d5d1f858 100644
---
a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarBroadcastExchangeExec.scala
+++
b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarBroadcastExchangeExec.scala
@@ -76,7 +76,11 @@ case class ColumnarBroadcastExchangeExec(mode:
BroadcastMode, child: SparkPlan)
child,
longMetric("numOutputRows"),
longMetric("dataSize"),
- metrics.getOrElse("buildThreads", null))
+ metrics.getOrElse("buildThreads", null),
+ metrics.getOrElse("driverBuildHashTableTime", null),
+ metrics.getOrElse("driverSerializeHashTableTime", null),
+ metrics.getOrElse("serializedHashTableSize", null)
+ )
}
val broadcasted = GlutenTimeMetric.millis(longMetric("broadcastTime"))
{
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
