From e844a0e39db66dff2128773e3bb4010508f89918 Mon Sep 17 00:00:00 2001 From: Kazuyuki Tanimura Date: Wed, 15 May 2024 14:56:43 -0700 Subject: [PATCH 1/3] test: Fix explain with exteded info comet test --- .../test/scala/org/apache/comet/CometExpressionSuite.scala | 2 +- spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index dbb46e078..40d15281b 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -1399,7 +1399,7 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { CometConf.COMET_EXEC_ENABLED.key -> "true", CometConf.COMET_SHUFFLE_ENFORCE_MODE_ENABLED.key -> "true", CometConf.COMET_EXEC_ALL_OPERATOR_ENABLED.key -> "true", - "spark.sql.extendedExplainProvider" -> "org.apache.comet.ExtendedExplainInfo") { + "spark.sql.extendedExplainProviders" -> "org.apache.comet.ExtendedExplainInfo") { val table = "test" withTable(table) { sql(s"create table $table(c0 int, c1 int , c2 float) using parquet") diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala index 8e05bf26a..31c53508c 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -249,7 +249,7 @@ abstract class CometTestBase var dfSpark: Dataset[Row] = null withSQLConf( CometConf.COMET_ENABLED.key -> "false", - "spark.sql.extendedExplainProvider" -> "") { + "spark.sql.extendedExplainProviders" -> "") { dfSpark = Dataset.ofRows(spark, df.logicalPlan) expected = dfSpark.collect() } @@ -259,7 +259,7 @@ abstract class CometTestBase dfSpark.queryExecution.explainString(ExtendedMode), dfComet.queryExecution.explainString(ExtendedMode)) if (supportsExtendedExplainInfo(dfSpark.queryExecution)) { - assert(diff.contains(expectedInfo)) + assert(expectedInfo.exists(s => diff.contains(s))) } val extendedInfo = new ExtendedExplainInfo().generateExtendedInfo(dfComet.queryExecution.executedPlan) From f5380e556ce08d7d887cbffca0c6376de474d99d Mon Sep 17 00:00:00 2001 From: Kazuyuki Tanimura Date: Thu, 16 May 2024 10:08:27 -0700 Subject: [PATCH 2/3] address review comments --- .../apache/comet/shims/ShimCometSparkSessionExtensions.scala | 2 ++ .../test/scala/org/apache/comet/CometExpressionSuite.scala | 2 +- spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala | 4 ++-- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala index ffec1bd40..22c567bb4 100644 --- a/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala +++ b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala @@ -64,4 +64,6 @@ object ShimCometSparkSessionExtensions { } true } + + protected val EXTENDED_EXPLAIN_PROVIDERS_KEY = "spark.sql.extendedExplainProviders" } diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index 40d15281b..f3fd50e9e 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -1399,7 +1399,7 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { CometConf.COMET_EXEC_ENABLED.key -> "true", CometConf.COMET_SHUFFLE_ENFORCE_MODE_ENABLED.key -> "true", CometConf.COMET_EXEC_ALL_OPERATOR_ENABLED.key -> "true", - "spark.sql.extendedExplainProviders" -> "org.apache.comet.ExtendedExplainInfo") { + EXTENDED_EXPLAIN_PROVIDERS_KEY -> "org.apache.comet.ExtendedExplainInfo") { val table = "test" withTable(table) { sql(s"create table $table(c0 int, c1 int , c2 float) using parquet") diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala index 31c53508c..68905ac34 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -249,7 +249,7 @@ abstract class CometTestBase var dfSpark: Dataset[Row] = null withSQLConf( CometConf.COMET_ENABLED.key -> "false", - "spark.sql.extendedExplainProviders" -> "") { + EXTENDED_EXPLAIN_PROVIDERS_KEY -> "") { dfSpark = Dataset.ofRows(spark, df.logicalPlan) expected = dfSpark.collect() } @@ -259,7 +259,7 @@ abstract class CometTestBase dfSpark.queryExecution.explainString(ExtendedMode), dfComet.queryExecution.explainString(ExtendedMode)) if (supportsExtendedExplainInfo(dfSpark.queryExecution)) { - assert(expectedInfo.exists(s => diff.contains(s))) + assert(expectedInfo.forall(s => diff.contains(s))) } val extendedInfo = new ExtendedExplainInfo().generateExtendedInfo(dfComet.queryExecution.executedPlan) From 1bde6ea56114507a2aa2526f22530c881288356e Mon Sep 17 00:00:00 2001 From: Kazuyuki Tanimura Date: Thu, 16 May 2024 10:25:55 -0700 Subject: [PATCH 3/3] address review comments --- .../comet/shims/ShimCometSparkSessionExtensions.scala | 7 +++++-- .../test/scala/org/apache/spark/sql/CometTestBase.scala | 4 +--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala index 22c567bb4..eb04c68ab 100644 --- a/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala +++ b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometSparkSessionExtensions.scala @@ -40,6 +40,11 @@ trait ShimCometSparkSessionExtensions { */ def getOffset(limit: LimitExec): Int = getOffsetOpt(limit).getOrElse(0) + /** + * TODO: delete after dropping Spark 3.x support and directly call + * SQLConf.EXTENDED_EXPLAIN_PROVIDERS.key + */ + protected val EXTENDED_EXPLAIN_PROVIDERS_KEY = "spark.sql.extendedExplainProviders" } object ShimCometSparkSessionExtensions { @@ -64,6 +69,4 @@ object ShimCometSparkSessionExtensions { } true } - - protected val EXTENDED_EXPLAIN_PROVIDERS_KEY = "spark.sql.extendedExplainProviders" } diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala index 68905ac34..112d35b13 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -247,9 +247,7 @@ abstract class CometTestBase expectedInfo: Set[String]): Unit = { var expected: Array[Row] = Array.empty var dfSpark: Dataset[Row] = null - withSQLConf( - CometConf.COMET_ENABLED.key -> "false", - EXTENDED_EXPLAIN_PROVIDERS_KEY -> "") { + withSQLConf(CometConf.COMET_ENABLED.key -> "false", EXTENDED_EXPLAIN_PROVIDERS_KEY -> "") { dfSpark = Dataset.ofRows(spark, df.logicalPlan) expected = dfSpark.collect() }