diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6e7bb4271b..a81c28f69e 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -668,6 +668,14 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_TRIAL_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.trial.enabled") + .category(CATEGORY_EXEC_EXPLAIN) + .doc("Analyze each query as if Comet were enabled and emit a coverage report, but do " + + "not actually offload execution to native.") + .booleanConf + .createWithDefault(false) + val COMET_EXPLAIN_CODEGEN_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.explain.codegen.enabled") .withAlternative("spark.comet.explainCodegen.enabled") diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index ef2f37371c..58517a87a7 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -49,7 +49,7 @@ import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ -import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo} +import org.apache.comet.{CometConf, CometCoverageStats, CometExplainInfo, ExtendedExplainInfo} import org.apache.comet.CometConf.{COMET_SPARK_TO_ARROW_ENABLED, COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST} import org.apache.comet.CometSparkSessionExtensions._ import org.apache.comet.rules.CometExecRule.allExecs @@ -617,6 +617,20 @@ case class CometExecRule(session: SparkSession) } } + // Trial mode: log how much of this query Comet would accelerate (measured from the + // fully-converted plan), then execute on Spark by reverting the native scans that + // CometScanRule produced, so nothing is actually offloaded to native. The rest of the + // native conversion only exists in `newPlan`, which we discard here. + if (CometConf.COMET_TRIAL_ENABLED.get()) { + logWarning( + s"[Comet trial] ${CometCoverageStats.forPlan(newPlan)}\n" + + new ExtendedExplainInfo().generateExtendedInfo(newPlan)) + return plan.transformUp { + case s: CometScanExec => s.wrapped + case s: CometBatchScanExec => s.wrapped.copy(runtimeFilters = s.runtimeFilters) + } + } + // Remove placeholders newPlan = newPlan.transform { case CometSinkPlaceHolder(_, _, s) => s diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala index 5444a89fa3..0f70959a9d 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -797,4 +797,36 @@ class CometExecRuleSuite extends CometTestBase { } } + test("trial mode analyzes the plan but executes entirely on Spark") { + withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") { + val query = "SELECT _2, count(*) FROM tbl GROUP BY _2" + withSQLConf( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_TRIAL_ENABLED.key -> "true") { + val df = sql(query) + // Results are identical to vanilla Spark. + checkSparkAnswer(df) + // Nothing is offloaded: the executed plan contains no Comet operators. + val cometNodes = stripAQEPlan(df.queryExecution.executedPlan).collect { + case p: CometPlan => p + } + assert( + cometNodes.isEmpty, + s"trial mode must not offload, but found Comet operators: $cometNodes") + } + + // Sanity check: with trial mode off, the same query is accelerated by Comet. + withSQLConf( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_TRIAL_ENABLED.key -> "false") { + val cometNodes = stripAQEPlan(sql(query).queryExecution.executedPlan).collect { + case p: CometPlan => p + } + assert(cometNodes.nonEmpty, "expected Comet operators when trial mode is disabled") + } + } + } + }