This is an automated email from the ASF dual-hosted git repository.
philo-he pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 3b08a23683 [GLUTEN-12251][VL] Support merging broadcast batches for
BHJ performance (#12259)
3b08a23683 is described below
commit 3b08a23683cf18fe34c04a99a0609c315c4d8c08
Author: Shilong Duan <[email protected]>
AuthorDate: Wed Jun 17 13:16:18 2026 +0800
[GLUTEN-12251][VL] Support merging broadcast batches for BHJ performance
(#12259)
---
.../org/apache/gluten/config/VeloxConfig.scala | 14 ++++++
.../spark/sql/execution/BroadcastUtils.scala | 53 ++++++++++++++++++++++
.../gluten/execution/VeloxHashJoinSuite.scala | 32 +++++++++++++
cpp/core/jni/JniWrapper.cc | 28 ++++++++++++
docs/velox-configuration.md | 1 +
.../ColumnarBatchSerializerJniWrapper.java | 2 +
6 files changed, 130 insertions(+)
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 951eb6b444..13ec2dfebe 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
@@ -65,6 +65,9 @@ class VeloxConfig(conf: SQLConf) extends GlutenConfig(conf) {
def enableBroadcastBuildRelationInOffheap: Boolean =
getConf(VELOX_BROADCAST_BUILD_RELATION_USE_OFFHEAP)
+ def broadcastBuildMergeBatches: Boolean =
+ getConf(VELOX_BROADCAST_BUILD_MERGE_BATCHES)
+
def enableBroadcastBuildOncePerExecutor: Boolean =
getConf(VELOX_BROADCAST_BUILD_HASHTABLE_ONCE_PER_EXECUTOR)
@@ -617,6 +620,17 @@ object VeloxConfig extends ConfigRegistry {
.booleanConf
.createWithDefault(false)
+ val VELOX_BROADCAST_BUILD_MERGE_BATCHES =
+ buildConf("spark.gluten.velox.broadcastBuild.mergeBatches")
+ .doc(
+ "If enabled, all columnar batches in a broadcast build relation will
be " +
+ "serialized into a single buffer to reduce the number of addInput
calls in " +
+ "HashBuild operator. This can significantly improve BHJ performance
when " +
+ "the broadcast table has many small batches, but may increase
driver-side " +
+ "peak memory and is not suitable for very large broadcasts.")
+ .booleanConf
+ .createWithDefault(false)
+
val VELOX_HASHMAP_ABANDON_BUILD_DUPHASH_MIN_ROWS =
buildConf("spark.gluten.velox.abandonDedupHashMap.minRows")
.experimental()
diff --git
a/backends-velox/src/main/scala/org/apache/spark/sql/execution/BroadcastUtils.scala
b/backends-velox/src/main/scala/org/apache/spark/sql/execution/BroadcastUtils.scala
index cf3f9ccca4..9c092e1070 100644
---
a/backends-velox/src/main/scala/org/apache/spark/sql/execution/BroadcastUtils.scala
+++
b/backends-velox/src/main/scala/org/apache/spark/sql/execution/BroadcastUtils.scala
@@ -150,6 +150,7 @@ object BroadcastUtils {
}
def serializeStream(batches: Iterator[ColumnarBatch]):
ColumnarBatchSerializeResult = {
+ val mergeBatches = VeloxConfig.get.broadcastBuildMergeBatches
val filtered = batches
.filter(_.numRows() != 0)
.map(
@@ -157,6 +158,58 @@ object BroadcastUtils {
ColumnarBatches.retain(b)
b
})
+
+ if (mergeBatches) {
+ serializeStreamMerged(filtered)
+ } else {
+ serializeStreamPerBatch(filtered)
+ }
+ }
+
+ private def serializeStreamMerged(
+ filtered: Iterator[ColumnarBatch]): ColumnarBatchSerializeResult = {
+ var numRows: Long = 0
+ val handles = new ArrayBuffer[Long]()
+ val retainedBatches = new ArrayBuffer[ColumnarBatch]()
+ filtered.foreach {
+ b =>
+ numRows += b.numRows()
+ handles +=
ColumnarBatches.getNativeHandle(BackendsApiManager.getBackendName, b)
+ retainedBatches += b
+ }
+ if (handles.nonEmpty) {
+ try {
+ val jniWrapper = ColumnarBatchSerializerJniWrapper
+ .create(
+ Runtimes
+ .contextInstance(BackendsApiManager.getBackendName,
"BroadcastUtils#serializeStream"))
+ val merged = jniWrapper.serializeAll(handles.toArray)
+
+ // Merging produces a single payload. Spark/Gluten broadcast
serialization cannot handle
+ // a single buffer larger than 2GB, so fall back to per-batch
serialization in that case.
+ if (merged.size() > Integer.MAX_VALUE) {
+ merged.toUnsafeByteArray.release()
+ val batchesToSerialize = retainedBatches.toArray
+ retainedBatches.clear()
+ return serializeStreamPerBatch(batchesToSerialize.iterator)
+ }
+
+ val useOffheapBroadcastBuildRelation =
+ VeloxConfig.get.enableBroadcastBuildRelationInOffheap
+ new ColumnarBatchSerializeResult(
+ useOffheapBroadcastBuildRelation,
+ numRows,
+ java.util.Collections.singletonList(merged))
+ } finally {
+ retainedBatches.foreach(ColumnarBatches.release)
+ }
+ } else {
+ ColumnarBatchSerializeResult.EMPTY
+ }
+ }
+
+ private def serializeStreamPerBatch(
+ filtered: Iterator[ColumnarBatch]): ColumnarBatchSerializeResult = {
var numRows: Long = 0
val values = filtered
.map(
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 86565aa42b..9c3f4f6e4e 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
@@ -279,6 +279,38 @@ class VeloxHashJoinSuite extends
VeloxWholeStageTransformerSuite {
}
}
+ test("Broadcast build mergeBatches: merged vs per-batch produce equivalent
results") {
+ Seq("true", "false").foreach {
+ mergeBatches =>
+ withSQLConf(
+ VeloxConfig.VELOX_BROADCAST_BUILD_MERGE_BATCHES.key -> mergeBatches,
+ "spark.sql.autoBroadcastJoinThreshold" -> "10MB",
+ "spark.sql.adaptive.enabled" -> "false",
+ // Force small batches so the build side has multiple batches to
merge.
+ GlutenConfig.COLUMNAR_MAX_BATCH_SIZE.key -> "16"
+ ) {
+ withTable("t_probe", "t_build") {
+ spark.range(200).selectExpr("id as key", "id * 2 as
v").write.saveAsTable("t_probe")
+ spark.range(50).selectExpr("id as key", "id + 1 as
v").write.saveAsTable("t_build")
+
+ val query =
+ """
+ |SELECT p.key, p.v, b.v AS bv
+ |FROM t_probe p JOIN t_build b ON p.key = b.key
+ |ORDER BY p.key
+ |""".stripMargin
+
+ runQueryAndCompare(query) {
+ df =>
+ val plan = df.queryExecution.executedPlan
+ val bhj = plan.collect { case j:
BroadcastHashJoinExecTransformer => j }
+ assert(bhj.nonEmpty, s"Should use BHJ when
mergeBatches=$mergeBatches")
+ }
+ }
+ }
+ }
+ }
+
test("Broadcast join with multiple cast expressions in join keys") {
withSQLConf(
("spark.sql.autoBroadcastJoinThreshold", "10MB"),
diff --git a/cpp/core/jni/JniWrapper.cc b/cpp/core/jni/JniWrapper.cc
index 827c5ad8bd..86da59d2ef 100644
--- a/cpp/core/jni/JniWrapper.cc
+++ b/cpp/core/jni/JniWrapper.cc
@@ -1310,6 +1310,34 @@ JNIEXPORT jobject JNICALL
Java_org_apache_gluten_vectorized_ColumnarBatchSeriali
JNI_METHOD_END(nullptr)
}
+JNIEXPORT jobject JNICALL
Java_org_apache_gluten_vectorized_ColumnarBatchSerializerJniWrapper_serializeAll(
// NOLINT
+ JNIEnv* env,
+ jobject wrapper,
+ jlongArray handles) {
+ JNI_METHOD_START
+ auto ctx = getRuntime(env, wrapper);
+ GLUTEN_CHECK(handles != nullptr, "serializeAll requires non-null handles
array");
+ const jsize numBatches = env->GetArrayLength(handles);
+ GLUTEN_CHECK(numBatches > 0, "serializeAll requires at least one batch");
+
+ auto safeArray = getLongArrayElementsSafe(env, handles);
+ auto serializer = ctx->createColumnarBatchSerializer(nullptr);
+ for (int32_t i = 0; i < numBatches; i++) {
+ auto batch = ObjectStore::retrieve<ColumnarBatch>(safeArray.elems()[i]);
+ GLUTEN_DCHECK(
+ batch != nullptr, "Cannot find the ColumnarBatch with handle " +
std::to_string(safeArray.elems()[i]));
+ serializer->append(batch);
+ }
+ auto serializedSize = serializer->maxSerializedSize();
+ auto byteBuffer = env->CallStaticObjectMethod(jniUnsafeByteBufferClass,
jniUnsafeByteBufferAllocate, serializedSize);
+ auto byteBufferAddress = env->CallLongMethod(byteBuffer,
jniUnsafeByteBufferAddress);
+ auto byteBufferSize = env->CallLongMethod(byteBuffer,
jniUnsafeByteBufferSize);
+ serializer->serializeTo(reinterpret_cast<uint8_t*>(byteBufferAddress),
byteBufferSize);
+
+ return byteBuffer;
+ JNI_METHOD_END(nullptr)
+}
+
// Framed [magic | statsLen | statsBlob | bytesLen | bytesBlob] entry point.
Uses the
// ColumnarBatchSerializer::framedSerializeWithStats virtual hook; non-Velox
backends inherit
// the default empty-vector return so callers fall back to the legacy
serialize() path.
diff --git a/docs/velox-configuration.md b/docs/velox-configuration.md
index 952a6c7f9c..5b67ef4afe 100644
--- a/docs/velox-configuration.md
+++ b/docs/velox-configuration.md
@@ -81,6 +81,7 @@ nav_order: 16
| spark.gluten.sql.enable.enhancedFeatures
| 🔄 Dynamic | true | Enable some features including iceberg
native write and other features.
[...]
| spark.gluten.sql.rewrite.castArrayToString
| 🔄 Dynamic | true | When true, rewrite `cast(array as
String)` to `concat('[', array_join(array, ', ', null), ']')` to allow
offloading to Velox.
[...]
| spark.gluten.velox.broadcast.build.targetBytesPerThread
| âš“ Static | 32MB | It is used to calculate the number of
hash table build threads. Based on our testing across various thresholds (1MB
to 128MB), we recommend a value of 32MB or 64MB, as these consistently provided
the most significant performance gains.
[...]
+| spark.gluten.velox.broadcastBuild.mergeBatches
| 🔄 Dynamic | false | If enabled, all columnar batches in a
broadcast build relation will be serialized into a single buffer to reduce the
number of addInput calls in HashBuild operator. This can significantly improve
BHJ performance when the broadcast table has many small batches, but may
increase driver-side peak memory and is not suitable for very large broadcasts.
|
| spark.gluten.velox.castFromVarcharAddTrimNode
| 🔄 Dynamic | false | If true, will add a trim node which
has the same semantic as vanilla Spark to CAST-from-varchar.Otherwise, do
nothing.
[...]
## Gluten Velox backend *experimental* configurations
diff --git
a/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ColumnarBatchSerializerJniWrapper.java
b/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ColumnarBatchSerializerJniWrapper.java
index 9a5247c823..513eb02ecf 100644
---
a/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ColumnarBatchSerializerJniWrapper.java
+++
b/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ColumnarBatchSerializerJniWrapper.java
@@ -39,6 +39,8 @@ public class ColumnarBatchSerializerJniWrapper implements
RuntimeAware {
public native JniUnsafeByteBuffer serialize(long handle);
+ public native JniUnsafeByteBuffer serializeAll(long[] handles);
+
// Framed [magic | statsLen | statsBlob | bytesLen | bytesBlob] payload
produced by
// VeloxColumnarBatchSerializer::framedSerializeWithStats. Returns byte[]
(not
// JniUnsafeByteBuffer) because the framed wire is small enough that the
simpler return type
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
