add checks

apache · Aug 20, 2024 · 2cd649e · 2cd649e
1 parent db7d4be
commit 2cd649e
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
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,9 +539,18 @@ 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 GlutenException(
+        s"Cannot broadcast the table (${totalBatchesBytes}) 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(
+      throw new GlutenException(
         s"Cannot broadcast the table that is larger than 8GB: ${rawSize >> 30} GB")
     }
     val rowCount = countsAndBytes.map(_._1).sum