add checks

apache · Aug 20, 2024 · 9509d63 · 9509d63
1 parent db7d4be
commit 9509d63
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHBackend.scala b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHBackend.scala
@@ -407,4 +407,11 @@ object CHBackendSettings extends BackendSettingsApi with Logging {
       }
   }
 
+  def getBroadcastThreshold: Long = {
+    val conf = SQLConf.get
+    conf
+      .getConf(SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD)
+      .getOrElse(conf.autoBroadcastJoinThreshold)
+  }
+
 }
diff --git a/...ickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHSparkPlanExecApi.scala b/...ickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHSparkPlanExecApi.scala
@@ -539,6 +539,15 @@ class CHSparkPlanExecApi extends SparkPlanExecApi {
       CHExecUtil.buildSideRDD(dataSize, newChild).collect
 
     val batches = countsAndBytes.map(_._2)
+    val totalBatchesBytes = batches.map(_.length).sum
+    if (
+      totalBatchesBytes < 0 || totalBatchesBytes.toLong > CHBackendSettings.getBroadcastThreshold
+    ) {
+      throw new SparkException(
+        s"Cannot broadcast the table that is larger than threshold:" +
+          s" ${CHBackendSettings.getBroadcastThreshold}. Ensure the shuffle written" +
+          s"bytes is collected properly.")
+    }
     val rawSize = dataSize.value
     if (rawSize >= BroadcastExchangeExec.MAX_BROADCAST_TABLE_BYTES) {
       throw new SparkException(