From 17637248d789fdfdeff1e98964508d09984b376e Mon Sep 17 00:00:00 2001 From: Ke Jia Date: Wed, 1 Jul 2026 02:59:15 +0100 Subject: [PATCH 1/3] Build hash table in driver side --- .../clickhouse/CHSparkPlanExecApi.scala | 5 +- .../gluten/vectorized/HashJoinBuilder.java | 13 + .../backendsapi/velox/VeloxMetricsApi.scala | 16 +- .../velox/VeloxSparkPlanExecApi.scala | 227 ++++++++++++++++-- .../apache/gluten/config/VeloxConfig.scala | 12 + .../execution/HashJoinExecTransformer.scala | 76 +++++- .../SerializedBroadcastHashTable.scala | 155 ++++++++++++ .../VeloxBroadcastBuildSideCache.scala | 138 ++++++++++- .../VeloxSerializedBroadcastRDD.scala | 66 +++++ .../gluten/metrics/JoinMetricsUpdater.scala | 10 +- .../execution/ColumnarBuildSideRelation.scala | 125 ++++++++-- ...SerializedHashTableBroadcastRelation.scala | 137 +++++++++++ .../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 | 111 ++++++++- .../operators/hashjoin/HashTableBuilder.cc | 5 +- .../operators/hashjoin/HashTableSerializer.cc | 62 +++++ .../operators/hashjoin/HashTableSerializer.h | 67 ++++++ docs/velox-configuration.md | 1 + .../gluten/memory/NativeMemoryManager.scala | 1 + .../org/apache/gluten/runtime/Runtime.scala | 1 + .../extension/BroadcastJoinContextTag.scala | 46 ++++ .../extension/GlutenJoinKeysCapture.scala | 139 +++++++++-- .../gluten/backendsapi/SparkPlanExecApi.scala | 5 +- .../execution/JoinExecTransformer.scala | 3 + .../gluten/substrait/SubstraitContext.scala | 3 + .../ColumnarBroadcastExchangeExec.scala | 6 +- 32 files changed, 1715 insertions(+), 85 deletions(-) create mode 100644 backends-velox/src/main/scala/org/apache/gluten/execution/SerializedBroadcastHashTable.scala create mode 100644 backends-velox/src/main/scala/org/apache/gluten/execution/VeloxSerializedBroadcastRDD.scala create mode 100644 backends-velox/src/main/scala/org/apache/spark/sql/execution/SerializedHashTableBroadcastRelation.scala create mode 100644 cpp/velox/operators/hashjoin/HashTableSerializer.cc create mode 100644 cpp/velox/operators/hashjoin/HashTableSerializer.h create mode 100644 gluten-core/src/main/scala/org/apache/gluten/extension/BroadcastJoinContextTag.scala 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 36c9813fbb8..c00738434dc 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 106d8b98170..dbd8a3a5f70 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 long rtHandle() { 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 83e83555fab..272574fe79b 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 2448e04e0cd..d320fa01798 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 dcb79a462fe..0295d223902 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(true) + 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 06f59acf03f..e77b2c61154 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 00000000000..a59326daef8 --- /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 451c633be4b..0d492817b14 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 00000000000..761d0000978 --- /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 b056cd36a8e..b7d32f94bd4 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 7042f34e936..7c406653b05 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 @@ -212,24 +212,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 00000000000..19949d41079 --- /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 fa3dcc69340..3288f301659 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, _) => @@ -182,24 +184,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 ddd76f917db..523aae65c51 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 c80950aa1fd..64dccf5bb73 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"), @@ -481,4 +622,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 c881d77ed10..89471119e3e 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 bbb20a84611..d935a12d048 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 11873471575..d7259879c59 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 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*>(hashTable.get()); + if (hashTableFalse != nullptr) { + return HashTableSerializer::serializedSize(hashTableFalse); + } + + auto* hashTableTrue = dynamic_cast*>(hashTable.get()); + VELOX_CHECK_NOT_NULL(hashTableTrue, "Hash table must be either HashTable or HashTable"); + return HashTableSerializer::serializedSize(hashTableTrue); +} + +void serializeHashTableTo(std::shared_ptr builder, uint8_t* data, size_t size) { + VELOX_CHECK_NOT_NULL(builder, "Hash table builder cannot be null"); + VELOX_CHECK_NOT_NULL(data, "Serialized buffer cannot be null"); + + auto hashTable = builder->hashTable(); + VELOX_CHECK_NOT_NULL(hashTable, "Hash table cannot be null"); + + auto* hashTableFalse = dynamic_cast*>(hashTable.get()); + if (hashTableFalse != nullptr) { + HashTableSerializer::serializeTo(hashTableFalse, data, size); + return; + } + + auto* hashTableTrue = dynamic_cast*>(hashTable.get()); + VELOX_CHECK_NOT_NULL(hashTableTrue, "Hash table must be either HashTable or HashTable"); + HashTableSerializer::serializeTo(hashTableTrue, data, size); +} + +std::shared_ptr +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 hashTable; + if (ignoreNullKeys) { + auto derived = HashTableSerializer::deserialize(data, size, poolPtr); + hashTable = std::move(derived); + } else { + auto derived = HashTableSerializer::deserialize(data, size, poolPtr); + hashTable = std::move(derived); + } + + std::vector> emptyKeys; + std::vector emptyChannels; + + auto keyTypes = hashTable->rows()->keyTypes(); + std::vector 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( + 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 47f89d17996..d75d2d663bc 100644 --- a/cpp/velox/jni/JniHashTable.h +++ b/cpp/velox/jni/JniHashTable.h @@ -91,6 +91,16 @@ std::shared_ptr nativeHashTableBuild( long getJoin(const std::string& hashTableId); +// Return the exact serialized hash table size for direct buffer allocation. +size_t serializedHashTableSize(std::shared_ptr builder); + +// Serialize hash table directly to a caller-provided buffer. +void serializeHashTableTo(std::shared_ptr builder, uint8_t* data, size_t size); + +// Deserialize hash table from broadcast data with explicit ignoreNullKeys parameter +std::shared_ptr +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 369d6716ed4..50a79f3af99 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,103 @@ 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(address), + static_cast(size), + static_cast(ignoreNullKeys), + static_cast(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(hashTableHandle); + auto hashTable = builder->hashTable(); + VELOX_CHECK_NOT_NULL(hashTable, "Hash table cannot be null"); + return static_cast(dynamic_cast*>(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(hashTableHandle); + return static_cast(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(hashTableHandle); + auto hashTable = builder->hashTable(); + VELOX_CHECK_NOT_NULL(hashTable, "Hash table cannot be null"); + + auto* baseTable = dynamic_cast(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(bloomFilter.get()); + if (bfFilter != nullptr) { + bloomFilterBlocksByteSize += bfFilter->blocksByteSize(); + } + } + return static_cast(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(hashTableHandle); + return static_cast(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(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), + serializedSize, + "Hash table buffer size mismatch"); + gluten::serializeHashTableTo(builder, reinterpret_cast(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 363edea4ecc..4c3129fccf8 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 00000000000..76365d76a42 --- /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 +size_t HashTableSerializer::serializedSize(const facebook::velox::exec::HashTable* hashTable) { + VELOX_CHECK_NOT_NULL(hashTable, "Hash table cannot be null"); + return hashTable->serializedSize(); +} + +template +void HashTableSerializer::serializeTo( + const facebook::velox::exec::HashTable* 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 +std::unique_ptr> +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::deserializeFrom(data, size, pool); +} + +template size_t HashTableSerializer::serializedSize(const facebook::velox::exec::HashTable*); + +template size_t HashTableSerializer::serializedSize(const facebook::velox::exec::HashTable*); + +template void HashTableSerializer::serializeTo(const facebook::velox::exec::HashTable*, uint8_t*, size_t); + +template void HashTableSerializer::serializeTo(const facebook::velox::exec::HashTable*, uint8_t*, size_t); + +template std::unique_ptr> +HashTableSerializer::deserialize(const uint8_t*, size_t, facebook::velox::memory::MemoryPool*); + +template std::unique_ptr> +HashTableSerializer::deserialize(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 00000000000..5dbcdf272e4 --- /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 +#include +#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 + static size_t serializedSize(const facebook::velox::exec::HashTable* 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 + static void + serializeTo(const facebook::velox::exec::HashTable* 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 + static std::unique_ptr> + 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 9a50e5ec8aa..a363d865416 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 | true | 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 3308f357a7c..b1a9920cce1 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 885dd831e84..3a2572563aa 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 00000000000..16d2c55150f --- /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 5d1cb8d90a8..50d4096dbe0 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 79f3d67c0ec..b0efc00fdfe 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 eac9e6ab81a..ba0b55bcbfb 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 1ceb2d4155a..8e448ade83c 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 2ead6bf07d5..f8d5d1f8580 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")) { From 8c4d8a0c75b23a69b546d238328abcde1393ec0c Mon Sep 17 00:00:00 2001 From: Ke Jia Date: Fri, 3 Jul 2026 05:37:04 +0100 Subject: [PATCH 2/3] format --- .../main/scala/org/apache/gluten/config/VeloxConfig.scala | 2 +- cpp/velox/jni/VeloxJniWrapper.cc | 5 +---- 2 files changed, 2 insertions(+), 5 deletions(-) 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 0295d223902..dacb835b8c5 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 @@ -692,7 +692,7 @@ object VeloxConfig extends ConfigRegistry { "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(true) + .createWithDefault(false) val QUERY_TRACE_ENABLED = buildConf("spark.gluten.sql.columnar.backend.velox.queryTraceEnabled") .doc("Enable query tracing flag.") diff --git a/cpp/velox/jni/VeloxJniWrapper.cc b/cpp/velox/jni/VeloxJniWrapper.cc index 50a79f3af99..36798f5cded 100644 --- a/cpp/velox/jni/VeloxJniWrapper.cc +++ b/cpp/velox/jni/VeloxJniWrapper.cc @@ -1247,10 +1247,7 @@ JNIEXPORT void JNICALL Java_org_apache_gluten_vectorized_HashJoinBuilder_seriali 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), - serializedSize, - "Hash table buffer size mismatch"); + VELOX_CHECK_EQ(static_cast(size), serializedSize, "Hash table buffer size mismatch"); gluten::serializeHashTableTo(builder, reinterpret_cast(address), serializedSize); JNI_METHOD_END() } From 016737df7286b14244c7892cda3e8713947adac1 Mon Sep 17 00:00:00 2001 From: Ke Jia Date: Fri, 3 Jul 2026 06:12:20 +0100 Subject: [PATCH 3/3] fix --- docs/velox-configuration.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/velox-configuration.md b/docs/velox-configuration.md index a363d865416..97720808878 100644 --- a/docs/velox-configuration.md +++ b/docs/velox-configuration.md @@ -27,7 +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 | true | 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.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 |