Skip to content

[VL] Reuse driver-built broadcast hash tables across queries #12872

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
zhouyuan wants to merge 1 commit into from
Closed

[VL] Reuse driver-built broadcast hash tables across queries #12872

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
265 changes: 165 additions & 100 deletions ...ends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ import org.apache.gluten.substrait.SubstraitContext
import org.apache.gluten.substrait.expression.{ExpressionBuilder, ExpressionNode, WindowFunctionNode}
import org.apache.gluten.vectorized.{ColumnarBatchSerializer, ColumnarBatchSerializeResult}

import org.apache.spark.{ShuffleDependency, SparkEnv, SparkException}
import org.apache.spark.{broadcast, ShuffleDependency, SparkEnv, SparkException}
import org.apache.spark.api.python.{ColumnarArrowEvalPythonExec, PullOutArrowEvalPythonPreProjectHelper}
import org.apache.spark.internal.Logging
import org.apache.spark.memory.SparkMemoryUtil
Expand Down Expand Up @@ -879,57 +879,15 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with Logging {
val useOffheapBroadcastBuildRelation =
VeloxConfig.get.enableBroadcastBuildRelationInOffheap

val serialized: Seq[ColumnarBatchSerializeResult] = newChild
.executeColumnar()
.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(
"Cannot broadcast the table that is larger than " +
s"${SparkMemoryUtil.bytesToString(GlutenConfig.get.maxBroadcastTableSize)}: " +
s"${SparkMemoryUtil.bytesToString(rawSize)}")
}
numOutputRows += buildSideRowCount
dataSize += rawSize

val rawThreads =
math
.ceil(dataSize.value.toDouble / VeloxConfig.get.veloxBroadcastHashTableBuildTargetBytes)
.toInt
val buildThreadsValue = if (rawThreads < 1) 1 else rawThreads
buildThreads += buildThreadsValue

// Create the base ColumnarBuildSideRelation first
val columnarRelation = if (useOffheapBroadcastBuildRelation) {
TaskResources.runUnsafe {
UnsafeColumnarBuildSideRelation(
newOutput,
serialized.flatMap(_.offHeapData().asScala),
mode,
newBuildKeys,
offload,
buildThreadsValue)
}
} else {
ColumnarBuildSideRelation(
newOutput,
serialized.flatMap(_.onHeapData().asScala).toArray,
mode,
newBuildKeys,
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) {
// The join context is resolved before the build side is collected, because it is also part of
// the key of the driver-side broadcast relation cache: a cache hit skips the collect.
val buildContextOpt: Option[BroadcastHashJoinContext] = 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.
Expand All @@ -951,8 +909,13 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with Logging {
}
}

joinContextOpt match {
case Some(joinContext) =>
if (joinContextOpt.isEmpty) {
// No join context available - fall back to executor-side build
logInfo(s"No broadcast join context found in logical plan, using executor-side build")
}

joinContextOpt.map {
joinContext =>
// We have join context information - build hash table on driver
logInfo(
s"Building hash table on driver in BroadcastExchangeExec " +
Expand Down Expand Up @@ -1009,7 +972,7 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with Logging {
} else {
joinContext.originalLeftKeys
}
val buildContext = BroadcastHashJoinContext(
BroadcastHashJoinContext(
buildSideJoinKeys = if (newBuildKeys.nonEmpty) newBuildKeys else joinKeys,
substraitJoinType = substraitJoinType,
buildRight = joinContext.buildRight,
Expand All @@ -1026,65 +989,167 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with Logging {
serializeHashTableTimeMetric = Option(serializeHashTableTimeMetric),
serializedHashTableSizeMetric = Option(serializedHashTableSizeMetric)
)
}
} else {
None
}

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
)
}
// A hash table built by a previous query can be reused as long as the build side plan and the
// join properties it was built with are the same. This is what makes repeated workloads, e.g.
// the concurrent streams of a TPC-DS throughput run, skip the whole build side pipeline.
val relationCacheKeyOpt = buildContextOpt
.filter(_ => VeloxDriverBroadcastRelationCache.enabled)
.flatMap {
buildContext =>
VeloxDriverBroadcastRelationCache.keyOf(
newChild,
newOutput,
newBuildKeys,
buildContext,
useOffheapBroadcastBuildRelation)
}

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
}
val cachedRelation = relationCacheKeyOpt.flatMap(VeloxDriverBroadcastRelationCache.getIfPresent)
if (cachedRelation.isDefined) {
val cached = cachedRelation.get
logInfo(
s"Reusing the broadcast relation built by a previous query: " +
s"size=${cached.serializedSize} bytes, rows=${cached.numRows}")
// The build side job is skipped, so report the metrics of the job that built the relation.
numOutputRows += cached.numRows
dataSize += cached.rawSize
Option(buildThreads).foreach(_ += cached.buildThreads)
Option(serializedHashTableSizeMetric).foreach(_.set(cached.serializedSize))
return cached.relation
}

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
val serialized: Seq[ColumnarBatchSerializeResult] = newChild
.executeColumnar()
.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(
"Cannot broadcast the table that is larger than " +
s"${SparkMemoryUtil.bytesToString(GlutenConfig.get.maxBroadcastTableSize)}: " +
s"${SparkMemoryUtil.bytesToString(rawSize)}")
}
numOutputRows += buildSideRowCount
dataSize += rawSize

val rawThreads =
math
.ceil(dataSize.value.toDouble / VeloxConfig.get.veloxBroadcastHashTableBuildTargetBytes)
.toInt
val buildThreadsValue = if (rawThreads < 1) 1 else rawThreads
buildThreads += buildThreadsValue

// Create the base ColumnarBuildSideRelation first
val columnarRelation = if (useOffheapBroadcastBuildRelation) {
TaskResources.runUnsafe {
UnsafeColumnarBuildSideRelation(
newOutput,
serialized.flatMap(_.offHeapData().asScala),
mode,
newBuildKeys,
offload,
buildThreadsValue)
}
} else {
// Return ColumnarBuildSideRelation for executor-side build (legacy approach)
columnarRelation
ColumnarBuildSideRelation(
newOutput,
serialized.flatMap(_.onHeapData().asScala).toArray,
mode,
newBuildKeys,
offload,
buildThreadsValue)
}

buildContextOpt match {
case Some(buildContext) =>
try {
// Build and serialize hash table on driver. When the relation is kept by the driver-side
// relation cache, that cache owns the serialized hash table, so it is not tracked by the
// per-exchange cache as well.
val trackSerializedHashTable = relationCacheKeyOpt.isEmpty
val (serializedHashTable, safeMode) = columnarRelation match {
case rel: ColumnarBuildSideRelation =>
(
VeloxBroadcastBuildSideCache
.buildAndSerializeOnDriverInBroadcastExchange(
rel,
buildContext,
buildSideRowCount,
trackSerializedHashTable),
rel.safeBroadcastMode
)
case rel: UnsafeColumnarBuildSideRelation =>
(
VeloxBroadcastBuildSideCache
.buildAndSerializeOnDriverInBroadcastExchange(
rel,
buildContext,
buildSideRowCount,
trackSerializedHashTable),
rel.getSafeBroadcastMode
)
}

logInfo(
s"Successfully built hash table on driver: " +
s"size=${serializedHashTable.sizeInBytes} bytes, " +
s"rows=${serializedHashTable.numRows}, " +
s"joinType=${buildContext.substraitJoinType}, " +
s"broadcastId=${buildContext.buildHashTableId}")

// Return SerializedHashTableBroadcastRelation
val relation = SerializedHashTableBroadcastRelation(
serializedHashTable,
safeMode,
newOutput,
0L, // buildTimeMs - tracked inside SerializedBroadcastHashTable
0L // serializeTimeMs - tracked inside SerializedBroadcastHashTable
)

relationCacheKeyOpt.foreach {
key =>
VeloxDriverBroadcastRelationCache.put(
key,
relation,
buildSideRowCount,
rawSize,
buildThreadsValue,
serializedHashTable.sizeInBytes)
}

relation
} catch {
case e: Exception =>
logWarning(
s"Failed to build hash table on driver for " +
s"broadcastId=${buildContext.buildHashTableId}, " +
s"falling back to executor-side build: ${e.getMessage}",
e
)
columnarRelation
}

case None =>
// Return ColumnarBuildSideRelation for executor-side build (legacy approach)
columnarRelation
}
}

override def cachedBroadcast(relation: BuildSideRelation): Option[broadcast.Broadcast[Any]] =
VeloxDriverBroadcastRelationCache.cachedBroadcast(relation)

override def cacheBroadcast(
relation: BuildSideRelation,
broadcasted: broadcast.Broadcast[Any]): Unit =
VeloxDriverBroadcastRelationCache.registerBroadcast(relation, broadcasted)

override def doCanonicalizeForBroadcastMode(mode: BroadcastMode): BroadcastMode = {
mode match {
case hash: HashedRelationBroadcastMode
Expand Down
30 changes: 30 additions & 0 deletions backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,12 @@ class VeloxConfig(conf: SQLConf) extends GlutenConfig(conf) {
def enableDriverSideBroadcastHashTableBuild: Boolean =
getConf(VELOX_DRIVER_SIDE_BROADCAST_HASH_TABLE_BUILD)

def enableDriverSideBroadcastHashTableCache: Boolean =
getConf(VELOX_DRIVER_SIDE_BROADCAST_HASH_TABLE_CACHE_ENABLED)

def driverSideBroadcastHashTableCacheMaxSize: Long =
getConf(VELOX_DRIVER_SIDE_BROADCAST_HASH_TABLE_CACHE_MAX_SIZE)

def enableGpuAsyncShuffleReader: Boolean = getConf(ENABLE_GPU_ASYNC_SHUFFLE_READER)

def gpuAsyncReaderMaxPrefetchBytes: Long = getConf(GPU_ASYNC_SHUFFLE_READER_MAX_PREFETCH_BYTES)
Expand Down Expand Up @@ -773,6 +779,30 @@ object VeloxConfig extends ConfigRegistry {
.booleanConf
.createWithDefault(false)

val VELOX_DRIVER_SIDE_BROADCAST_HASH_TABLE_CACHE_ENABLED =
buildConf("spark.gluten.sql.columnar.backend.velox.driverSideBroadcastHashTableCache.enabled")
.doc(
"Reuse the hash tables built by driver-side broadcast hash table build across queries of " +
"the same application. Cache entries are keyed by the canonicalized build side plan " +
"plus the join properties the hash table was built with, so identical broadcast build " +
"sides coming from different queries, e.g. the concurrent streams of a TPC-DS " +
"throughput run, share one hash table instead of collecting, building and " +
"broadcasting it again. Only takes effect when " +
"spark.gluten.sql.columnar.backend.velox.driverSideBroadcastHashTableBuild is enabled. " +
"Reusing assumes the data behind a build side does not change while the application " +
"is running, and the cached hash tables hold driver off-heap memory until they are " +
"evicted, so this is disabled by default.")
.booleanConf
.createWithDefault(false)

val VELOX_DRIVER_SIDE_BROADCAST_HASH_TABLE_CACHE_MAX_SIZE =
buildConf("spark.gluten.sql.columnar.backend.velox.driverSideBroadcastHashTableCache.maxSize")
.doc(
"The maximum total size of the serialized hash tables kept by the driver-side broadcast " +
"hash table cache. Least-recently-used entries are evicted once the limit is exceeded.")
.bytesConf(ByteUnit.BYTE)
.createWithDefaultString("1GB")

val QUERY_TRACE_ENABLED = buildConf("spark.gluten.sql.columnar.backend.velox.queryTraceEnabled")
.doc("Enable query tracing flag.")
.booleanConf
Expand Down
13 changes: 11 additions & 2 deletions backends-velox/src/main/scala/org/apache/gluten/execution/VeloxBroadcastBuildSideCache.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,11 +111,17 @@ object VeloxBroadcastBuildSideCache
* BroadcastExchangeExec and doesn't need a broadcast variable.
*
* This is the Spark-native approach where hash table is built in BroadcastExchangeExec.
*
* @param trackInDriverCache
* whether the serialized hash table is kept in `driverSerializedCache`, which also owns its
* off-heap memory. Set to false when the caller keeps the built relation alive itself, e.g.
* `VeloxDriverBroadcastRelationCache`, so that the hash table has a single owner.
*/
def buildAndSerializeOnDriverInBroadcastExchange(
relation: BuildSideRelation,
broadcastContext: BroadcastHashJoinContext,
numRows: Long): SerializedBroadcastHashTable = {
numRows: Long,
trackInDriverCache: Boolean = true): SerializedBroadcastHashTable = {

val broadcastId = broadcastContext.buildHashTableId

Expand Down Expand Up @@ -180,7 +186,9 @@ object VeloxBroadcastBuildSideCache
broadcastContext.serializeHashTableTimeMetric.foreach(_ += serializeTimeMs)
broadcastContext.serializedHashTableSizeMetric.foreach(_ += result.sizeInBytes)

driverSerializedCache.put(broadcastId, result)
if (trackInDriverCache) {
driverSerializedCache.put(broadcastId, result)
}
result
} finally {
resetRelation(droppedDuplicates)
Expand Down Expand Up @@ -233,6 +241,7 @@ object VeloxBroadcastBuildSideCache
def cleanAll(): Unit = {
buildSideRelationCache.invalidateAll()
driverSerializedCache.invalidateAll()
VeloxDriverBroadcastRelationCache.cleanAll()
}

override def onRemoval(
Expand Down
Loading
Loading