Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat: Enable Comet broadcast by default #213

Merged
merged 13 commits into from
Apr 5, 2024
Merged

feat: Enable Comet broadcast by default #213

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
5fa8781
feat: Remove COMET_EXEC_BROADCAST_ENABLED
viirya Mar 18, 2024
c804574
Fix
viirya Mar 29, 2024
2c4597a
Fix
viirya Mar 29, 2024
fd8e343
Update plan stability
viirya Mar 30, 2024
6cb63ed
Fix
viirya Apr 1, 2024
25eec59
Remove unused import and class
viirya Apr 3, 2024
437a126
Fix
viirya Apr 3, 2024
db5961d
Remove unused imports
viirya Apr 3, 2024
96b4aa4
Fix
viirya Apr 4, 2024
778466a
Fix scala style
viirya Apr 4, 2024
3628bd6
fix
viirya Apr 4, 2024
1712f0e
Fix
viirya Apr 5, 2024
32f3ae1
Update diff
viirya Apr 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions common/src/main/scala/org/apache/comet/CometConf.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
33 changes: 22 additions & 11 deletions spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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, _) =>
Expand Down Expand Up @@ -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 =
Expand Down
4 changes: 2 additions & 2 deletions spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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(
Expand Down
Loading