feat: Remove COMET_EXEC_BROADCAST_ENABLED

apache · Mar 29, 2024 · 5fa8781 · 5fa8781
1 parent aa6ddc5
commit 5fa8781
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 17 deletions.
diff --git a/common/src/main/scala/org/apache/comet/CometConf.scala b/common/src/main/scala/org/apache/comet/CometConf.scala
@@ -139,12 +139,13 @@ object CometConf {
     .booleanConf
     .createWithDefault(false)
 
-  val COMET_EXEC_BROADCAST_ENABLED: ConfigEntry[Boolean] =
+  val COMET_EXEC_BROADCAST_FORCE_ENABLED: ConfigEntry[Boolean] =
     conf(s"$COMET_EXEC_CONFIG_PREFIX.broadcast.enabled")
       .doc(
-        "Whether to enable broadcasting for Comet native operators. By default, " +
-          "this config is false. Note that this feature is not fully supported yet " +
-          "and only enabled for test purpose.")
+        "Whether to force enabling broadcasting for Comet native operators. By default, " +
+          "this config is false. Comet broadcast feature will be enabled automatically by " +
+          "Comet extension. But for unit tests, we need this feature to force enabling it " +
+          "for invalid cases. So this config is only used for unit test.")
       .booleanConf
       .createWithDefault(false)
 

diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
@@ -43,7 +43,7 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 
 import org.apache.comet.CometConf._
-import org.apache.comet.CometSparkSessionExtensions.{isANSIEnabled, isCometBroadCastEnabled, isCometColumnarShuffleEnabled, isCometEnabled, isCometExecEnabled, isCometOperatorEnabled, isCometScan, isCometScanEnabled, isCometShuffleEnabled, isSchemaSupported}
+import org.apache.comet.CometSparkSessionExtensions.{isANSIEnabled, isCometBroadCastForceEnabled, isCometColumnarShuffleEnabled, isCometEnabled, isCometExecEnabled, isCometOperatorEnabled, isCometScan, isCometScanEnabled, isCometShuffleEnabled, isSchemaSupported}
 import org.apache.comet.parquet.{CometParquetScan, SupportsComet}
 import org.apache.comet.serde.OperatorOuterClass.Operator
 import org.apache.comet.serde.QueryPlanSerde
@@ -442,15 +442,26 @@ class CometSparkSessionExtensions
               s
           }
 
-        case b: BroadcastExchangeExec
-            if isCometNative(b.child) && isCometOperatorEnabled(conf, "broadcastExchangeExec") &&
-              isCometBroadCastEnabled(conf) =>
-          QueryPlanSerde.operator2Proto(b) match {
-            case Some(nativeOp) =>
-              val cometOp = CometBroadcastExchangeExec(b, b.child)
-              CometSinkPlaceHolder(nativeOp, b, cometOp)
-            case None => b
+        // `CometBroadcastExchangeExec`'s broadcast output is not compatible with Spark's broadcast
+        // exchange. It is only used for Comet native execution. We only transform Spark broadcast
+        // exchange to Comet broadcast exchange if its downstream is a Comet native plan or if the
+        // broadcast exchange is forced to be enabled by Comet config.
+        case plan
+            if (isCometNative(plan) || isCometBroadCastForceEnabled(conf)) &&
+              plan.children.exists(_.isInstanceOf[BroadcastExchangeExec]) =>
+          val newChildren = plan.children.map {
+            case b: BroadcastExchangeExec
+                if isCometNative(b.child) &&
+                  isCometOperatorEnabled(conf, "broadcastExchangeExec") =>
+              QueryPlanSerde.operator2Proto(b) match {
+                case Some(nativeOp) =>
+                  val cometOp = CometBroadcastExchangeExec(b, b.child)
+                  CometSinkPlaceHolder(nativeOp, b, cometOp)
+                case None => b
+              }
+            case other => other
           }
+          plan.withNewChildren(newChildren)
 
         // For AQE shuffle stage on a Comet shuffle exchange
         case s @ ShuffleQueryStageExec(_, _: CometShuffleExchangeExec, _) =>
@@ -651,8 +662,8 @@ object CometSparkSessionExtensions extends Logging {
     !conf.getConfString(operatorDisabledFlag, "false").toBoolean
   }
 
-  private[comet] def isCometBroadCastEnabled(conf: SQLConf): Boolean = {
-    COMET_EXEC_BROADCAST_ENABLED.get(conf)
+  private[comet] def isCometBroadCastForceEnabled(conf: SQLConf): Boolean = {
+    COMET_EXEC_BROADCAST_FORCE_ENABLED.get(conf)
   }
 
   private[comet] def isCometShuffleEnabled(conf: SQLConf): Boolean =

diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
@@ -73,7 +73,7 @@ class CometExecSuite extends CometTestBase {
 
   test("CometBroadcastExchangeExec") {
     assume(isSpark34Plus, "ChunkedByteBuffer is not serializable before Spark 3.4+")
-    withSQLConf(CometConf.COMET_EXEC_BROADCAST_ENABLED.key -> "true") {
+    withSQLConf(CometConf.COMET_EXEC_BROADCAST_FORCE_ENABLED.key -> "true") {
       withParquetTable((0 until 5).map(i => (i, i + 1)), "tbl_a") {
         withParquetTable((0 until 5).map(i => (i, i + 1)), "tbl_b") {
           val df = sql(
@@ -99,7 +99,7 @@ class CometExecSuite extends CometTestBase {
 
   test("CometBroadcastExchangeExec: empty broadcast") {
     assume(isSpark34Plus, "ChunkedByteBuffer is not serializable before Spark 3.4+")
-    withSQLConf(CometConf.COMET_EXEC_BROADCAST_ENABLED.key -> "true") {
+    withSQLConf(CometConf.COMET_EXEC_BROADCAST_FORCE_ENABLED.key -> "true") {
       withParquetTable((0 until 5).map(i => (i, i + 1)), "tbl_a") {
         withParquetTable((0 until 5).map(i => (i, i + 1)), "tbl_b") {
           val df = sql(