Google Git Google Git
Sign in
apache/datafusion-comet/refs/heads/codex/comet-df-bitmap/./dev/diffs/4.2.0.diff
blob: 3335dfb3713831ac152b37287a658d3702116236 [file]
diff --git a/core/src/test/scala/org/apache/spark/storage/FallbackStorageSuite.scala b/core/src/test/scala/org/apache/spark/storage/FallbackStorageSuite.scala
index 6df8bc85b51..dabb75e2b75 100644
--- a/core/src/test/scala/org/apache/spark/storage/FallbackStorageSuite.scala
+++ b/core/src/test/scala/org/apache/spark/storage/FallbackStorageSuite.scala
@@ -268,6 +268,11 @@ class FallbackStorageSuite extends SparkFunSuite with LocalSparkContext {
}
test("Upload from all decommissioned executors") {
+ // Comet replaces Spark's shuffle with its own native shuffle, which is incompatible with
+ // the fallback storage migration path used by BlockManagerDecommissioner.
+ val cometEnv = System.getenv("ENABLE_COMET")
+ assume(cometEnv == null || cometEnv == "0" || cometEnv == "false",
+ "Skipped when Comet is enabled: incompatible with Comet native shuffle storage")
sc = new SparkContext(getSparkConf(2, 2))
withSpark(sc) { sc =>
TestUtils.waitUntilExecutorsUp(sc, 2, 60000)
@@ -298,6 +303,11 @@ class FallbackStorageSuite extends SparkFunSuite with LocalSparkContext {
}
test("Upload multi stages") {
+ // Comet replaces Spark's shuffle with its own native shuffle, which is incompatible with
+ // the fallback storage migration path used by BlockManagerDecommissioner.
+ val cometEnv = System.getenv("ENABLE_COMET")
+ assume(cometEnv == null || cometEnv == "0" || cometEnv == "false",
+ "Skipped when Comet is enabled: incompatible with Comet native shuffle storage")
sc = new SparkContext(getSparkConf())
withSpark(sc) { sc =>
TestUtils.waitUntilExecutorsUp(sc, 1, 60000)
@@ -332,6 +342,11 @@ class FallbackStorageSuite extends SparkFunSuite with LocalSparkContext {
CompressionCodec.shortCompressionCodecNames.keys.foreach { codec =>
test(s"$codec - Newly added executors should access old data from remote storage") {
+ // Comet replaces Spark's shuffle with its own native shuffle, which is incompatible with
+ // the fallback storage migration path used by BlockManagerDecommissioner.
+ val cometEnv = System.getenv("ENABLE_COMET")
+ assume(cometEnv == null || cometEnv == "0" || cometEnv == "false",
+ "Skipped when Comet is enabled: incompatible with Comet native shuffle storage")
sc = new SparkContext(getSparkConf(2, 0).set(IO_COMPRESSION_CODEC, codec))
withSpark(sc) { sc =>
TestUtils.waitUntilExecutorsUp(sc, 2, 60000)
diff --git a/pom.xml b/pom.xml
index 46558134f41..7fa990be776 100644
--- a/pom.xml
+++ b/pom.xml
@@ -154,6 +154,8 @@
<kryo.version>4.0.3</kryo.version>
<ivy.version>2.5.3</ivy.version>
<oro.version>2.0.8</oro.version>
+ <spark.version.short>4.2</spark.version.short>
+ <comet.version>1.2.0-SNAPSHOT</comet.version>
<!--
If you change codahale.metrics.version, you also need to change
the link to metrics.dropwizard.io in docs/monitoring.md.
@@ -2646,6 +2648,25 @@
<artifactId>arpack</artifactId>
<version>${netlib.ludovic.dev.version}</version>
</dependency>
+ <dependency>
+ <groupId>org.apache.datafusion</groupId>
+ <artifactId>comet-spark-spark${spark.version.short}_${scala.binary.version}</artifactId>
+ <version>${comet.version}</version>
+ <exclusions>
+ <exclusion>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-sql_${scala.binary.version}</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-core_${scala.binary.version}</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-catalyst_${scala.binary.version}</artifactId>
+ </exclusion>
+ </exclusions>
+ </dependency>
<!-- SPARK-16484 add `datasketches-java` for support Datasketches HllSketch -->
<dependency>
<groupId>org.apache.datasketches</groupId>
diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala
index 7c5e49f3a69..28f83381fa0 100644
--- a/project/SparkBuild.scala
+++ b/project/SparkBuild.scala
@@ -455,9 +455,11 @@ object SparkBuild extends PomBuild {
/* Spark SQL Core settings */
enable(SQL.settings)(sql)
+ enable(CometTestSettings.settings)(sql)
/* Hive console settings */
enable(Hive.settings)(hive)
+ enable(CometTestSettings.settings)(hive)
enable(HiveThriftServer.settings)(hiveThriftServer)
@@ -1426,6 +1428,25 @@ object SQL {
}
}
+/**
+ * Comet disables itself when its shuffle manager is not registered. When Comet is enabled, make
+ * that shuffle manager the default for every test SparkConf in the modules that have Comet on
+ * their classpath, so that suites which build their own SparkSession or SparkContext run Comet
+ * too, not just the ones that go through SharedSparkSession or TestHive.
+ */
+object CometTestSettings {
+ private val cometEnabled =
+ sys.env.get("ENABLE_COMET").forall(v => v == "1" || v.equalsIgnoreCase("true"))
+
+ lazy val settings: Seq[Setting[_]] =
+ if (cometEnabled && !sys.props.contains("spark.shuffle.manager")) {
+ Seq((Test / javaOptions) +=
+ "-Dspark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
+ } else {
+ Seq.empty
+ }
+}
+
object Hive {
lazy val settings = Seq(
diff --git a/sql/core/pom.xml b/sql/core/pom.xml
index 82810f181ac..21a83831188 100644
--- a/sql/core/pom.xml
+++ b/sql/core/pom.xml
@@ -97,6 +97,10 @@
<groupId>org.apache.spark</groupId>
<artifactId>spark-tags_${scala.binary.version}</artifactId>
</dependency>
+ <dependency>
+ <groupId>org.apache.datafusion</groupId>
+ <artifactId>comet-spark-spark${spark.version.short}_${scala.binary.version}</artifactId>
+ </dependency>
<!--
This spark-tags test-dep is needed even though it isn't used in this module, otherwise testing-cmds that exclude
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala
index f03b4796314..d32d2c49ce3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala
@@ -1248,6 +1248,23 @@ object SparkSession extends SparkSessionCompanion with Logging {
extensions
}
+ /**
+ * Whether Comet extension is enabled
+ */
+ def isCometEnabled: Boolean = {
+ val v = System.getenv("ENABLE_COMET")
+ v == null || v == "1" || v.toBoolean
+ }
+
+
+ private def loadCometExtension(sparkContext: SparkContext): Seq[String] = {
+ if (sparkContext.getConf.getBoolean("spark.comet.enabled", isCometEnabled)) {
+ Seq("org.apache.comet.CometSparkSessionExtensions")
+ } else {
+ Seq.empty
+ }
+ }
+
/**
* Initialize extensions specified in [[StaticSQLConf]]. The classes will be applied to the
* extensions passed into this function.
@@ -1257,7 +1274,8 @@ object SparkSession extends SparkSessionCompanion with Logging {
extensions: SparkSessionExtensions): SparkSessionExtensions = {
val extensionConfClassNames = sparkContext.conf.get(StaticSQLConf.SPARK_SESSION_EXTENSIONS)
.getOrElse(Seq.empty)
- extensionConfClassNames.foreach { extensionConfClassName =>
+ val extensionClassNames = extensionConfClassNames ++ loadCometExtension(sparkContext)
+ extensionClassNames.foreach { extensionConfClassName =>
try {
val extensionConfClass = Utils.classForName(extensionConfClassName)
val extensionConf = extensionConfClass.getConstructor().newInstance()
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
index 4410fe50912..43bcce2a038 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql.execution
import org.apache.spark.annotation.DeveloperApi
import org.apache.spark.sql.catalyst.plans.logical.{EmptyRelation, LogicalPlan}
+import org.apache.spark.sql.comet.CometScanExec
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec}
import org.apache.spark.sql.execution.adaptive.LogicalQueryStage
import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
@@ -84,6 +85,7 @@ private[execution] object SparkPlanInfo {
// dump the file scan metadata (e.g file path) to event log
val metadata = plan match {
case fileScan: FileSourceScanLike => fileScan.metadata
+ case cometScan: CometScanExec => cometScan.metadata
case _ => Map[String, String]()
}
val childrenInfo = children.flatMap {
diff --git a/sql/core/src/test/resources/sql-tests/inputs/decimalArithmeticOperations.sql b/sql/core/src/test/resources/sql-tests/inputs/decimalArithmeticOperations.sql
index 13bbd9d81b7..541cdfb1e04 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/decimalArithmeticOperations.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/decimalArithmeticOperations.sql
@@ -15,6 +15,12 @@
-- limitations under the License.
--
+-- TODO: Disabled due to one of the test failed for Spark4.0
+-- TODO: https://github.com/apache/datafusion-comet/issues/1948
+-- The following query failed
+-- select /*+ COALESCE(1) */ id, a+b, a-b, a*b, a/b from decimals_test order by id
+--SET spark.comet.enabled = false
+
CREATE TEMPORARY VIEW t AS SELECT 1.0 as a, 0.0 as b;
-- division, remainder and pmod by 0 return NULL
diff --git a/sql/core/src/test/resources/sql-tests/inputs/explain-aqe.sql b/sql/core/src/test/resources/sql-tests/inputs/explain-aqe.sql
index 7aef901da4f..f3d6e18926d 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/explain-aqe.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/explain-aqe.sql
@@ -2,3 +2,4 @@
--SET spark.sql.adaptive.enabled=true
--SET spark.sql.maxMetadataStringLength = 500
+--SET spark.comet.enabled = false
diff --git a/sql/core/src/test/resources/sql-tests/inputs/explain-cbo.sql b/sql/core/src/test/resources/sql-tests/inputs/explain-cbo.sql
index eeb2180f7a5..afd1b5ec289 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/explain-cbo.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/explain-cbo.sql
@@ -1,5 +1,6 @@
--SET spark.sql.cbo.enabled=true
--SET spark.sql.maxMetadataStringLength = 500
+--SET spark.comet.enabled = false
CREATE TABLE explain_temp1(a INT, b INT) USING PARQUET;
CREATE TABLE explain_temp2(c INT, d INT) USING PARQUET;
diff --git a/sql/core/src/test/resources/sql-tests/inputs/explain.sql b/sql/core/src/test/resources/sql-tests/inputs/explain.sql
index 96dddafd82a..efedbbb3c0e 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/explain.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/explain.sql
@@ -1,6 +1,7 @@
--SET spark.sql.codegen.wholeStage = true
--SET spark.sql.adaptive.enabled = false
--SET spark.sql.maxMetadataStringLength = 500
+--SET spark.comet.enabled = false
-- Test tables
CREATE table explain_temp1 (key int, val int) USING PARQUET;
diff --git a/sql/core/src/test/resources/sql-tests/inputs/join-nearest-by.sql b/sql/core/src/test/resources/sql-tests/inputs/join-nearest-by.sql
index 40cfa87c4cd..4e1d5407100 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/join-nearest-by.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/join-nearest-by.sql
@@ -1,5 +1,10 @@
-- Test cases for NEAREST BY top-K ranking join.
+-- The EXPLAIN queries in this file record Spark's operator names in the golden file, which
+-- Comet replaces with its own operators, so run the whole file without Comet (same approach as
+-- explain.sql / explain-aqe.sql / explain-cbo.sql).
+--SET spark.comet.enabled = false
+
CREATE VIEW users(user_id, score) AS VALUES (1, 10.0), (2, 20.0), (3, 30.0);
CREATE VIEW products(product, pscore) AS VALUES ('A', 11.0), ('B', 22.0), ('C', 5.0);
diff --git a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part3.sql b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part3.sql
index 41fd4de2a09..162d5a817b6 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part3.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/aggregates_part3.sql
@@ -6,6 +6,10 @@
-- https://github.com/postgres/postgres/blob/REL_12_BETA2/src/test/regress/sql/aggregates.sql#L352-L605
-- Test aggregate operator with codegen on and off.
+
+-- Floating-point precision difference between DataFusion and JVM for FILTER aggregates
+--SET spark.comet.enabled = false
+
--CONFIG_DIM1 spark.sql.codegen.wholeStage=true
--CONFIG_DIM1 spark.sql.codegen.wholeStage=false,spark.sql.codegen.factoryMode=CODEGEN_ONLY
--CONFIG_DIM1 spark.sql.codegen.wholeStage=false,spark.sql.codegen.factoryMode=NO_CODEGEN
diff --git a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/int4.sql b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/int4.sql
index 3a409eea348..26e9aaf215c 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/int4.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/int4.sql
@@ -6,6 +6,9 @@
-- https://github.com/postgres/postgres/blob/REL_12_BETA2/src/test/regress/sql/int4.sql
--
+-- TODO: https://github.com/apache/datafusion-comet/issues/551
+--SET spark.comet.enabled = false
+
CREATE TABLE INT4_TBL(f1 int) USING parquet;
-- [SPARK-28023] Trim the string when cast string type to other types
diff --git a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/int8.sql b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/int8.sql
index fac23b4a26f..98b12ae5ccc 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/int8.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/int8.sql
@@ -6,6 +6,10 @@
-- Test int8 64-bit integers.
-- https://github.com/postgres/postgres/blob/REL_12_BETA2/src/test/regress/sql/int8.sql
--
+
+-- TODO: https://github.com/apache/datafusion-comet/issues/551
+--SET spark.comet.enabled = false
+
CREATE TABLE INT8_TBL(q1 bigint, q2 bigint) USING parquet;
-- PostgreSQL implicitly casts string literals to data with integral types, but
diff --git a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/select_having.sql b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/select_having.sql
index 0efe0877e9b..f9df0400c99 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/select_having.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/postgreSQL/select_having.sql
@@ -6,6 +6,9 @@
-- https://github.com/postgres/postgres/blob/REL_12_BETA2/src/test/regress/sql/select_having.sql
--
+-- TODO: https://github.com/apache/datafusion-comet/issues/551
+--SET spark.comet.enabled = false
+
-- load test data
CREATE TABLE test_having (a int, b int, c string, d string) USING parquet;
INSERT INTO test_having VALUES (0, 1, 'XXXX', 'A');
diff --git a/sql/core/src/test/resources/sql-tests/inputs/replacing-missing-expression-with-alias.sql b/sql/core/src/test/resources/sql-tests/inputs/replacing-missing-expression-with-alias.sql
index 0f238033fad..590faebe77a 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/replacing-missing-expression-with-alias.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/replacing-missing-expression-with-alias.sql
@@ -1,3 +1,7 @@
+-- Comet: ORDER BY ties can produce a different row order from the Spark golden files.
+-- https://github.com/apache/datafusion-comet/issues/5570
+--SET spark.comet.enabled=false
+
-- Replace expression with alias that has semantically equal child if expression is not in output
SELECT col1 + 1 AS a FROM VALUES(1) GROUP BY a ORDER BY col1 + 1;
SELECT col1 + 1 AS a, a AS b FROM VALUES(1) GROUP BY a ORDER BY col1 + 1;
diff --git a/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-limit.sql b/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-limit.sql
index 7c816d8a416..b1551a2b296 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-limit.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-limit.sql
@@ -1,6 +1,23 @@
-- A test suite for IN LIMIT in parent side, subquery, and both predicate subquery
-- It includes correlated cases.
+-- TODO: Disabled due to one of the test failed for Spark4.0
+-- TODO: https://github.com/apache/datafusion-comet/issues/1948
+-- The following query failed
+-- SELECT Count(DISTINCT( t1a )),
+-- t1b
+-- FROM t1
+-- WHERE t1d NOT IN (SELECT t2d
+-- FROM t2
+-- WHERE t2b > t1b
+-- ORDER BY t2b DESC nulls first, t2d
+-- LIMIT 1
+-- OFFSET 1)
+-- GROUP BY t1b
+-- ORDER BY t1b NULLS last
+-- LIMIT 1
+-- OFFSET 1;
+--SET spark.comet.enabled = false
--CONFIG_DIM1 spark.sql.optimizeNullAwareAntiJoin=true
--CONFIG_DIM1 spark.sql.optimizeNullAwareAntiJoin=false
@@ -61,6 +78,7 @@ WHERE t1a IN (SELECT t2a
WHERE t1d = t2d)
LIMIT 2;
+--SET spark.sql.cbo.enabled=true
-- correlated IN subquery
-- LIMIT on both parent and subquery sides
SELECT *
diff --git a/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-order-by.sql b/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-order-by.sql
index 7fbb1c12924..a612d40813a 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-order-by.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-order-by.sql
@@ -1,5 +1,11 @@
-- A test suite for ORDER BY in parent side, subquery, and both predicate subquery
-- It includes correlated cases.
+
+-- Comet: the `... and t1c IN (SELECT t2c FROM t2 ORDER BY t2b DESC nulls last) ORDER BY t1c DESC
+-- nulls last` query has ties on the sort key (two rows with t1c = 12), so the relative order of
+-- the tied rows is implementation defined and Comet's sort does not reproduce the order recorded
+-- in the golden file.
+--SET spark.comet.enabled = false
--SET spark.sql.autoBroadcastJoinThreshold=-1
-- Test sort operator with codegen on and off.
diff --git a/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-set-operations.sql b/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-set-operations.sql
index 3df495e0524..82a95c278bb 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-set-operations.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-set-operations.sql
@@ -1,3 +1,7 @@
+-- Comet: ORDER BY ties can produce a different row order from the Spark golden files.
+-- https://github.com/apache/datafusion-comet/issues/5570
+--SET spark.comet.enabled=false
+
-- A test suite for set-operations in parent side, subquery, and both predicate subquery
-- It includes correlated cases.
--ONLY_IF spark
diff --git a/sql/core/src/test/resources/sql-tests/inputs/view-schema-binding-config.sql b/sql/core/src/test/resources/sql-tests/inputs/view-schema-binding-config.sql
index e803254ea64..74db78aee38 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/view-schema-binding-config.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/view-schema-binding-config.sql
@@ -1,6 +1,9 @@
-- This test suits check the spark.sql.viewSchemaBindingMode configuration.
-- It can be DISABLED and COMPENSATION
+-- TODO: https://github.com/apache/datafusion-comet/issues/551
+--SET spark.comet.enabled = false
+
-- Verify the default binding is true
SET spark.sql.legacy.viewSchemaBindingMode;
diff --git a/sql/core/src/test/resources/sql-tests/inputs/view-schema-compensation.sql b/sql/core/src/test/resources/sql-tests/inputs/view-schema-compensation.sql
index 21a3ce1e122..f4762ab98f0 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/view-schema-compensation.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/view-schema-compensation.sql
@@ -1,5 +1,9 @@
-- This test suite checks the WITH SCHEMA COMPENSATION clause
-- Disable ANSI mode to ensure we are forcing it explicitly in the CASTS
+
+-- TODO: https://github.com/apache/datafusion-comet/issues/551
+--SET spark.comet.enabled = false
+
SET spark.sql.ansi.enabled = false;
-- In COMPENSATION views get invalidated if the type can't cast
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
index 085dbcd8046..3090d321b6c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
@@ -49,7 +49,7 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEProp
import org.apache.spark.sql.execution.columnar._
import org.apache.spark.sql.execution.command.CommandUtils
import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
-import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
@@ -564,7 +564,8 @@ class CachedTableSuite extends SharedSparkSession
df.collect()
}
assert(
- collect(df.queryExecution.executedPlan) { case e: ShuffleExchangeExec => e }.size == expected)
+ collect(df.queryExecution.executedPlan) {
+ case _: ShuffleExchangeLike => 1 }.size == expected)
}
test("A cached table preserves the partitioning and ordering of its cached SparkPlan") {
@@ -1703,9 +1704,18 @@ class CachedTableSuite extends SharedSparkSession
_.nodeName.contains("TableCacheQueryStage"))
val aqeNode = findNodeInSparkPlanInfo(inMemoryScanNode.get,
_.nodeName.contains("AdaptiveSparkPlan"))
- val aqePlanRoot = findNodeInSparkPlanInfo(inMemoryScanNode.get,
- _.nodeName.contains("ResultQueryStage"))
- aqePlanRoot.get.children.head.nodeName == "AQEShuffleRead"
+ // Spark 4.0 wraps results in ResultQueryStage. The coalescing indicator is AQEShuffleRead
+ // as the direct child of InputAdapter.
+ // AdaptiveSparkPlan -> ResultQueryStage -> WholestageCodegen ->
+ // CometColumnarToRow -> InputAdapter -> AQEShuffleRead (if coalesced)
+ val resultStage = aqeNode.get.children.head // ResultQueryStage
+ val wsc = resultStage.children.head // WholeStageCodegen
+ val c2r = wsc.children.head // ColumnarToRow or CometColumnarToRow
+ val inputAdapter = c2r.children.head // InputAdapter
+ resultStage.nodeName == "ResultQueryStage" &&
+ wsc.nodeName.startsWith("WholeStageCodegen") && // could be "WholeStageCodegen (1)"
+ (c2r.nodeName == "CometColumnarToRow" || c2r.nodeName == "ColumnarToRow") &&
+ inputAdapter.children.head.nodeName == "AQEShuffleRead"
}
withTempView("t0", "t1", "t2") {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
index 5b8154d2900..eb3de47f387 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
@@ -30,8 +30,9 @@ import org.apache.spark.sql.catalyst.util.AUTO_GENERATED_ALIAS
import org.apache.spark.sql.errors.DataTypeErrors.toSQLId
import org.apache.spark.sql.execution.WholeStageCodegenExec
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.comet.CometHashAggregateExec
import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
-import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
@@ -667,7 +668,9 @@ class DataFrameAggregateSuite extends SharedSparkSession
df.selectExpr("sort_array(collect_set(b) RESPECT NULLS)"), Seq(Row(Seq(null, 2))))
}
- test("SPARK-57298: collect_set normalizes NaN and -0.0 for floating-point types") {
+ test("SPARK-57298: collect_set normalizes NaN and -0.0 for floating-point types",
+ // Comet's native collect_set does not yet normalize NaN / -0.0.
+ IgnoreComet("https://github.com/apache/datafusion-comet/issues/4966")) {
checkAnswer(
sql("SELECT collect_set(v) FROM VALUES (double('NaN')), (double('NaN')) AS t(v)"),
Row(Seq(Double.NaN)))
@@ -932,7 +935,9 @@ class DataFrameAggregateSuite extends SharedSparkSession
case _ => false
}.isDefined)
} else {
- assert(stripAQEPlan(hashAggPlan).isInstanceOf[HashAggregateExec])
+ val strippedPlan = stripAQEPlan(hashAggPlan)
+ assert(strippedPlan.isInstanceOf[HashAggregateExec] ||
+ strippedPlan.exists(_.isInstanceOf[CometHashAggregateExec]))
}
// test case for ObjectHashAggregate and SortAggregate
@@ -941,7 +946,8 @@ class DataFrameAggregateSuite extends SharedSparkSession
val objHashAggOrSortAggPlan =
stripAQEPlan(objHashAggOrSortAggDF.queryExecution.executedPlan)
if (useObjectHashAgg) {
- assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec])
+ assert(objHashAggOrSortAggPlan.isInstanceOf[ObjectHashAggregateExec] ||
+ objHashAggOrSortAggPlan.isInstanceOf[CometHashAggregateExec])
} else {
assert(objHashAggOrSortAggPlan.isInstanceOf[SortAggregateExec])
}
@@ -991,12 +997,12 @@ class DataFrameAggregateSuite extends SharedSparkSession
assert(sortAggPlans.isEmpty)
val objHashAggPlans = collect(aggPlan) {
- case objHashAgg: ObjectHashAggregateExec => objHashAgg
+ case objHashAgg @ (_: ObjectHashAggregateExec | _: CometHashAggregateExec) => objHashAgg
}
assert(objHashAggPlans.nonEmpty)
val exchangePlans = collect(aggPlan) {
- case shuffle: ShuffleExchangeExec => shuffle
+ case shuffle: ShuffleExchangeLike => shuffle
}
assert(exchangePlans.length == 1)
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala
index 9733d51a91c..395a108abc8 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala
@@ -434,7 +434,9 @@ class DataFrameJoinSuite extends SharedSparkSession
withTempDatabase { dbName =>
withTable(table1Name, table2Name) {
- withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ "spark.comet.enabled" -> "false") {
spark.range(50).write.saveAsTable(s"$dbName.$table1Name")
spark.range(100).write.saveAsTable(s"$dbName.$table2Name")
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala
index d838ba4c234..cb0573d56d0 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala
@@ -23,10 +23,11 @@ import java.util.Locale
import org.apache.spark.sql.catalyst.optimizer.RemoveNoopUnion
import org.apache.spark.sql.catalyst.plans.logical.Union
import org.apache.spark.sql.catalyst.plans.physical.UnknownPartitioning
+import org.apache.spark.sql.comet.CometUnionExec
import org.apache.spark.sql.execution.{SparkPlan, UnionExec}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
-import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeLike}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.{ExamplePoint, ExamplePointUDT, SharedSparkSession, SQLTestData}
@@ -1518,11 +1519,12 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP
val union = df1.repartition($"a").union(df2.repartition($"a"))
val unionExec = union.queryExecution.executedPlan.collect {
case u: UnionExec => u
+ case u: CometUnionExec => u
}
assert(unionExec.size == 1)
val shuffle = df1.repartition($"a").queryExecution.executedPlan.collect {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
}
assert(shuffle.size == 1)
@@ -1553,11 +1555,12 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP
val union = df1.repartition($"a").union(df2.repartition($"d"))
val unionExec = union.queryExecution.executedPlan.collect {
case u: UnionExec => u
+ case u: CometUnionExec => u
}
assert(unionExec.size == 1)
val shuffle = df1.repartition($"a").queryExecution.executedPlan.collect {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
}
assert(shuffle.size == 1)
@@ -1572,10 +1575,10 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP
// Avoid unnecessary shuffle if union output partitioning is enabled
val shuffledUnion = union.repartition($"a")
val shuffleNumBefore = union.queryExecution.executedPlan.collect {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
}
val shuffleNumAfter = shuffledUnion.queryExecution.executedPlan.collect {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
}
if (enabled) {
@@ -1604,6 +1607,7 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP
val union = df1.repartitionByRange($"a").union(df2.repartitionByRange($"d"))
val unionExec = union.queryExecution.executedPlan.collect {
case u: UnionExec => u
+ case u: CometUnionExec => u
}
assert(unionExec.size == 1)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
index 577ae025f1d..a0c95263931 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
@@ -36,11 +36,12 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference,
import org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation
import org.apache.spark.sql.catalyst.parser.ParseException
import org.apache.spark.sql.catalyst.plans.logical.{Filter, LeafNode, LocalRelation, LogicalPlan, OneRowRelation}
+import org.apache.spark.sql.comet.CometBroadcastExchangeExec
import org.apache.spark.sql.connector.FakeV2Provider
import org.apache.spark.sql.execution.{FilterExec, LogicalRDD, QueryExecution, SortExec, WholeStageCodegenExec}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.aggregate.HashAggregateExec
-import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike}
+import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec, ShuffleExchangeLike}
import org.apache.spark.sql.expressions.{Aggregator, Window}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
@@ -1492,7 +1493,7 @@ class DataFrameSuite extends SharedSparkSession
fail("Should not have back to back Aggregates")
}
atFirstAgg = true
- case e: ShuffleExchangeExec => atFirstAgg = false
+ case e: ShuffleExchangeLike => atFirstAgg = false
case _ =>
}
}
@@ -1682,7 +1683,7 @@ class DataFrameSuite extends SharedSparkSession
checkAnswer(join, df)
assert(
collect(join.queryExecution.executedPlan) {
- case e: ShuffleExchangeExec => true }.size === 1)
+ case _: ShuffleExchangeLike => true }.size === 1)
assert(
collect(join.queryExecution.executedPlan) { case e: ReusedExchangeExec => true }.size === 1)
val broadcasted = broadcast(join)
@@ -1690,10 +1691,12 @@ class DataFrameSuite extends SharedSparkSession
checkAnswer(join2, df)
assert(
collect(join2.queryExecution.executedPlan) {
- case e: ShuffleExchangeExec => true }.size == 1)
+ case _: ShuffleExchangeLike => true }.size == 1)
assert(
collect(join2.queryExecution.executedPlan) {
- case e: BroadcastExchangeExec => true }.size === 1)
+ case e: BroadcastExchangeExec => true
+ case _: CometBroadcastExchangeExec => true
+ }.size === 1)
assert(
collect(join2.queryExecution.executedPlan) { case e: ReusedExchangeExec => true }.size == 4)
}
@@ -2091,7 +2094,7 @@ class DataFrameSuite extends SharedSparkSession
// Assert that no extra shuffle introduced by cogroup.
val exchanges = collect(df3.queryExecution.executedPlan) {
- case h: ShuffleExchangeExec => h
+ case h: ShuffleExchangeLike => h
}
assert(exchanges.size == 2)
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala
index f79824de8ff..5432984960f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala
@@ -25,8 +25,10 @@ import org.apache.spark.sql.catalyst.optimizer.TransposeWindow
import org.apache.spark.sql.catalyst.plans.logical.{Window => LogicalWindow}
import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning
import org.apache.spark.sql.catalyst.trees.UnaryLike
+import org.apache.spark.sql.comet.CometWindowExec
+import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
-import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, Exchange, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, Exchange, ShuffleExchangeExec, ShuffleExchangeLike}
import org.apache.spark.sql.execution.window.WindowExec
import org.apache.spark.sql.expressions.{Aggregator, MutableAggregationBuffer, UserDefinedAggregateFunction, Window}
import org.apache.spark.sql.functions._
@@ -975,7 +977,8 @@ class DataFrameWindowFunctionsSuite extends SharedSparkSession
}
}
- test("Window spill with more than the inMemoryThreshold and spillThreshold") {
+ test("Window spill with more than the inMemoryThreshold and spillThreshold",
+ IgnoreComet("Comet does not support spilling")) {
val df = Seq((1, "1"), (2, "2"), (1, "3"), (2, "4")).toDF("key", "value")
val window = Window.partitionBy($"key").orderBy($"value")
@@ -987,7 +990,7 @@ class DataFrameWindowFunctionsSuite extends SharedSparkSession
}
}
- test("SPARK-21258: complex object in combination with spilling") {
+ test("SPARK-21258: complex object in combination with spilling", IgnoreComet("Comet does not support spilling")) {
// Make sure we trigger the spilling path.
withSQLConf(SQLConf.WINDOW_EXEC_BUFFER_IN_MEMORY_THRESHOLD.key -> "1",
SQLConf.WINDOW_EXEC_BUFFER_SPILL_THRESHOLD.key -> "17") {
@@ -1336,10 +1339,12 @@ class DataFrameWindowFunctionsSuite extends SharedSparkSession
}
def isShuffleExecByRequirement(
- plan: ShuffleExchangeExec,
+ plan: ShuffleExchangeLike,
desiredClusterColumns: Seq[String]): Boolean = plan match {
case ShuffleExchangeExec(op: HashPartitioning, _, ENSURE_REQUIREMENTS, _) =>
partitionExpressionsColumns(op.expressions) === desiredClusterColumns
+ case CometShuffleExchangeExec(op: HashPartitioning, _, _, ENSURE_REQUIREMENTS, _, _) =>
+ partitionExpressionsColumns(op.expressions) === desiredClusterColumns
case _ => false
}
@@ -1362,7 +1367,12 @@ class DataFrameWindowFunctionsSuite extends SharedSparkSession
val shuffleByRequirement = windowed.queryExecution.executedPlan.exists {
case w: WindowExec =>
w.child.exists {
- case s: ShuffleExchangeExec => isShuffleExecByRequirement(s, Seq("key1", "key2"))
+ case s: ShuffleExchangeLike => isShuffleExecByRequirement(s, Seq("key1", "key2"))
+ case _ => false
+ }
+ case w: CometWindowExec =>
+ w.child.exists {
+ case s: ShuffleExchangeLike => isShuffleExecByRequirement(s, Seq("key1", "key2"))
case _ => false
}
case _ => false
@@ -1814,7 +1824,8 @@ class DataFrameWindowFunctionsSuite extends SharedSparkSession
}
}
- test("SPARK-49386: Window spill with more than the inMemoryThreshold and spillSizeThreshold") {
+ test("SPARK-49386: Window spill with more than the inMemoryThreshold and spillSizeThreshold",
+ IgnoreComet("Comet does not support spilling")) {
val df = Seq((1, "1"), (2, "2"), (1, "3"), (2, "4")).toDF("key", "value")
val window = Window.partitionBy($"key").orderBy($"value")
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index 879569045b6..f3ff89067d2 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -46,7 +46,7 @@ import org.apache.spark.sql.catalyst.trees.DataFrameQueryContext
import org.apache.spark.sql.catalyst.util.sideBySide
import org.apache.spark.sql.execution.{LogicalRDD, RDDScanExec, SQLExecution}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
-import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike}
import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.functions._
@@ -2521,7 +2521,7 @@ class DatasetSuite extends SharedSparkSession
// Assert that no extra shuffle introduced by cogroup.
val exchanges = collect(df3.queryExecution.executedPlan) {
- case h: ShuffleExchangeExec => h
+ case h: ShuffleExchangeLike => h
}
assert(exchanges.size == 2)
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
index c68f64d52f2..32489f2cb15 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
@@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen
import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression}
import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._
import org.apache.spark.sql.catalyst.plans.ExistenceJoin
+import org.apache.spark.sql.comet.{CometNativeScanExec, CometScanExec, CometSubqueryBroadcastExec}
import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog}
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.adaptive._
@@ -192,6 +193,7 @@ abstract class DynamicPartitionPruningSuiteBase
}
val subqueryBroadcast = dpExprs.collect {
case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b
+ case InSubqueryExec(_, b: CometSubqueryBroadcastExec, _, _, _, _) => b
}
val hasFilter = if (withSubquery) "Should" else "Shouldn't"
@@ -246,6 +248,8 @@ abstract class DynamicPartitionPruningSuiteBase
val buf = collectDynamicPruningExpressions(df.queryExecution.executedPlan).collect {
case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) =>
b.indices.map(idx => b.buildKeys(idx))
+ case InSubqueryExec(_, b: CometSubqueryBroadcastExec, _, _, _, _) =>
+ b.indices.map(idx => b.buildKeys(idx))
}
assert(buf.distinct.size == n)
}
@@ -261,6 +265,12 @@ abstract class DynamicPartitionPruningSuiteBase
case s: BatchScanExec => s.runtimeFilters.collect {
case d: DynamicPruningExpression => d.child
}
+ case s: CometScanExec => s.partitionFilters.collect {
+ case d: DynamicPruningExpression => d.child
+ }
+ case s: CometNativeScanExec => s.partitionFilters.collect {
+ case d: DynamicPruningExpression => d.child
+ }
case _ => Nil
}
}
@@ -1203,10 +1213,16 @@ abstract class DynamicPartitionPruningSuiteBase
val plan = df.queryExecution.executedPlan
val countSubqueryBroadcasts =
- collectWithSubqueries(plan)({ case _: SubqueryBroadcastExec => 1 }).sum
+ collectWithSubqueries(plan)({
+ case _: SubqueryBroadcastExec => 1
+ case _: CometSubqueryBroadcastExec => 1
+ }).sum
val countReusedSubqueryBroadcasts =
- collectWithSubqueries(plan)({ case ReusedSubqueryExec(_: SubqueryBroadcastExec) => 1}).sum
+ collectWithSubqueries(plan)({
+ case ReusedSubqueryExec(_: SubqueryBroadcastExec) => 1
+ case ReusedSubqueryExec(_: CometSubqueryBroadcastExec) => 1
+ }).sum
assert(countSubqueryBroadcasts == 1)
assert(countReusedSubqueryBroadcasts == 1)
@@ -1578,6 +1594,7 @@ abstract class DynamicPartitionPruningSuiteBase
val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) {
case s: SubqueryBroadcastExec => s
+ case s: CometSubqueryBroadcastExec => s
}
assert(subqueryBroadcastExecs.size === 1)
subqueryBroadcastExecs.foreach { subqueryBroadcastExec =>
@@ -1730,6 +1747,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat
case s: BatchScanExec =>
// we use f1 col for v2 tables due to schema pruning
s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1")))
+ case s: CometScanExec =>
+ s.output.exists(_.exists(_.argString(maxFields = 100).contains("fid")))
+ case s: CometNativeScanExec =>
+ s.output.exists(_.exists(_.argString(maxFields = 100).contains("fid")))
case _ => false
}
assert(scanOption.isDefined)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala
index 618f0ab675e..efcd4853251 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala
@@ -268,7 +268,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite
}
}
- test("SPARK-33853: explain codegen - check presence of subquery") {
+ test("SPARK-33853: explain codegen - check presence of subquery",
+ IgnoreComet("Comet plan has a different WholeStageCodegen subtree count")) {
withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") {
withTempView("df") {
val df1 = spark.range(1, 100)
@@ -339,7 +340,7 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite
val sqlText = "EXPLAIN FORMATTED SELECT (SELECT min(id) FROM df) as v"
val expected_pattern =
- "Subquery:1 Hosting operator id = 2 Hosting Expression = Subquery subquery#x"
+ "Subquery:1 Hosting operator id = \\d+ Hosting Expression = Subquery subquery#x"
withNormalizedExplain(sqlText) { normalizedOutput =>
assert(expected_pattern.r.findAllMatchIn(normalizedOutput).length == 1)
@@ -471,7 +472,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite
}
}
- test("Explain formatted output for scan operator for datasource V2") {
+ test("Explain formatted output for scan operator for datasource V2",
+ IgnoreComet("Comet explain output is different")) {
withTempDir { dir =>
Seq("parquet", "orc", "csv", "json").foreach { fmt =>
val basePath = dir.getCanonicalPath + "/" + fmt
@@ -549,7 +551,9 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite
}
}
-class ExplainSuiteAE extends ExplainSuiteHelper with EnableAdaptiveExecutionSuite {
+// Ignored when Comet is enabled. Comet changes expected query plans.
+class ExplainSuiteAE extends ExplainSuiteHelper with EnableAdaptiveExecutionSuite
+ with IgnoreCometSuite {
import testImplicits._
override protected def sparkConf =
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/FileBasedDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/FileBasedDataSourceSuite.scala
index 1fc45e9703f..b54410841af 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/FileBasedDataSourceSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/FileBasedDataSourceSuite.scala
@@ -33,6 +33,7 @@ import org.apache.spark.sql.catalyst.expressions.{AttributeReference, GreaterTha
import org.apache.spark.sql.catalyst.expressions.IntegralLiteralTestUtils.{negativeInt, positiveInt}
import org.apache.spark.sql.catalyst.plans.logical.Filter
import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeScanExec, CometScanExec, CometSortMergeJoinExec}
import org.apache.spark.sql.execution.{FileSourceScanLike, SimpleMode}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.datasources.FilePartition
@@ -654,18 +655,25 @@ class FileBasedDataSourceSuite extends SharedSparkSession
checkAnswer(sql(s"select A from $tableName"), data.select("A"))
// RuntimeException is triggered at executor side, which is then wrapped as
- // SparkException at driver side
+ // SparkException at driver side. Comet native readers throw
+ // SparkRuntimeException directly without the SparkException wrapper.
+ def getDuplicateFieldError(query: String): SparkRuntimeException = {
+ try {
+ sql(query).collect()
+ fail("Expected an exception").asInstanceOf[SparkRuntimeException]
+ } catch {
+ case e: SparkException =>
+ e.getCause.asInstanceOf[SparkRuntimeException]
+ case e: SparkRuntimeException => e
+ }
+ }
checkError(
- exception = intercept[SparkException] {
- sql(s"select b from $tableName").collect()
- }.getCause.asInstanceOf[SparkRuntimeException],
+ exception = getDuplicateFieldError(s"select b from $tableName"),
condition = "_LEGACY_ERROR_TEMP_2093",
parameters = Map("requiredFieldName" -> "b", "matchedOrcFields" -> "[b, B]")
)
checkError(
- exception = intercept[SparkException] {
- sql(s"select B from $tableName").collect()
- }.getCause.asInstanceOf[SparkRuntimeException],
+ exception = getDuplicateFieldError(s"select B from $tableName"),
condition = "_LEGACY_ERROR_TEMP_2093",
parameters = Map("requiredFieldName" -> "b", "matchedOrcFields" -> "[b, B]")
)
@@ -976,6 +984,7 @@ class FileBasedDataSourceSuite extends SharedSparkSession
assert(bJoinExec.isEmpty)
val smJoinExec = collect(joinedDF.queryExecution.executedPlan) {
case smJoin: SortMergeJoinExec => smJoin
+ case smJoin: CometSortMergeJoinExec => smJoin
}
assert(smJoinExec.nonEmpty)
}
@@ -1036,6 +1045,7 @@ class FileBasedDataSourceSuite extends SharedSparkSession
val fileScan = df.queryExecution.executedPlan collectFirst {
case BatchScanExec(_, f: FileScan, _, _, _, _) => f
+ case CometBatchScanExec(BatchScanExec(_, f: FileScan, _, _, _, _), _, _) => f
}
assert(fileScan.nonEmpty)
assert(fileScan.get.partitionFilters.nonEmpty)
@@ -1077,6 +1087,7 @@ class FileBasedDataSourceSuite extends SharedSparkSession
val fileScan = df.queryExecution.executedPlan collectFirst {
case BatchScanExec(_, f: FileScan, _, _, _, _) => f
+ case CometBatchScanExec(BatchScanExec(_, f: FileScan, _, _, _, _), _, _) => f
}
assert(fileScan.nonEmpty)
assert(fileScan.get.partitionFilters.isEmpty)
@@ -1261,6 +1272,9 @@ class FileBasedDataSourceSuite extends SharedSparkSession
val filters = df.queryExecution.executedPlan.collect {
case f: FileSourceScanLike => f.dataFilters
case b: BatchScanExec => b.scan.asInstanceOf[FileScan].dataFilters
+ case b: CometScanExec => b.dataFilters
+ case b: CometNativeScanExec => b.dataFilters
+ case b: CometBatchScanExec => b.scan.asInstanceOf[FileScan].dataFilters
}.flatten
assert(filters.contains(GreaterThan(scan.logicalPlan.output.head, Literal(5L))))
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala
new file mode 100644
index 00000000000..4b31bea33de
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/IgnoreComet.scala
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql
+
+import org.scalactic.source.Position
+import org.scalatest.Tag
+
+import org.apache.spark.sql.test.SQLTestUtils
+
+/**
+ * Tests with this tag will be ignored when Comet is enabled (e.g., via `ENABLE_COMET`).
+ */
+case class IgnoreComet(reason: String) extends Tag("DisableComet")
+
+/**
+ * Helper trait that disables Comet for all tests regardless of default config values.
+ */
+trait IgnoreCometSuite extends SQLTestUtils {
+ override protected def test(testName: String, testTags: Tag*)(testFun: => Any)
+ (implicit pos: Position): Unit = {
+ if (isCometEnabled) {
+ ignore(testName + " (disabled when Comet is on)", testTags: _*)(testFun)
+ } else {
+ super.test(testName, testTags: _*)(testFun)
+ }
+ }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala
index df9f1d1ed2d..ad3dd43b652 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/JoinHintSuite.scala
@@ -22,6 +22,7 @@ import org.apache.logging.log4j.Level
import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide, EliminateResolvedHint}
import org.apache.spark.sql.catalyst.plans.logical._
import org.apache.spark.sql.catalyst.rules.RuleExecutor
+import org.apache.spark.sql.comet.{CometHashJoinExec, CometSortMergeJoinExec}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.joins._
import org.apache.spark.sql.internal.SQLConf
@@ -361,6 +362,7 @@ class JoinHintSuite extends SharedSparkSession with AdaptiveSparkPlanHelper {
val executedPlan = df.queryExecution.executedPlan
val shuffleHashJoins = collect(executedPlan) {
case s: ShuffledHashJoinExec => s
+ case c: CometHashJoinExec => c.originalPlan.asInstanceOf[ShuffledHashJoinExec]
}
assert(shuffleHashJoins.size == 1)
assert(shuffleHashJoins.head.buildSide == buildSide)
@@ -370,6 +372,7 @@ class JoinHintSuite extends SharedSparkSession with AdaptiveSparkPlanHelper {
val executedPlan = df.queryExecution.executedPlan
val shuffleMergeJoins = collect(executedPlan) {
case s: SortMergeJoinExec => s
+ case c: CometSortMergeJoinExec => c
}
assert(shuffleMergeJoins.size == 1)
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala
index 7f695e90df8..57fea6f8324 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala
@@ -29,7 +29,8 @@ import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation
import org.apache.spark.sql.catalyst.expressions.{Ascending, GenericRow, SortOrder}
import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, JoinSelectionHelper}
import org.apache.spark.sql.catalyst.plans.logical.{Filter, HintInfo, Join, JoinHint, NO_BROADCAST_AND_REPLICATION}
-import org.apache.spark.sql.execution.{BinaryExecNode, FilterExec, ProjectExec, SortExec, SparkPlan, WholeStageCodegenExec}
+import org.apache.spark.sql.comet._
+import org.apache.spark.sql.execution.{BinaryExecNode, ColumnarToRowExec, FilterExec, InputAdapter, ProjectExec, SortExec, SparkPlan, WholeStageCodegenExec}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec, ShuffleExchangeLike}
import org.apache.spark.sql.execution.joins._
@@ -808,7 +809,8 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
}
}
- test("test SortMergeJoin (with spill)") {
+ test("test SortMergeJoin (with spill)",
+ IgnoreComet("TODO: Comet SMJ doesn't support spill yet")) {
withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1",
SQLConf.SORT_MERGE_JOIN_EXEC_BUFFER_IN_MEMORY_THRESHOLD.key -> "0",
SQLConf.SORT_MERGE_JOIN_EXEC_BUFFER_SPILL_THRESHOLD.key -> "1") {
@@ -816,7 +818,8 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
}
}
- test("SPARK-49386: test SortMergeJoin (with spill by size threshold)") {
+ test("SPARK-49386: test SortMergeJoin (with spill by size threshold)",
+ IgnoreComet("TODO: Comet SMJ doesn't support spill yet")) {
withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1",
SQLConf.SORT_MERGE_JOIN_EXEC_BUFFER_IN_MEMORY_THRESHOLD.key -> "0",
SQLConf.SORT_MERGE_JOIN_EXEC_BUFFER_SPILL_THRESHOLD.key -> Int.MaxValue.toString,
@@ -946,10 +949,12 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
val physical = df.queryExecution.sparkPlan
val physicalJoins = physical.collect {
case j: SortMergeJoinExec => j
+ case j: CometSortMergeJoinExec => j.originalPlan.asInstanceOf[SortMergeJoinExec]
}
val executed = df.queryExecution.executedPlan
val executedJoins = collect(executed) {
case j: SortMergeJoinExec => j
+ case j: CometSortMergeJoinExec => j.originalPlan.asInstanceOf[SortMergeJoinExec]
}
// This only applies to the above tested queries, in which a child SortMergeJoin always
// contains the SortOrder required by its parent SortMergeJoin. Thus, SortExec should never
@@ -1195,9 +1200,11 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
val plan = df1.join(df2.hint("SHUFFLE_HASH"), $"k1" === $"k2", joinType)
.groupBy($"k1").count()
.queryExecution.executedPlan
- assert(collect(plan) { case _: ShuffledHashJoinExec => true }.size === 1)
+ assert(collect(plan) {
+ case _: ShuffledHashJoinExec | _: CometHashJoinExec => true }.size === 1)
// No extra shuffle before aggregate
- assert(collect(plan) { case _: ShuffleExchangeExec => true }.size === 2)
+ assert(collect(plan) {
+ case _: ShuffleExchangeLike => true }.size === 2)
})
}
@@ -1214,10 +1221,11 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
.join(df4.hint("SHUFFLE_MERGE"), $"k1" === $"k4", joinType)
.queryExecution
.executedPlan
- assert(collect(plan) { case _: SortMergeJoinExec => true }.size === 2)
+ assert(collect(plan) {
+ case _: SortMergeJoinExec | _: CometSortMergeJoinExec => true }.size === 2)
assert(collect(plan) { case _: BroadcastHashJoinExec => true }.size === 1)
// No extra sort before last sort merge join
- assert(collect(plan) { case _: SortExec => true }.size === 3)
+ assert(collect(plan) { case _: SortExec | _: CometSortExec => true }.size === 3)
})
// Test shuffled hash join
@@ -1227,10 +1235,13 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
.join(df4.hint("SHUFFLE_MERGE"), $"k1" === $"k4", joinType)
.queryExecution
.executedPlan
- assert(collect(plan) { case _: SortMergeJoinExec => true }.size === 2)
- assert(collect(plan) { case _: ShuffledHashJoinExec => true }.size === 1)
+ assert(collect(plan) {
+ case _: SortMergeJoinExec | _: CometSortMergeJoinExec => true }.size === 2)
+ assert(collect(plan) {
+ case _: ShuffledHashJoinExec | _: CometHashJoinExec => true }.size === 1)
// No extra sort before last sort merge join
- assert(collect(plan) { case _: SortExec => true }.size === 3)
+ assert(collect(plan) {
+ case _: SortExec | _: CometSortExec => true }.size === 3)
})
}
@@ -1349,12 +1360,12 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
inputDFs.foreach { case (df1, df2, joinExprs) =>
val smjDF = df1.join(df2.hint("SHUFFLE_MERGE"), joinExprs, "full")
assert(collect(smjDF.queryExecution.executedPlan) {
- case _: SortMergeJoinExec => true }.size === 1)
+ case _: SortMergeJoinExec | _: CometSortMergeJoinExec => true }.size === 1)
val smjResult = smjDF.collect()
val shjDF = df1.join(df2.hint("SHUFFLE_HASH"), joinExprs, "full")
assert(collect(shjDF.queryExecution.executedPlan) {
- case _: ShuffledHashJoinExec => true }.size === 1)
+ case _: ShuffledHashJoinExec | _: CometHashJoinExec => true }.size === 1)
// Same result between shuffled hash join and sort merge join
checkAnswer(shjDF, smjResult)
}
@@ -1413,12 +1424,14 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
val smjDF = df1.hint("SHUFFLE_MERGE").join(df2, joinExprs, "leftouter")
assert(collect(smjDF.queryExecution.executedPlan) {
case _: SortMergeJoinExec => true
+ case _: CometSortMergeJoinExec => true
}.size === 1)
val smjResult = smjDF.collect()
val shjDF = df1.hint("SHUFFLE_HASH").join(df2, joinExprs, "leftouter")
assert(collect(shjDF.queryExecution.executedPlan) {
case _: ShuffledHashJoinExec => true
+ case _: CometHashJoinExec => true
}.size === 1)
// Same result between shuffled hash join and sort merge join
checkAnswer(shjDF, smjResult)
@@ -1429,12 +1442,14 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
val smjDF = df2.join(df1.hint("SHUFFLE_MERGE"), joinExprs, "rightouter")
assert(collect(smjDF.queryExecution.executedPlan) {
case _: SortMergeJoinExec => true
+ case _: CometSortMergeJoinExec => true
}.size === 1)
val smjResult = smjDF.collect()
val shjDF = df2.join(df1.hint("SHUFFLE_HASH"), joinExprs, "rightouter")
assert(collect(shjDF.queryExecution.executedPlan) {
case _: ShuffledHashJoinExec => true
+ case _: CometHashJoinExec => true
}.size === 1)
// Same result between shuffled hash join and sort merge join
checkAnswer(shjDF, smjResult)
@@ -1478,13 +1493,20 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
assert(shjCodegenDF.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_ : ShuffledHashJoinExec) => true
case WholeStageCodegenExec(ProjectExec(_, _ : ShuffledHashJoinExec)) => true
+ case WholeStageCodegenExec(ColumnarToRowExec(InputAdapter(_: CometHashJoinExec))) =>
+ true
+ case WholeStageCodegenExec(ColumnarToRowExec(
+ InputAdapter(CometProjectExec(_, _, _, _, _: CometHashJoinExec, _)))) => true
+ case _: CometHashJoinExec => true
}.size === 1)
checkAnswer(shjCodegenDF, Seq.empty)
withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false") {
val shjNonCodegenDF = df1.join(df2.hint("SHUFFLE_HASH"), $"k1" === $"k2", joinType)
assert(shjNonCodegenDF.queryExecution.executedPlan.collect {
- case _: ShuffledHashJoinExec => true }.size === 1)
+ case _: ShuffledHashJoinExec => true
+ case _: CometHashJoinExec => true
+ }.size === 1)
checkAnswer(shjNonCodegenDF, Seq.empty)
}
}
@@ -1516,7 +1538,10 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
"/*+ BROADCAST(t2) */ t1.k as k"
}
val plan = sql(getAggQuery(selectExpr, joinType)).queryExecution.executedPlan
- assert(collect(plan) { case _: BroadcastNestedLoopJoinExec => true }.size === 1)
+ assert(collect(plan) {
+ case _: BroadcastNestedLoopJoinExec | _: CometBroadcastNestedLoopJoinExec =>
+ true
+ }.size === 1)
// No extra shuffle before aggregation
assert(collect(plan) { case _: ShuffleExchangeExec => true }.size === 0)
}
@@ -1532,7 +1557,8 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
val plan = sql(getAggQuery(selectExpr, joinType)).queryExecution.executedPlan
assert(collect(plan) { case _: BroadcastNestedLoopJoinExec => true }.size === 1)
// Have shuffle before aggregation
- assert(collect(plan) { case _: ShuffleExchangeExec => true }.size === 1)
+ assert(collect(plan) {
+ case _: ShuffleExchangeLike => true }.size === 1)
}
def getJoinQuery(selectExpr: String, joinType: String): String = {
@@ -1560,10 +1586,16 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
"/*+ BROADCAST(right_t) */ k1 as k0"
}
val plan = sql(getJoinQuery(selectExpr, joinType)).queryExecution.executedPlan
- assert(collect(plan) { case _: BroadcastNestedLoopJoinExec => true }.size === 1)
- assert(collect(plan) { case _: SortMergeJoinExec => true }.size === 3)
+ assert(collect(plan) {
+ case _: BroadcastNestedLoopJoinExec | _: CometBroadcastNestedLoopJoinExec =>
+ true
+ }.size === 1)
+ assert(collect(plan) {
+ case _: SortMergeJoinExec => true
+ case _: CometSortMergeJoinExec => true
+ }.size === 3)
// No extra sort on left side before last sort merge join
- assert(collect(plan) { case _: SortExec => true }.size === 5)
+ assert(collect(plan) { case _: SortExec | _: CometSortExec => true }.size === 5)
}
// Test output ordering is not preserved
@@ -1572,9 +1604,12 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
val selectExpr = "/*+ BROADCAST(left_t) */ k1 as k0"
val plan = sql(getJoinQuery(selectExpr, joinType)).queryExecution.executedPlan
assert(collect(plan) { case _: BroadcastNestedLoopJoinExec => true }.size === 1)
- assert(collect(plan) { case _: SortMergeJoinExec => true }.size === 3)
+ assert(collect(plan) {
+ case _: SortMergeJoinExec => true
+ case _: CometSortMergeJoinExec => true
+ }.size === 3)
// Have sort on left side before last sort merge join
- assert(collect(plan) { case _: SortExec => true }.size === 6)
+ assert(collect(plan) { case _: SortExec | _: CometSortExec => true }.size === 6)
}
// Test singe partition
@@ -1584,7 +1619,8 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
|FROM range(0, 10, 1, 1) t1 FULL OUTER JOIN range(0, 10, 1, 1) t2
|""".stripMargin)
val plan = fullJoinDF.queryExecution.executedPlan
- assert(collect(plan) { case _: ShuffleExchangeExec => true}.size == 1)
+ assert(collect(plan) {
+ case _: ShuffleExchangeLike => true}.size == 1)
checkAnswer(fullJoinDF, Row(100))
}
}
@@ -1657,6 +1693,9 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
Seq(semiJoinDF, antiJoinDF).foreach { df =>
assert(collect(df.queryExecution.executedPlan) {
case j: ShuffledHashJoinExec if j.ignoreDuplicatedKey == ignoreDuplicatedKey => true
+ case j: CometHashJoinExec
+ if j.originalPlan.asInstanceOf[ShuffledHashJoinExec].ignoreDuplicatedKey ==
+ ignoreDuplicatedKey => true
}.size == 1)
}
}
@@ -1701,14 +1740,20 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper
test("SPARK-43113: Full outer join with duplicate stream-side references in condition (SMJ)") {
def check(plan: SparkPlan): Unit = {
- assert(collect(plan) { case _: SortMergeJoinExec => true }.size === 1)
+ assert(collect(plan) {
+ case _: SortMergeJoinExec => true
+ case _: CometSortMergeJoinExec => true
+ }.size === 1)
}
dupStreamSideColTest("MERGE", check)
}
test("SPARK-43113: Full outer join with duplicate stream-side references in condition (SHJ)") {
def check(plan: SparkPlan): Unit = {
- assert(collect(plan) { case _: ShuffledHashJoinExec => true }.size === 1)
+ assert(collect(plan) {
+ case _: ShuffledHashJoinExec => true
+ case _: CometHashJoinExec => true
+ }.size === 1)
}
dupStreamSideColTest("SHUFFLE_HASH", check)
}
@@ -1847,7 +1892,8 @@ class ThreadLeakInSortMergeJoinSuite
sparkConf.set(SHUFFLE_SPILL_NUM_ELEMENTS_FORCE_SPILL_THRESHOLD, 20))
}
- test("SPARK-47146: thread leak when doing SortMergeJoin (with spill)") {
+ test("SPARK-47146: thread leak when doing SortMergeJoin (with spill)",
+ IgnoreComet("Comet SMJ doesn't spill yet")) {
withSQLConf(
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/MapStatusEndToEndSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/MapStatusEndToEndSuite.scala
index 0708b2c9392..d6ff2e1555c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/MapStatusEndToEndSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/MapStatusEndToEndSuite.scala
@@ -36,7 +36,8 @@ class MapStatusEndToEndSuite extends QueryTest {
SparkSession.clearDefaultSession()
}
- test("Propagate checksum from executor to driver") {
+ test("Propagate checksum from executor to driver",
+ IgnoreComet("https://github.com/apache/datafusion-comet/issues/6414")) {
assert(spark.sparkContext.conf.get(SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key) == "5")
assert(spark.conf.get(SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key) == "5")
assert(spark.sparkContext.conf.get(SQLConf.CLASSIC_SHUFFLE_DEPENDENCY_FILE_CLEANUP_ENABLED.key)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/PlanStabilitySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/PlanStabilitySuite.scala
index 6cd49948630..b6fcb716a2a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/PlanStabilitySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/PlanStabilitySuite.scala
@@ -69,7 +69,7 @@ import org.apache.spark.util.Utils
* }}}
*/
// scalastyle:on line.size.limit
-trait PlanStabilitySuite extends DisableAdaptiveExecutionSuite {
+trait PlanStabilitySuite extends DisableAdaptiveExecutionSuite with IgnoreCometSuite {
protected val baseResourcePath = {
// use the same way as `SQLQueryTestSuite` to get the resource path
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala b/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala
index 291aa7cab72..7783c37683e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala
@@ -42,6 +42,7 @@ import org.apache.spark.sql.catalyst.plans._
import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
import org.apache.spark.sql.catalyst.util._
import org.apache.spark.sql.classic.ClassicConversions._
+import org.apache.spark.sql.comet.{CometFilterExec, CometProjectExec}
import org.apache.spark.sql.execution.{FilterExec, QueryExecution, SparkPlan, SQLExecution}
import org.apache.spark.sql.execution.adaptive.DisableAdaptiveExecution
import org.apache.spark.sql.execution.columnar.InMemoryRelation
@@ -295,6 +296,11 @@ trait QueryTestBase
self.spark.asInstanceOf[classic.SparkSession].converter
}
+ /**
+ * Whether Comet extension is enabled
+ */
+ protected def isCometEnabled: Boolean = classic.SparkSession.isCometEnabled
+
protected override def withSQLConf[T](pairs: (String, String)*)(f: => T): T = {
SparkSession.setActiveSession(spark)
super.withSQLConf(pairs: _*)(f)
@@ -514,6 +520,8 @@ trait QueryTestBase
val schema = df.schema
val withoutFilters = df.queryExecution.executedPlan.transform {
case FilterExec(_, child) => child
+ case CometFilterExec(_, _, _, _, child, _) => child
+ case CometProjectExec(_, _, _, _, CometFilterExec(_, _, _, _, child, _), _) => child
}
spark.asInstanceOf[classic.SparkSession]
@@ -757,6 +765,11 @@ trait QueryTest extends SparkFunSuite with QueryTestBase with PlanTest {
override protected def test(testName: String, testTags: Tag*)(testFun: => Any)
(implicit pos: Position): Unit = {
+ // Check Comet skip tags first, before DisableAdaptiveExecution handling
+ if (isCometEnabled && testTags.exists(_.isInstanceOf[IgnoreComet])) {
+ ignore(testName + " (disabled when Comet is on)", testTags: _*)(testFun)
+ return
+ }
if (testTags.exists(_.isInstanceOf[DisableAdaptiveExecution])) {
super.test(testName, testTags: _*) {
withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
index da6f6aca2ad..12f48c4bb5d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
@@ -1529,7 +1529,8 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper
checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001")))
}
- test("external sorting updates peak execution memory") {
+ test("external sorting updates peak execution memory",
+ IgnoreComet("TODO: native CometSort does not update peak execution memory")) {
AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") {
sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect()
}
@@ -1990,8 +1991,15 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper
countAcc.add(1)
x
})
+ // Comet's `CometProject` and `CometHashAggregate` do not implement Spark's cross-sibling
+ // subexpression elimination over `ScalaUDF`, so each reference invokes the UDF body
+ // separately. The other call sites in this test pass against Comet because the source
+ // (`testData2`, a `LocalRelation`) is not Comet-scannable and the project runs on Spark's
+ // path; the `agg` case routes through `CometHashAggregate` once an Exchange enters the plan.
+ // Tracking issue: https://github.com/apache/datafusion-comet/issues/4516
verifyCallCount(
- df.agg(sum(testUdf($"b") + testUdf($"b") + testUdf($"b"))), Row(3.0), 1)
+ df.agg(sum(testUdf($"b") + testUdf($"b") + testUdf($"b"))), Row(3.0),
+ if (isCometEnabled) 3 else 1)
verifyCallCount(
df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1)
@@ -4755,7 +4763,8 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper
)
}
- test("SPARK-47939: Explain should work with parameterized queries") {
+ test("SPARK-47939: Explain should work with parameterized queries",
+ IgnoreComet("Comet explain output is different")) {
def checkQueryPlan(df: DataFrame, plan: String): Unit = assert(
df.collect()
.map(_.getString(0))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala
index 395cb67f441..33ac6ed19af 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala
@@ -190,6 +190,33 @@ class SQLQueryTestSuite extends SharedSparkSession with SQLHelper
protected def ignoreList: Set[String] = Set(
"ignored.sql" // Do NOT remove this one. It is here to test the ignore functionality.
) ++ otherIgnoreList
+
+ // Comet: hll.sql, thetasketch.sql and tuplesketch.sql each create one of these tables and fill
+ // it with invalid UTF-8, which Spark allows in a STRING column and Comet's native scan rejects.
+ // Only the queries that name one of them run with the native scan disabled, so every other
+ // query in those files still runs with Comet.
+ // TODO: https://github.com/apache/datafusion-comet/issues/4121
+ private val cometInvalidUtf8Tables = Seq("hll_string_test", "t_string_collation")
+
+ override protected def getNormalizedQueryExecutionResult(
+ session: SparkSession, sql: String): (String, Seq[String]) = {
+ if (isCometEnabled && cometInvalidUtf8Tables.exists(sql.contains)) {
+ val key = "spark.comet.scan.enabled"
+ val previous = session.conf.getOption(key)
+ session.conf.set(key, "false")
+ try {
+ super.getNormalizedQueryExecutionResult(session, sql)
+ } finally {
+ previous match {
+ case Some(value) => session.conf.set(key, value)
+ case None => session.conf.unset(key)
+ }
+ }
+ } else {
+ super.getNormalizedQueryExecutionResult(session, sql)
+ }
+ }
+
/** List of test cases that require TPCDS table schemas to be loaded. */
private def requireTPCDSCases: Seq[String] = Seq("pipe-operators.sql")
/** List of TPCDS table names and schemas to load from the [[TPCDSSchema]] base class. */
@@ -829,9 +856,24 @@ class SQLQueryTestSuite extends SharedSparkSession with SQLHelper
s"Schema did not match for query #$i\n${expected.sql}: $output") {
output.schema
}
- assertResult(expected.output, s"Result did not match" +
+ // Comet may surface errors as `CometNativeException` instead of the matching Spark
+ // exception class when a `ScalaUDF` dispatched into the native plan evaluates a
+ // divide-by-zero (DataFusion wraps the typed error so the JNI bridge cannot downcast it).
+ // Same category, different surface. Collapse both sides to a placeholder when this happens
+ // so the literal compare passes.
+ // Tracking issue: https://github.com/apache/datafusion-comet/issues/4517
+ val (expectedOut, actualOut) = if (isCometEnabled &&
+ expected.output.startsWith("org.apache.spark.SparkArithmeticException") &&
+ expected.output.contains("\"DIVIDE_BY_ZERO\"") &&
+ output.output.startsWith("org.apache.comet.CometNativeException") &&
+ output.output.contains("DivideByZero")) {
+ ("[DIVIDE_BY_ZERO]", "[DIVIDE_BY_ZERO]")
+ } else {
+ (expected.output, output.output)
+ }
+ assertResult(expectedOut, s"Result did not match" +
s" for query #$i\n${expected.sql}") {
- output.output
+ actualOut
}
}
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
index bfcf583a705..6c177bff0e3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
@@ -217,7 +217,15 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp
}
}
+ // Comet: Comet's own columnar rules take over these plans, so the rules the tests below
+ // inject never see the Spark operators they act on or check for.
+ private def assumeCometDisabled(): Unit = {
+ assume(!org.apache.spark.sql.classic.SparkSession.isCometEnabled,
+ "Skipped when Comet is enabled: Comet replaces the operators the injected rules act on")
+ }
+
test("inject adaptive query prep rule") {
+ assumeCometDisabled()
val extensions = create { extensions =>
// inject rule that will run during AQE query stage preparation and will add custom tags
// to the plan
@@ -311,6 +319,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp
}
private def testInjectColumnar(enableAQE: Boolean): Unit = {
+ assumeCometDisabled()
def collectPlanSteps(plan: SparkPlan): Seq[Int] = plan match {
case a: AdaptiveSparkPlanExec =>
assert(a.toString.startsWith("AdaptiveSparkPlan isFinalPlan=true"))
@@ -366,6 +375,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp
}
test("reset column vectors") {
+ assumeCometDisabled()
val session = SparkSession.builder()
.master("local[1]")
.config(COLUMN_BATCH_SIZE.key, 2)
@@ -553,6 +563,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp
}
test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE Optimizer") {
+ assumeCometDisabled()
def executedPlan(df: Dataset[java.lang.Long]): SparkPlan = {
assert(df.queryExecution.executedPlan.isInstanceOf[AdaptiveSparkPlanExec])
stripAQEPlan(df.queryExecution.executedPlan)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala
index d7b2511eac2..d5f5b940b94 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala
@@ -82,6 +82,10 @@ class SparkSessionJobTaggingAndCancellationSuite
}
test("Tags set from session are prefixed with session UUID") {
+ // This test relies on job scheduling order which is timing-dependent and becomes unreliable
+ // when Comet is enabled due to changes in async execution behaviour.
+ assume(!classic.SparkSession.isCometEnabled,
+ "Skipped when Comet is enabled: test results are timing-dependent")
sc = new SparkContext("local[2]", "test")
val session = classic.SparkSession.builder().sparkContext(sc).getOrCreate()
import session.implicits._
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala
index 63589472854..f8c07a9b037 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala
@@ -17,6 +17,8 @@
package org.apache.spark.sql
+import org.apache.comet.CometConf
+
import org.apache.spark.{SPARK_DOC_ROOT, SparkIllegalArgumentException, SparkRuntimeException}
import org.apache.spark.sql.catalyst.expressions.Cast._
import org.apache.spark.sql.catalyst.expressions.IsNotNull
@@ -179,29 +181,31 @@ class StringFunctionsSuite extends SharedSparkSession {
}
test("string regex_replace / regex_extract") {
- val df = Seq(
- ("100-200", "(\\d+)-(\\d+)", "300"),
- ("100-200", "(\\d+)-(\\d+)", "400"),
- ("100-200", "(\\d+)", "400")).toDF("a", "b", "c")
+ withSQLConf(CometConf.getExprAllowIncompatConfigKey("regexp") -> "true") {
+ val df = Seq(
+ ("100-200", "(\\d+)-(\\d+)", "300"),
+ ("100-200", "(\\d+)-(\\d+)", "400"),
+ ("100-200", "(\\d+)", "400")).toDF("a", "b", "c")
- checkAnswer(
- df.select(
- regexp_replace($"a", "(\\d+)", "num"),
- regexp_replace($"a", $"b", $"c"),
- regexp_extract($"a", "(\\d+)-(\\d+)", 1)),
- Row("num-num", "300", "100") :: Row("num-num", "400", "100") ::
- Row("num-num", "400-400", "100") :: Nil)
-
- // for testing the mutable state of the expression in code gen.
- // This is a hack way to enable the codegen, thus the codegen is enable by default,
- // it will still use the interpretProjection if projection followed by a LocalRelation,
- // hence we add a filter operator.
- // See the optimizer rule `ConvertToLocalRelation`
- checkAnswer(
- df.filter("isnotnull(a)").selectExpr(
- "regexp_replace(a, b, c)",
- "regexp_extract(a, b, 1)"),
- Row("300", "100") :: Row("400", "100") :: Row("400-400", "100") :: Nil)
+ checkAnswer(
+ df.select(
+ regexp_replace($"a", "(\\d+)", "num"),
+ regexp_replace($"a", $"b", $"c"),
+ regexp_extract($"a", "(\\d+)-(\\d+)", 1)),
+ Row("num-num", "300", "100") :: Row("num-num", "400", "100") ::
+ Row("num-num", "400-400", "100") :: Nil)
+
+ // for testing the mutable state of the expression in code gen.
+ // This is a hack way to enable the codegen, thus the codegen is enable by default,
+ // it will still use the interpretProjection if projection followed by a LocalRelation,
+ // hence we add a filter operator.
+ // See the optimizer rule `ConvertToLocalRelation`
+ checkAnswer(
+ df.filter("isnotnull(a)").selectExpr(
+ "regexp_replace(a, b, c)",
+ "regexp_extract(a, b, 1)"),
+ Row("300", "100") :: Row("400", "100") :: Row("400-400", "100") :: Nil)
+ }
}
test("non-matching optional group") {
@@ -1425,7 +1429,12 @@ class StringFunctionsSuite extends SharedSparkSession {
s"'$$3 $$1') FROM $tableName"
val df = sql(query)
val plan = df.queryExecution.executedPlan
- assert(plan.isInstanceOf[WholeStageCodegenExec] == (codegenMode == "CODEGEN_ONLY"))
+ // Comet routes regexp_replace through the codegen dispatcher, so the executed plan is a
+ // Comet operator rather than WholeStageCodegenExec. The exception assertions below still
+ // hold; only this Spark-internal plan-shape check is skipped under Comet.
+ if (!isCometEnabled) {
+ assert(plan.isInstanceOf[WholeStageCodegenExec] == (codegenMode == "CODEGEN_ONLY"))
+ }
val exception = intercept[SparkRuntimeException] {
df.collect()
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
index cd3e389d765..354c5b97e8b 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala
@@ -23,10 +23,11 @@ import org.apache.spark.SparkRuntimeException
import org.apache.spark.sql.catalyst.expressions.{EqualTo, NamedExpression, OuterReference, SubqueryExpression}
import org.apache.spark.sql.catalyst.plans.{LeftAnti, LeftSemi}
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project, Sort, Union}
+import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec, CometNativeScanExec, CometScanExec}
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecution}
import org.apache.spark.sql.execution.datasources.FileScanRDD
-import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
@@ -1528,6 +1529,30 @@ class SubquerySuite extends SharedSparkSession
fs.inputRDDs().forall(
_.asInstanceOf[FileScanRDD].filePartitions.forall(
_.files.forall(_.urlEncodedPath.contains("p=0"))))
+ case WholeStageCodegenExec(ColumnarToRowExec(InputAdapter(
+ fs @ CometScanExec(_, _, _, partitionFilters, _, _, _, _, _, _)))) =>
+ partitionFilters.exists(ExecSubqueryExpression.hasSubquery) &&
+ fs.inputRDDs().forall(
+ _.asInstanceOf[FileScanRDD].filePartitions.forall(
+ _.files.forall(_.urlEncodedPath.contains("p=0"))))
+ case CometNativeColumnarToRowExec(
+ fs: CometNativeScanExec) =>
+ fs.partitionFilters.exists(ExecSubqueryExpression.hasSubquery) &&
+ fs.inputRDDs().forall(
+ _.asInstanceOf[FileScanRDD].filePartitions.forall(
+ _.files.forall(_.urlEncodedPath.contains("p=0"))))
+ case WholeStageCodegenExec(CometColumnarToRowExec(InputAdapter(
+ fs @ CometScanExec(_, _, _, partitionFilters, _, _, _, _, _, _)))) =>
+ partitionFilters.exists(ExecSubqueryExpression.hasSubquery) &&
+ fs.inputRDDs().forall(
+ _.asInstanceOf[FileScanRDD].filePartitions.forall(
+ _.files.forall(_.urlEncodedPath.contains("p=0"))))
+ case WholeStageCodegenExec(CometColumnarToRowExec(InputAdapter(
+ fs: CometNativeScanExec))) =>
+ fs.partitionFilters.exists(ExecSubqueryExpression.hasSubquery) &&
+ fs.inputRDDs().forall(
+ _.asInstanceOf[FileScanRDD].filePartitions.forall(
+ _.files.forall(_.urlEncodedPath.contains("p=0"))))
case _ => false
})
}
@@ -2093,7 +2118,7 @@ class SubquerySuite extends SharedSparkSession
df.collect()
val exchanges = collect(df.queryExecution.executedPlan) {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
}
assert(exchanges.size === 1)
}
@@ -2464,18 +2489,29 @@ class SubquerySuite extends SharedSparkSession
def checkFileSourceScan(query: String, answer: Seq[Row]): Unit = {
val df = sql(query)
checkAnswer(df, answer)
- val fileSourceScanExec = collect(df.queryExecution.executedPlan) {
- case f: FileSourceScanExec => f
+ val dataSourceScanExec = collect(df.queryExecution.executedPlan) {
+ case f: FileSourceScanLike => f
+ case c: CometScanExec => c
+ case n: CometNativeScanExec => n
}
sparkContext.listenerBus.waitUntilEmpty()
- assert(fileSourceScanExec.size === 1)
- val scalarSubquery = fileSourceScanExec.head.dataFilters.flatMap(_.collect {
- case s: ScalarSubquery => s
- })
+ assert(dataSourceScanExec.size === 1)
+ val scalarSubquery = dataSourceScanExec.head match {
+ case f: FileSourceScanLike =>
+ f.dataFilters.flatMap(_.collect { case s: ScalarSubquery => s })
+ case c: CometScanExec =>
+ c.dataFilters.flatMap(_.collect { case s: ScalarSubquery => s })
+ case n: CometNativeScanExec =>
+ n.dataFilters.flatMap(_.collect { case s: ScalarSubquery => s })
+ }
assert(scalarSubquery.length === 1)
assert(scalarSubquery.head.plan.isInstanceOf[ReusedSubqueryExec])
- assert(fileSourceScanExec.head.metrics("numFiles").value === 1)
- assert(fileSourceScanExec.head.metrics("numOutputRows").value === answer.size)
+ assert(dataSourceScanExec.head.metrics("numFiles").value === 1)
+ val numOutputRows = dataSourceScanExec.head match {
+ case n: CometNativeScanExec => n.metrics("output_rows").value
+ case other => other.metrics("numOutputRows").value
+ }
+ assert(numOutputRows === answer.size)
}
withTable("t1", "t2") {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala
index 2d26356890d..2c5994f5fbc 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala
@@ -21,6 +21,7 @@ import org.apache.spark.sql.QueryTest.sameRows
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Cast, Literal}
import org.apache.spark.sql.catalyst.expressions.variant.{ToVariantObject, VariantExpressionEvalUtils}
+import org.apache.spark.sql.comet.CometNativeColumnarToRowExec
import org.apache.spark.sql.execution.WholeStageCodegenExec
import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
import org.apache.spark.sql.functions._
@@ -391,7 +392,8 @@ class VariantEndToEndSuite extends SharedSparkSession {
s"cast(to_variant_object(s) as ${schema(2).dataType.sql})")
checkAnswer(df, input)
val plan = df.queryExecution.executedPlan
- assert(plan.isInstanceOf[WholeStageCodegenExec] == (codegenMode == "CODEGEN_ONLY"))
+ assert(plan.isInstanceOf[WholeStageCodegenExec] == (codegenMode == "CODEGEN_ONLY")
+ || plan.isInstanceOf[CometNativeColumnarToRowExec])
}
}
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala
index 37684c7fce3..f3574dec867 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala
@@ -26,6 +26,8 @@ import org.apache.spark.sql.{AnalysisException, Row}
import org.apache.spark.sql.catalyst.ExtendedAnalysisException
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.util.CollationFactory
+import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashJoinExec, CometSortMergeJoinExec}
+import org.apache.spark.sql.comet.CometHashAggregateExec
import org.apache.spark.sql.connector.{DatasourceV2SQLBase, FakeV2ProviderWithCustomSchema}
import org.apache.spark.sql.connector.catalog.{CatalogV2Util, Identifier, InMemoryTable}
import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper
@@ -61,7 +63,9 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
assert(
collectFirst(queryPlan) {
case _: SortMergeJoinExec => assert(isSortMergeForced)
+ case _: CometSortMergeJoinExec => assert(isSortMergeForced)
case _: HashJoin => assert(!isSortMergeForced)
+ case _: CometHashJoinExec | _: CometBroadcastHashJoinExec => assert(!isSortMergeForced)
}.nonEmpty
)
}
@@ -469,6 +473,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
val dfBinary = sql(s"SELECT COUNT(*), c FROM $tableNameBinary GROUP BY c")
assert(collectFirst(dfBinary.queryExecution.executedPlan) {
case _: HashAggregateExec | _: ObjectHashAggregateExec => ()
+ case _: CometHashAggregateExec => ()
}.nonEmpty)
}
}
@@ -1607,10 +1612,14 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
if (!CollationFactory.fetchCollation(t.collation).supportsBinaryEquality) {
assert(collectFirst(queryPlan) {
case b: HashJoin => b.leftKeys.head
+ case ch: CometHashJoinExec => ch.leftKeys.head
+ case cbh: CometBroadcastHashJoinExec => cbh.leftKeys.head
}.head.isInstanceOf[CollationKey])
} else {
assert(!collectFirst(queryPlan) {
case b: HashJoin => b.leftKeys.head
+ case ch: CometHashJoinExec => ch.leftKeys.head
+ case cbh: CometBroadcastHashJoinExec => cbh.leftKeys.head
}.head.isInstanceOf[CollationKey])
}
}
@@ -1666,11 +1675,13 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
if (!CollationFactory.fetchCollation(t.collation).supportsBinaryEquality) {
assert(collectFirst(queryPlan) {
case b: BroadcastHashJoinExec => b.leftKeys.head
+ case b: CometBroadcastHashJoinExec => b.leftKeys.head
}.head.asInstanceOf[ArrayTransform].function.asInstanceOf[LambdaFunction].
function.isInstanceOf[CollationKey])
} else {
assert(!collectFirst(queryPlan) {
case b: BroadcastHashJoinExec => b.leftKeys.head
+ case b: CometBroadcastHashJoinExec => b.leftKeys.head
}.head.isInstanceOf[ArrayTransform])
}
}
@@ -1736,6 +1747,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper {
} else {
assert(!collectFirst(queryPlan) {
case b: BroadcastHashJoinExec => b.leftKeys.head
+ case b: CometBroadcastHashJoinExec => b.leftKeys.head
}.head.isInstanceOf[ArrayTransform])
}
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala
index 5ae23bc3338..5c2c3fff284 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala
@@ -29,6 +29,7 @@ import org.apache.spark.SparkUnsupportedOperationException
import org.apache.spark.sql.{AnalysisException, DataFrame, Row}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{AttributeReference, GreaterThan => CatalystGreaterThan, Literal => CatalystLiteral, ScalarSubquery}
+import org.apache.spark.sql.comet.CometSortExec
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Project}
import org.apache.spark.sql.connector.catalog.{PartitionInternalRow, SupportsRead, Table, TableCapability, TableProvider}
import org.apache.spark.sql.connector.catalog.TableCapability._
@@ -41,7 +42,7 @@ import org.apache.spark.sql.execution.SortExec
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2Relation, DataSourceV2ScanRelation, V2ScanPartitioningAndOrdering}
import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Implicits._
-import org.apache.spark.sql.execution.exchange.{Exchange, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.exchange.{Exchange, ShuffleExchangeExec, ShuffleExchangeLike}
import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions._
@@ -283,13 +284,13 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper
val groupByColJ = df.groupBy($"j").agg(sum($"i"))
checkAnswer(groupByColJ, Seq(Row(2, 8), Row(4, 2), Row(6, 5)))
assert(collectFirst(groupByColJ.queryExecution.executedPlan) {
- case e: ShuffleExchangeExec => e
+ case e: ShuffleExchangeLike => e
}.isDefined)
val groupByIPlusJ = df.groupBy($"i" + $"j").agg(count("*"))
checkAnswer(groupByIPlusJ, Seq(Row(5, 2), Row(6, 2), Row(8, 1), Row(9, 1)))
assert(collectFirst(groupByIPlusJ.queryExecution.executedPlan) {
- case e: ShuffleExchangeExec => e
+ case e: ShuffleExchangeLike => e
}.isDefined)
}
}
@@ -349,10 +350,11 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper
val (shuffleExpected, sortExpected) = groupByExpects
assert(collectFirst(groupBy.queryExecution.executedPlan) {
- case e: ShuffleExchangeExec => e
+ case e: ShuffleExchangeLike => e
}.isDefined === shuffleExpected)
assert(collectFirst(groupBy.queryExecution.executedPlan) {
case e: SortExec => e
+ case c: CometSortExec => c
}.isDefined === sortExpected)
}
@@ -367,10 +369,11 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper
val (shuffleExpected, sortExpected) = windowFuncExpects
assert(collectFirst(windowPartByColIOrderByColJ.queryExecution.executedPlan) {
- case e: ShuffleExchangeExec => e
+ case e: ShuffleExchangeLike => e
}.isDefined === shuffleExpected)
assert(collectFirst(windowPartByColIOrderByColJ.queryExecution.executedPlan) {
case e: SortExec => e
+ case c: CometSortExec => c
}.isDefined === sortExpected)
}
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/FileDataSourceV2FallBackSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/FileDataSourceV2FallBackSuite.scala
index 0fbfaedfdca..7d8383d27be 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/FileDataSourceV2FallBackSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/FileDataSourceV2FallBackSuite.scala
@@ -20,6 +20,7 @@ import scala.collection.mutable.ArrayBuffer
import org.apache.spark.{SparkConf, SparkException}
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.comet.{CometNativeScanExec, CometScanExec}
import org.apache.spark.sql.connector.catalog.{SupportsRead, SupportsWrite, Table, TableCapability}
import org.apache.spark.sql.connector.read.ScanBuilder
import org.apache.spark.sql.connector.write.{LogicalWriteInfo, WriteBuilder}
@@ -190,7 +191,11 @@ class FileDataSourceV2FallBackSuite extends SharedSparkSession {
val df = spark.read.format(format).load(path.getCanonicalPath)
checkAnswer(df, inputData.toDF())
assert(
- df.queryExecution.executedPlan.exists(_.isInstanceOf[FileSourceScanExec]))
+ df.queryExecution.executedPlan.exists {
+ case _: FileSourceScanExec | _: CometScanExec | _: CometNativeScanExec => true
+ case _ => false
+ }
+ )
}
} finally {
spark.listenerManager.unregister(listener)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
index 38de6b043bc..939e4a8c4c3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
@@ -25,6 +25,8 @@ import org.apache.spark.sql.{DataFrame, ExplainSuiteHelper, Row}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, Literal, TransformExpression}
import org.apache.spark.sql.catalyst.plans.physical
+import org.apache.spark.sql.comet.CometSortMergeJoinExec
+import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
import org.apache.spark.sql.connector.catalog.{Column, Identifier, InMemoryTableCatalog}
import org.apache.spark.sql.connector.catalog.functions._
import org.apache.spark.sql.connector.distributions.Distributions
@@ -325,6 +327,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
protected def collectAllShuffles(plan: SparkPlan): Seq[ShuffleExchangeLike] = {
collect(plan) {
case s: ShuffleExchangeExec => s
+ case c: CometShuffleExchangeExec => c
}
}
@@ -338,9 +341,10 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
// here we skip collecting shuffle operators that are not associated with SMJ
collect(plan) {
case s: SortMergeJoinExec => s
+ case c: CometSortMergeJoinExec => c.originalPlan
}.flatMap(smj =>
collect(smj) {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
})
}.toSet.toSeq
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/WriteDistributionAndOrderingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/WriteDistributionAndOrderingSuite.scala
index ec1b34b8c21..ff4f2a70041 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/WriteDistributionAndOrderingSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/WriteDistributionAndOrderingSuite.scala
@@ -21,7 +21,7 @@ package org.apache.spark.sql.connector
import java.sql.Date
import java.util.Collections
-import org.apache.spark.sql.{catalyst, AnalysisException, DataFrame, Row}
+import org.apache.spark.sql.{catalyst, AnalysisException, DataFrame, IgnoreCometSuite, Row}
import org.apache.spark.sql.catalyst.expressions.{ApplyFunctionExpression, Cast, Literal, TransformExpression}
import org.apache.spark.sql.catalyst.expressions.objects.Invoke
import org.apache.spark.sql.catalyst.plans.physical
@@ -46,7 +46,8 @@ import org.apache.spark.sql.util.QueryExecutionListener
import org.apache.spark.tags.SlowSQLTest
@SlowSQLTest
-class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase {
+class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
+ with IgnoreCometSuite {
import testImplicits._
before {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionErrorsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionErrorsSuite.scala
index ce549da03b4..47adf7cf94d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionErrorsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionErrorsSuite.scala
@@ -31,7 +31,7 @@ import org.mockito.Mockito.{mock, spy, when}
import org.scalatest.time.SpanSugar._
import org.apache.spark._
-import org.apache.spark.sql.{AnalysisException, DataFrame, Dataset, Encoder, KryoData, Row, SaveMode}
+import org.apache.spark.sql.{AnalysisException, DataFrame, Dataset, Encoder, IgnoreComet, KryoData, Row, SaveMode}
import org.apache.spark.sql.catalyst.FunctionIdentifier
import org.apache.spark.sql.catalyst.analysis.{NamedParameter, UnresolvedGenerator}
import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, RowEncoder}
@@ -266,7 +266,8 @@ class QueryExecutionErrorsSuite
}
test("INCONSISTENT_BEHAVIOR_CROSS_VERSION: " +
- "compatibility with Spark 2.4/3.2 in reading/writing dates") {
+ "compatibility with Spark 2.4/3.2 in reading/writing dates",
+ IgnoreComet("Comet doesn't completely support datetime rebase mode yet")) {
// Fail to read ancient datetime values.
withSQLConf(SQLConf.PARQUET_REBASE_MODE_IN_READ.key -> EXCEPTION.toString) {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala
index 3ef22ccb77e..54f67b45387 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala
@@ -56,6 +56,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper
}
}
+ // Comet: AQE coalesces post-shuffle partitions by map output size, and Comet's shuffle
+ // writes Arrow IPC blocks whose sizes differ from Spark's, so the expected partition
+ // counts below only hold for Spark's shuffle.
+ private def isCometEnabled: Boolean = org.apache.spark.sql.classic.SparkSession.isCometEnabled
+
val numInputPartitions: Int = 10
def withSparkSession(
@@ -166,9 +171,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper
assert(shuffleReads.isEmpty)
case None =>
- assert(shuffleReads.length === 2)
- shuffleReads.foreach { read =>
- assert(read.outputPartitioning.numPartitions === 2)
+ if (!isCometEnabled) {
+ assert(shuffleReads.length === 2)
+ shuffleReads.foreach { read =>
+ assert(read.outputPartitioning.numPartitions === 2)
+ }
}
}
}
@@ -221,9 +228,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper
assert(shuffleReads.isEmpty)
case None =>
- assert(shuffleReads.length === 2)
- shuffleReads.foreach { read =>
- assert(read.outputPartitioning.numPartitions === 2)
+ if (!isCometEnabled) {
+ assert(shuffleReads.length === 2)
+ shuffleReads.foreach { read =>
+ assert(read.outputPartitioning.numPartitions === 2)
+ }
}
}
}
@@ -271,9 +280,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper
assert(shuffleReads.isEmpty)
case None =>
- assert(shuffleReads.length === 2)
- shuffleReads.foreach { read =>
- assert(read.outputPartitioning.numPartitions === 3)
+ if (!isCometEnabled) {
+ assert(shuffleReads.length === 2)
+ shuffleReads.foreach { read =>
+ assert(read.outputPartitioning.numPartitions === 3)
+ }
}
}
}
@@ -530,10 +541,12 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper
// Shuffle partition coalescing of the join is performed independent of the non-grouping
// aggregate on the other side of the union.
val finalPlan = stripAQEPlan(resultDf.queryExecution.executedPlan)
- assert(
- finalPlan.collect {
- case r @ CoalescedShuffleRead() => r
- }.size == 2)
+ if (!isCometEnabled) {
+ assert(
+ finalPlan.collect {
+ case r @ CoalescedShuffleRead() => r
+ }.size == 2)
+ }
}
withSparkSession(test, 100, None)
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala
index 60bf3a21e96..7693107f107 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala
@@ -23,7 +23,7 @@ import scala.util.Random
import org.apache.hadoop.fs.Path
import org.apache.spark.SparkConf
-import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.{DataFrame, IgnoreComet}
import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
import org.apache.spark.sql.execution.datasources.v2.orc.OrcScan
import org.apache.spark.sql.internal.SQLConf
@@ -195,7 +195,7 @@ class DataSourceV2ScanExecRedactionSuite extends DataSourceScanRedactionTest {
}
}
- test("FileScan description") {
+ test("FileScan description", IgnoreComet("Comet doesn't use BatchScan")) {
Seq("json", "orc", "parquet").foreach { format =>
withTempPath { path =>
val dir = path.getCanonicalPath
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffsetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffsetSuite.scala
index faf0469ec20..dfed25d87fe 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffsetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffsetSuite.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql.execution
import org.apache.spark.sql.Dataset
import org.apache.spark.sql.IntegratedUDFTestUtils._
+import org.apache.spark.sql.comet.{CometCollectLimitExec, CometGlobalLimitExec, CometProjectExec, CometSortExec}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.functions.rand
import org.apache.spark.sql.internal.SQLConf
@@ -38,7 +39,7 @@ class InsertSortForLimitAndOffsetSuite extends SharedSparkSession
private def assertHasCollectLimitExec(plan: SparkPlan): Unit = {
assert(find(plan) {
- case _: CollectLimitExec => true
+ case _: CollectLimitExec | _: CometCollectLimitExec => true
case _ => false
}.isDefined)
}
@@ -46,6 +47,7 @@ class InsertSortForLimitAndOffsetSuite extends SharedSparkSession
private def assertHasGlobalLimitExec(plan: SparkPlan): Unit = {
assert(find(plan) {
case _: GlobalLimitExec => true
+ case _: CometGlobalLimitExec => true
case _ => false
}.isDefined)
}
@@ -54,6 +56,11 @@ class InsertSortForLimitAndOffsetSuite extends SharedSparkSession
find(plan) {
case GlobalLimitExec(_, s: SortExec, _) => !s.global
case GlobalLimitExec(_, ProjectExec(_, s: SortExec), _) => !s.global
+ case CometGlobalLimitExec(_, _, _, _, s: CometSortExec, _) =>
+ !s.originalPlan.asInstanceOf[SortExec].global
+ case CometGlobalLimitExec(_, _, _, _,
+ CometProjectExec(_, _, _, _, s: CometSortExec, _), _) =>
+ !s.originalPlan.asInstanceOf[SortExec].global
case _ => false
}.isDefined
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/LogicalPlanTagInSparkPlanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/LogicalPlanTagInSparkPlanSuite.scala
index 743ec41dbe7..9f30d6c8e04 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/LogicalPlanTagInSparkPlanSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/LogicalPlanTagInSparkPlanSuite.scala
@@ -53,6 +53,10 @@ class LogicalPlanTagInSparkPlanSuite extends TPCDSQuerySuite with DisableAdaptiv
case ColumnarToRowExec(i: InputAdapter) => isScanPlanTree(i.child)
case p: ProjectExec => isScanPlanTree(p.child)
case f: FilterExec => isScanPlanTree(f.child)
+ // Comet produces scan plan tree like:
+ // ColumnarToRow
+ // +- ReusedExchange
+ case _: ReusedExchangeExec => false
case _: LeafExecNode => true
case _ => false
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
index 7c434ebee5f..ac18158100d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.execution
import org.apache.spark.SparkUnsupportedOperationException
import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.{execution, DataFrame, Row}
+import org.apache.spark.sql.{execution, DataFrame, IgnoreCometSuite, Row}
import org.apache.spark.sql.AnalysisException
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
@@ -38,7 +38,9 @@ import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types._
-class PlannerSuite extends SharedSparkSession with AdaptiveSparkPlanHelper {
+// Ignore this suite when Comet is enabled. This suite tests the Spark planner and Comet planner
+// comes out with too many difference. Simply ignoring this suite for now.
+class PlannerSuite extends SharedSparkSession with AdaptiveSparkPlanHelper with IgnoreCometSuite {
import testImplicits._
setupTestData()
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala
index f7afdb5e6e5..e5a2eeb86a0 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala
@@ -22,7 +22,7 @@ import scala.io.Source
import scala.util.Try
import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent, SparkListenerJobStart}
-import org.apache.spark.sql.{AnalysisException, ExtendedExplainGenerator, FastOperator, SaveMode}
+import org.apache.spark.sql.{AnalysisException, ExtendedExplainGenerator, FastOperator, IgnoreComet, SaveMode}
import org.apache.spark.sql.catalyst.{QueryPlanningTracker, QueryPlanningTrackerCallback, TableIdentifier}
import org.apache.spark.sql.catalyst.analysis.{CurrentNamespace, UnresolvedFunction, UnresolvedRelation}
import org.apache.spark.sql.catalyst.expressions.{Alias, UnsafeRow}
@@ -461,7 +461,7 @@ class QueryExecutionSuite extends SharedSparkSession {
}
}
- test("SPARK-47289: extended explain info") {
+ test("SPARK-47289: extended explain info", IgnoreComet("Comet plan extended info is different")) {
val concat = new PlanStringConcat()
withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantProjectsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantProjectsSuite.scala
index 6d259ed19d0..e271f7a2914 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantProjectsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantProjectsSuite.scala
@@ -17,7 +17,10 @@
package org.apache.spark.sql.execution
+import org.apache.comet.CometConf
+
import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.comet.CometProjectExec
import org.apache.spark.sql.connector.SimpleWritableDataSource
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
import org.apache.spark.sql.internal.SQLConf
@@ -33,7 +36,10 @@ abstract class RemoveRedundantProjectsSuiteBase
private def assertProjectExecCount(df: DataFrame, expected: Int): Unit = {
withClue(df.queryExecution) {
val plan = df.queryExecution.executedPlan
- val actual = collectWithSubqueries(plan) { case p: ProjectExec => p }.size
+ val actual = collectWithSubqueries(plan) {
+ case p: ProjectExec => p
+ case p: CometProjectExec => p
+ }.size
assert(actual == expected)
}
}
@@ -133,12 +139,25 @@ abstract class RemoveRedundantProjectsSuiteBase
val df = data.selectExpr("a", "b", "key", "explode(array(key, a, b)) as d").filter("d > 0")
df.collect()
val plan = df.queryExecution.executedPlan
- val numProjects = collectWithSubqueries(plan) { case p: ProjectExec => p }.length
+ val numProjects = collectWithSubqueries(plan) {
+ case p: ProjectExec => p
+ case p: CometProjectExec => p
+ }.length
// Create a new plan that reverse the GenerateExec output and add a new ProjectExec between
// GenerateExec and its child. This is to test if the ProjectExec is removed, the output of
// the query will be incorrect.
- val newPlan = stripAQEPlan(plan) transform {
+
+ // Comet-specific change to get original Spark plan before applying
+ // a transformation to add a new ProjectExec
+ var sparkPlan: SparkPlan = null
+ withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "false") {
+ val df = data.selectExpr("a", "b", "key", "explode(array(key, a, b)) as d").filter("d > 0")
+ df.collect()
+ sparkPlan = df.queryExecution.executedPlan
+ }
+
+ val newPlan = stripAQEPlan(sparkPlan) transform {
case g @ GenerateExec(_, requiredChildOutput, _, _, child) =>
g.copy(requiredChildOutput = requiredChildOutput.reverse,
child = ProjectExec(requiredChildOutput.reverse, child))
@@ -150,6 +169,7 @@ abstract class RemoveRedundantProjectsSuiteBase
// The manually added ProjectExec node shouldn't be removed.
assert(collectWithSubqueries(newExecutedPlan) {
case p: ProjectExec => p
+ case p: CometProjectExec => p
}.size == numProjects + 1)
// Check the original plan's output and the new plan's output are the same.
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantSortsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantSortsSuite.scala
index 3cba30079cd..08513544cc0 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantSortsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantSortsSuite.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql.execution
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.catalyst.plans.physical.{RangePartitioning, UnknownPartitioning}
+import org.apache.spark.sql.comet.CometSortExec
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
import org.apache.spark.sql.execution.joins.ShuffledJoin
import org.apache.spark.sql.internal.SQLConf
@@ -32,7 +33,7 @@ abstract class RemoveRedundantSortsSuiteBase
private def checkNumSorts(df: DataFrame, count: Int): Unit = {
val plan = df.queryExecution.executedPlan
- assert(collectWithSubqueries(plan) { case s: SortExec => s }.length == count)
+ assert(collectWithSubqueries(plan) { case _: SortExec | _: CometSortExec => 1 }.length == count)
}
private def checkSorts(query: String, enabledCount: Int, disabledCount: Int): Unit = {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala
index 03e2f4942ed..b0a80fcf683 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/RemoveRedundantWindowGroupLimitsSuite.scala
@@ -18,6 +18,7 @@
package org.apache.spark.sql.execution
import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.comet.CometWindowGroupLimitExec
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
import org.apache.spark.sql.execution.window.WindowGroupLimitExec
import org.apache.spark.sql.functions.lit
@@ -29,7 +30,10 @@ abstract class RemoveRedundantWindowGroupLimitsSuiteBase
private def checkNumWindowGroupLimits(df: DataFrame, count: Int): Unit = {
val plan = df.queryExecution.executedPlan
- assert(collectWithSubqueries(plan) { case exec: WindowGroupLimitExec => exec }.length == count)
+ assert(collectWithSubqueries(plan) {
+ case exec: WindowGroupLimitExec => exec
+ case exec: CometWindowGroupLimitExec => exec
+ }.length == count)
}
private def checkWindowGroupLimits(query: String, count: Int): Unit = {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
index 41edb53441f..fabfa9c883e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReplaceHashWithSortAggSuite.scala
@@ -18,6 +18,7 @@
package org.apache.spark.sql.execution
import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.comet.{CometHashAggregateExec, CometSortAggregateExec}
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
import org.apache.spark.sql.internal.SQLConf
@@ -30,9 +31,11 @@ abstract class ReplaceHashWithSortAggSuiteBase
private def checkNumAggs(df: DataFrame, hashAggCount: Int, sortAggCount: Int): Unit = {
val plan = df.queryExecution.executedPlan
assert(collectWithSubqueries(plan) {
- case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec) => s
+ case s @ (_: HashAggregateExec | _: ObjectHashAggregateExec | _: CometHashAggregateExec ) => s
}.length == hashAggCount)
- assert(collectWithSubqueries(plan) { case s: SortAggregateExec => s }.length == sortAggCount)
+ assert(collectWithSubqueries(plan) {
+ case s @ (_: SortAggregateExec | _: CometSortAggregateExec) => s
+ }.length == sortAggCount)
}
private def checkAggs(
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala
index e8a39eb8c53..a621c7a7923 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLWindowFunctionSuite.scala
@@ -18,7 +18,7 @@
package org.apache.spark.sql.execution
import org.apache.spark.TestUtils.assertSpilled
-import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.{AnalysisException, IgnoreComet, Row}
import org.apache.spark.sql.internal.SQLConf.{WINDOW_EXEC_BUFFER_IN_MEMORY_THRESHOLD, WINDOW_EXEC_BUFFER_SPILL_THRESHOLD}
import org.apache.spark.sql.test.SharedSparkSession
@@ -470,7 +470,7 @@ class SQLWindowFunctionSuite extends SharedSparkSession {
Row(1, 3, null) :: Row(2, null, 4) :: Nil)
}
- test("test with low buffer spill threshold") {
+ test("test with low buffer spill threshold", IgnoreComet("Comet does not support spilling")) {
val nums = sparkContext.parallelize(1 to 10).map(x => (x, x % 2)).toDF("x", "y")
nums.createOrReplaceTempView("nums")
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala
index b167dd13dcb..b2b6b88ce66 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala
@@ -31,6 +31,7 @@ import org.apache.spark.sql.catalyst.expressions.{
import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode}
import org.apache.spark.sql.catalyst.plans.logical.Deduplicate
import org.apache.spark.sql.catalyst.trees.LeafLike
+import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec}
import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
@@ -166,7 +167,11 @@ class SparkPlanSuite extends SharedSparkSession {
spark.range(1).write.parquet(path.getAbsolutePath)
val df = spark.read.parquet(path.getAbsolutePath)
val columnarToRowExec =
- df.queryExecution.executedPlan.collectFirst { case p: ColumnarToRowExec => p }.get
+ df.queryExecution.executedPlan.collectFirst {
+ case p: ColumnarToRowExec => p
+ case p: CometColumnarToRowExec => p
+ case p: CometNativeColumnarToRowExec => p
+ }.get
try {
spark.range(1).foreach { _ =>
columnarToRowExec.canonicalized
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala
index 835c55fe4c4..57cce5f6b03 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala
@@ -20,7 +20,7 @@ package org.apache.spark.sql.execution
import java.util.concurrent.{CountDownLatch, Executors, TimeUnit}
import org.apache.spark.SparkConf
-import org.apache.spark.sql.{DataFrame, QueryTest, Row}
+import org.apache.spark.sql.{DataFrame, IgnoreComet, QueryTest, Row}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
@@ -605,7 +605,10 @@ class UnionCodegenSuite extends QueryTest with SharedSparkSession {
}
}
- test("SPARK-56482: partitioning-aware union falls back to non-codegen") {
+ test("SPARK-56482: partitioning-aware union falls back to non-codegen",
+ // Comet replaces Spark's UnionExec with CometUnionExec, so the executed plan does not
+ // expose a UnionExec for this WSCG-internal assertion to inspect.
+ IgnoreComet("https://github.com/apache/datafusion-comet/issues/4965")) {
// After repartition, both children expose a `HashPartitioning` on the same key,
// so `UnionExec.outputPartitioning` is non-Unknown and the codegen path is denied.
// AQE is disabled here so the executedPlan exposes the UnionExec directly
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
index d70bd715879..074a9fa29d9 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
@@ -22,9 +22,10 @@ import java.time.Duration
import org.apache.spark.SparkException
import org.apache.spark.rdd.MapPartitionsWithEvaluatorRDD
-import org.apache.spark.sql.{Dataset, Row, SaveMode}
+import org.apache.spark.sql.{Dataset, IgnoreCometSuite, Row, SaveMode}
import org.apache.spark.sql.catalyst.expressions.{And, Cast, CodegenObjectFactoryMode, Expression, IsNotNull}
import org.apache.spark.sql.catalyst.expressions.codegen.{ByteCodeStats, CodeAndComment, CodeGenerator}
+import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometHashJoinExec, CometSortExec, CometSortMergeJoinExec}
import org.apache.spark.sql.execution.adaptive.DisableAdaptiveExecutionSuite
import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, SortAggregateExec}
import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
@@ -37,7 +38,7 @@ import org.apache.spark.sql.types.{DayTimeIntervalType, DecimalType, IntegerType
// Disable AQE because the WholeStageCodegenExec is added when running QueryStageExec
class WholeStageCodegenSuite extends SharedSparkSession
- with DisableAdaptiveExecutionSuite {
+ with DisableAdaptiveExecutionSuite with IgnoreCometSuite {
import testImplicits._
@@ -176,6 +177,7 @@ class WholeStageCodegenSuite extends SharedSparkSession
val oneJoinDF = df1.join(df2.hint("SHUFFLE_HASH"), $"k1" === $"k2")
assert(oneJoinDF.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_ : ShuffledHashJoinExec) => true
+ case _: CometHashJoinExec => true
}.size === 1)
checkAnswer(oneJoinDF, Seq(Row(0, 0), Row(1, 1), Row(2, 2), Row(3, 3), Row(4, 4)))
@@ -184,6 +186,7 @@ class WholeStageCodegenSuite extends SharedSparkSession
.join(df3.hint("SHUFFLE_HASH"), $"k1" === $"k3")
assert(twoJoinsDF.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_ : ShuffledHashJoinExec) => true
+ case _: CometHashJoinExec => true
}.size === 2)
checkAnswer(twoJoinsDF,
Seq(Row(0, 0, 0), Row(1, 1, 1), Row(2, 2, 2), Row(3, 3, 3), Row(4, 4, 4)))
@@ -210,6 +213,8 @@ class WholeStageCodegenSuite extends SharedSparkSession
assert(joinUniqueDF.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_ : ShuffledHashJoinExec) if hint == "SHUFFLE_HASH" => true
case WholeStageCodegenExec(_ : SortMergeJoinExec) if hint == "SHUFFLE_MERGE" => true
+ case _: CometHashJoinExec if hint == "SHUFFLE_HASH" => true
+ case _: CometSortMergeJoinExec if hint == "SHUFFLE_MERGE" => true
}.size === 1)
checkAnswer(joinUniqueDF, Seq(Row(0, 0), Row(1, 1), Row(2, 2), Row(3, 3), Row(4, 4),
Row(null, 5), Row(null, 6), Row(null, 7), Row(null, 8), Row(null, 9)))
@@ -220,6 +225,8 @@ class WholeStageCodegenSuite extends SharedSparkSession
assert(joinNonUniqueDF.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_ : ShuffledHashJoinExec) if hint == "SHUFFLE_HASH" => true
case WholeStageCodegenExec(_ : SortMergeJoinExec) if hint == "SHUFFLE_MERGE" => true
+ case _: CometHashJoinExec if hint == "SHUFFLE_HASH" => true
+ case _: CometSortMergeJoinExec if hint == "SHUFFLE_MERGE" => true
}.size === 1)
checkAnswer(joinNonUniqueDF, Seq(Row(0, 0), Row(0, 3), Row(0, 6), Row(0, 9), Row(1, 1),
Row(1, 4), Row(1, 7), Row(2, 2), Row(2, 5), Row(2, 8), Row(3, null), Row(4, null)))
@@ -230,6 +237,8 @@ class WholeStageCodegenSuite extends SharedSparkSession
assert(joinWithNonEquiDF.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_ : ShuffledHashJoinExec) if hint == "SHUFFLE_HASH" => true
case WholeStageCodegenExec(_ : SortMergeJoinExec) if hint == "SHUFFLE_MERGE" => true
+ case _: CometHashJoinExec if hint == "SHUFFLE_HASH" => true
+ case _: CometSortMergeJoinExec if hint == "SHUFFLE_MERGE" => true
}.size === 1)
checkAnswer(joinWithNonEquiDF, Seq(Row(0, 0), Row(0, 6), Row(0, 9), Row(1, 1),
Row(1, 7), Row(2, 2), Row(2, 8), Row(3, null), Row(4, null), Row(null, 3), Row(null, 4),
@@ -241,6 +250,8 @@ class WholeStageCodegenSuite extends SharedSparkSession
assert(twoJoinsDF.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_ : ShuffledHashJoinExec) if hint == "SHUFFLE_HASH" => true
case WholeStageCodegenExec(_ : SortMergeJoinExec) if hint == "SHUFFLE_MERGE" => true
+ case _: CometHashJoinExec if hint == "SHUFFLE_HASH" => true
+ case _: CometSortMergeJoinExec if hint == "SHUFFLE_MERGE" => true
}.size === 2)
checkAnswer(twoJoinsDF,
Seq(Row(0, 0, 0), Row(1, 1, null), Row(2, 2, 2), Row(3, 3, null), Row(4, 4, null),
@@ -262,6 +273,8 @@ class WholeStageCodegenSuite extends SharedSparkSession
assert(rightJoinUniqueDf.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_: ShuffledHashJoinExec) if hint == "SHUFFLE_HASH" => true
case WholeStageCodegenExec(_: SortMergeJoinExec) if hint == "SHUFFLE_MERGE" => true
+ case _: CometHashJoinExec if hint == "SHUFFLE_HASH" => true
+ case _: CometSortMergeJoinExec if hint == "SHUFFLE_MERGE" => true
}.size === 1)
checkAnswer(rightJoinUniqueDf, Seq(Row(1, 1), Row(2, 2), Row(3, 3), Row(4, 4),
Row(null, 5), Row(null, 6), Row(null, 7), Row(null, 8), Row(null, 9),
@@ -273,6 +286,8 @@ class WholeStageCodegenSuite extends SharedSparkSession
assert(leftJoinUniqueDf.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_: ShuffledHashJoinExec) if hint == "SHUFFLE_HASH" => true
case WholeStageCodegenExec(_: SortMergeJoinExec) if hint == "SHUFFLE_MERGE" => true
+ case _: CometHashJoinExec if hint == "SHUFFLE_HASH" => true
+ case _: CometSortMergeJoinExec if hint == "SHUFFLE_MERGE" => true
}.size === 1)
checkAnswer(leftJoinUniqueDf, Seq(Row(0, null), Row(1, 1), Row(2, 2), Row(3, 3), Row(4, 4)))
assert(leftJoinUniqueDf.count() === 5)
@@ -282,6 +297,8 @@ class WholeStageCodegenSuite extends SharedSparkSession
assert(rightJoinNonUniqueDf.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_: ShuffledHashJoinExec) if hint == "SHUFFLE_HASH" => true
case WholeStageCodegenExec(_: SortMergeJoinExec) if hint == "SHUFFLE_MERGE" => true
+ case _: CometHashJoinExec if hint == "SHUFFLE_HASH" => true
+ case _: CometSortMergeJoinExec if hint == "SHUFFLE_MERGE" => true
}.size === 1)
checkAnswer(rightJoinNonUniqueDf, Seq(Row(0, 3), Row(0, 6), Row(0, 9), Row(1, 1),
Row(1, 4), Row(1, 7), Row(1, 10), Row(2, 2), Row(2, 5), Row(2, 8)))
@@ -291,6 +308,8 @@ class WholeStageCodegenSuite extends SharedSparkSession
assert(leftJoinNonUniqueDf.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_: ShuffledHashJoinExec) if hint == "SHUFFLE_HASH" => true
case WholeStageCodegenExec(_: SortMergeJoinExec) if hint == "SHUFFLE_MERGE" => true
+ case _: CometHashJoinExec if hint == "SHUFFLE_HASH" => true
+ case _: CometSortMergeJoinExec if hint == "SHUFFLE_MERGE" => true
}.size === 1)
checkAnswer(leftJoinNonUniqueDf, Seq(Row(0, 3), Row(0, 6), Row(0, 9), Row(1, 1),
Row(1, 4), Row(1, 7), Row(1, 10), Row(2, 2), Row(2, 5), Row(2, 8), Row(3, null),
@@ -302,6 +321,8 @@ class WholeStageCodegenSuite extends SharedSparkSession
assert(rightJoinWithNonEquiDf.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_: ShuffledHashJoinExec) if hint == "SHUFFLE_HASH" => true
case WholeStageCodegenExec(_: SortMergeJoinExec) if hint == "SHUFFLE_MERGE" => true
+ case _: CometHashJoinExec if hint == "SHUFFLE_HASH" => true
+ case _: CometSortMergeJoinExec if hint == "SHUFFLE_MERGE" => true
}.size === 1)
checkAnswer(rightJoinWithNonEquiDf, Seq(Row(0, 6), Row(0, 9), Row(1, 1), Row(1, 7),
Row(1, 10), Row(2, 2), Row(2, 8), Row(null, 3), Row(null, 4), Row(null, 5)))
@@ -312,6 +333,8 @@ class WholeStageCodegenSuite extends SharedSparkSession
assert(leftJoinWithNonEquiDf.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_: ShuffledHashJoinExec) if hint == "SHUFFLE_HASH" => true
case WholeStageCodegenExec(_: SortMergeJoinExec) if hint == "SHUFFLE_MERGE" => true
+ case _: CometHashJoinExec if hint == "SHUFFLE_HASH" => true
+ case _: CometSortMergeJoinExec if hint == "SHUFFLE_MERGE" => true
}.size === 1)
checkAnswer(leftJoinWithNonEquiDf, Seq(Row(0, 6), Row(0, 9), Row(1, 1), Row(1, 7),
Row(1, 10), Row(2, 2), Row(2, 8), Row(3, null), Row(4, null)))
@@ -322,6 +345,8 @@ class WholeStageCodegenSuite extends SharedSparkSession
assert(twoRightJoinsDf.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_: ShuffledHashJoinExec) if hint == "SHUFFLE_HASH" => true
case WholeStageCodegenExec(_: SortMergeJoinExec) if hint == "SHUFFLE_MERGE" => true
+ case _: CometHashJoinExec if hint == "SHUFFLE_HASH" => true
+ case _: CometSortMergeJoinExec if hint == "SHUFFLE_MERGE" => true
}.size === 2)
checkAnswer(twoRightJoinsDf, Seq(Row(2, 2, 2), Row(3, 3, 3), Row(4, 4, 4)))
@@ -331,6 +356,8 @@ class WholeStageCodegenSuite extends SharedSparkSession
assert(twoLeftJoinsDf.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_: ShuffledHashJoinExec) if hint == "SHUFFLE_HASH" => true
case WholeStageCodegenExec(_: SortMergeJoinExec) if hint == "SHUFFLE_MERGE" => true
+ case _: CometHashJoinExec if hint == "SHUFFLE_HASH" => true
+ case _: CometSortMergeJoinExec if hint == "SHUFFLE_MERGE" => true
}.size === 2)
checkAnswer(twoLeftJoinsDf,
Seq(Row(0, null, null), Row(1, 1, null), Row(2, 2, 2), Row(3, 3, 3), Row(4, 4, 4)))
@@ -347,6 +374,7 @@ class WholeStageCodegenSuite extends SharedSparkSession
val oneLeftOuterJoinDF = df1.join(df2.hint("SHUFFLE_MERGE"), $"k1" === $"k2", "left_outer")
assert(oneLeftOuterJoinDF.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_ : SortMergeJoinExec) => true
+ case _: CometSortMergeJoinExec => true
}.size === 1)
checkAnswer(oneLeftOuterJoinDF, Seq(Row(0, 0), Row(1, 1), Row(2, 2), Row(3, 3), Row(4, null),
Row(5, null), Row(6, null), Row(7, null), Row(8, null), Row(9, null)))
@@ -355,6 +383,7 @@ class WholeStageCodegenSuite extends SharedSparkSession
val oneRightOuterJoinDF = df2.join(df3.hint("SHUFFLE_MERGE"), $"k2" === $"k3", "right_outer")
assert(oneRightOuterJoinDF.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_ : SortMergeJoinExec) => true
+ case _: CometSortMergeJoinExec => true
}.size === 1)
checkAnswer(oneRightOuterJoinDF, Seq(Row(0, 0), Row(1, 1), Row(2, 2), Row(3, 3), Row(null, 4),
Row(null, 5)))
@@ -364,6 +393,7 @@ class WholeStageCodegenSuite extends SharedSparkSession
.join(df1.hint("SHUFFLE_MERGE"), $"k3" === $"k1", "right_outer")
assert(twoJoinsDF.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(_ : SortMergeJoinExec) => true
+ case _: CometSortMergeJoinExec => true
}.size === 2)
checkAnswer(twoJoinsDF,
Seq(Row(0, 0, 0), Row(1, 1, 1), Row(2, 2, 2), Row(3, 3, 3), Row(4, null, 4), Row(5, null, 5),
@@ -379,6 +409,7 @@ class WholeStageCodegenSuite extends SharedSparkSession
val oneJoinDF = df1.join(df2.hint("SHUFFLE_MERGE"), $"k1" === $"k2", "left_semi")
assert(oneJoinDF.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(ProjectExec(_, _ : SortMergeJoinExec)) => true
+ case _: CometSortMergeJoinExec => true
}.size === 1)
checkAnswer(oneJoinDF, Seq(Row(0), Row(1), Row(2), Row(3)))
@@ -386,8 +417,8 @@ class WholeStageCodegenSuite extends SharedSparkSession
val twoJoinsDF = df3.join(df2.hint("SHUFFLE_MERGE"), $"k3" === $"k2", "left_semi")
.join(df1.hint("SHUFFLE_MERGE"), $"k3" === $"k1", "left_semi")
assert(twoJoinsDF.queryExecution.executedPlan.collect {
- case WholeStageCodegenExec(ProjectExec(_, _ : SortMergeJoinExec)) |
- WholeStageCodegenExec(_ : SortMergeJoinExec) => true
+ case _: SortMergeJoinExec => true
+ case _: CometSortMergeJoinExec => true
}.size === 2)
checkAnswer(twoJoinsDF, Seq(Row(0), Row(1), Row(2), Row(3)))
}
@@ -401,6 +432,7 @@ class WholeStageCodegenSuite extends SharedSparkSession
val oneJoinDF = df1.join(df2.hint("SHUFFLE_MERGE"), $"k1" === $"k2", "left_anti")
assert(oneJoinDF.queryExecution.executedPlan.collect {
case WholeStageCodegenExec(ProjectExec(_, _ : SortMergeJoinExec)) => true
+ case _: CometSortMergeJoinExec => true
}.size === 1)
checkAnswer(oneJoinDF, Seq(Row(4), Row(5), Row(6), Row(7), Row(8), Row(9)))
@@ -408,8 +440,8 @@ class WholeStageCodegenSuite extends SharedSparkSession
val twoJoinsDF = df1.join(df2.hint("SHUFFLE_MERGE"), $"k1" === $"k2", "left_anti")
.join(df3.hint("SHUFFLE_MERGE"), $"k1" === $"k3", "left_anti")
assert(twoJoinsDF.queryExecution.executedPlan.collect {
- case WholeStageCodegenExec(ProjectExec(_, _ : SortMergeJoinExec)) |
- WholeStageCodegenExec(_ : SortMergeJoinExec) => true
+ case _: SortMergeJoinExec => true
+ case _: CometSortMergeJoinExec => true
}.size === 2)
checkAnswer(twoJoinsDF, Seq(Row(6), Row(7), Row(8), Row(9)))
}
@@ -542,7 +574,10 @@ class WholeStageCodegenSuite extends SharedSparkSession
val plan = df.queryExecution.executedPlan
assert(plan.exists(p =>
p.isInstanceOf[WholeStageCodegenExec] &&
- p.asInstanceOf[WholeStageCodegenExec].child.isInstanceOf[SortExec]))
+ p.asInstanceOf[WholeStageCodegenExec].collect {
+ case _: SortExec => true
+ case _: CometSortExec => true
+ }.nonEmpty))
assert(df.collect() === Array(Row(1), Row(2), Row(3)))
}
@@ -722,7 +757,9 @@ class WholeStageCodegenSuite extends SharedSparkSession
.write.mode(SaveMode.Overwrite).parquet(path)
withSQLConf(SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key -> "255",
- SQLConf.WHOLESTAGE_SPLIT_CONSUME_FUNC_BY_OPERATOR.key -> "true") {
+ SQLConf.WHOLESTAGE_SPLIT_CONSUME_FUNC_BY_OPERATOR.key -> "true",
+ // Disable Comet native execution because this checks wholestage codegen.
+ "spark.comet.exec.enabled" -> "false") {
val projection = Seq.tabulate(columnNum)(i => s"c$i + c$i as newC$i")
val df = spark.read.parquet(path).selectExpr(projection: _*)
@@ -842,6 +879,8 @@ class WholeStageCodegenSuite extends SharedSparkSession
assert(distinctWithId.queryExecution.executedPlan.exists {
case WholeStageCodegenExec(ProjectExec(_,
BroadcastHashJoinExec(_, _, _, _, _, _: HashAggregateExec, _, _, _))) => true
+ case WholeStageCodegenExec(ProjectExec(_,
+ BroadcastHashJoinExec(_, _, _, _, _, _: CometColumnarToRowExec, _, _, _))) => true
case _ => false
})
checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0)))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
index d6d19d21e65..85e2f441f81 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
@@ -27,13 +27,15 @@ import org.apache.spark.SparkException
import org.apache.spark.rdd.RDD
import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent, SparkListenerJobStart}
import org.apache.spark.shuffle.sort.SortShuffleManager
-import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}
+import org.apache.spark.sql.{DataFrame, Dataset, IgnoreComet, Row, SparkSession}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, EqualTo, IsNull, Or}
import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight}
import org.apache.spark.sql.catalyst.plans.{Inner, LeftAnti}
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Join, JoinHint, LocalRelation, LogicalPlan}
import org.apache.spark.sql.classic.Strategy
+import org.apache.spark.sql.comet._
+import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
import org.apache.spark.sql.execution.columnar.{InMemoryTableScanExec, InMemoryTableScanLike}
@@ -124,6 +126,7 @@ class AdaptiveQueryExecSuite
private def findTopLevelBroadcastHashJoin(plan: SparkPlan): Seq[BroadcastHashJoinExec] = {
collect(plan) {
case j: BroadcastHashJoinExec => j
+ case j: CometBroadcastHashJoinExec => j.originalPlan.asInstanceOf[BroadcastHashJoinExec]
}
}
@@ -136,36 +139,55 @@ class AdaptiveQueryExecSuite
private def findTopLevelSortMergeJoin(plan: SparkPlan): Seq[SortMergeJoinExec] = {
collect(plan) {
case j: SortMergeJoinExec => j
+ case j: CometSortMergeJoinExec =>
+ assert(j.originalPlan.isInstanceOf[SortMergeJoinExec])
+ j.originalPlan.asInstanceOf[SortMergeJoinExec]
}
}
private def findTopLevelShuffledHashJoin(plan: SparkPlan): Seq[ShuffledHashJoinExec] = {
collect(plan) {
case j: ShuffledHashJoinExec => j
+ case j: CometHashJoinExec => j.originalPlan.asInstanceOf[ShuffledHashJoinExec]
}
}
private def findTopLevelBaseJoin(plan: SparkPlan): Seq[BaseJoinExec] = {
collect(plan) {
case j: BaseJoinExec => j
+ case c: CometHashJoinExec => c.originalPlan.asInstanceOf[BaseJoinExec]
+ case c: CometSortMergeJoinExec => c.originalPlan.asInstanceOf[BaseJoinExec]
+ case c: CometBroadcastHashJoinExec => c.originalPlan.asInstanceOf[BaseJoinExec]
}
}
+ private def isEmptyRelation(plan: SparkPlan): Boolean = stripAQEPlan(plan) match {
+ case _: EmptyRelationExec | _: CometEmptyRelationExec => true
+ case CometScanWrapper(_, _: CometEmptyRelationExec) => true
+ case p: WholeStageCodegenExec => isEmptyRelation(p.child)
+ case p: InputAdapter => isEmptyRelation(p.child)
+ case p: ColumnarToRowTransition => isEmptyRelation(p.child)
+ case _ => false
+ }
+
private def findTopLevelSort(plan: SparkPlan): Seq[SortExec] = {
collect(plan) {
case s: SortExec => s
+ case s: CometSortExec => s.originalPlan.asInstanceOf[SortExec]
}
}
private def findTopLevelAggregate(plan: SparkPlan): Seq[BaseAggregateExec] = {
collect(plan) {
case agg: BaseAggregateExec => agg
+ case agg: CometBaseAggregateExec => agg.originalPlan.asInstanceOf[BaseAggregateExec]
}
}
private def findTopLevelLimit(plan: SparkPlan): Seq[CollectLimitExec] = {
collect(plan) {
case l: CollectLimitExec => l
+ case l: CometCollectLimitExec => l.originalPlan.asInstanceOf[CollectLimitExec]
}
}
@@ -209,6 +231,7 @@ class AdaptiveQueryExecSuite
val parts = rdd.partitions
assert(parts.forall(rdd.preferredLocations(_).nonEmpty))
}
+
assert(numShuffles === (numLocalReads.length + numShufflesWithoutLocalRead))
}
@@ -217,7 +240,7 @@ class AdaptiveQueryExecSuite
val plan = df.queryExecution.executedPlan
assert(plan.isInstanceOf[AdaptiveSparkPlanExec])
val shuffle = plan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan.collect {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
}
assert(shuffle.size == 1)
assert(shuffle(0).outputPartitioning.numPartitions == numPartition)
@@ -233,7 +256,8 @@ class AdaptiveQueryExecSuite
assert(smj.size == 1)
val bhj = findTopLevelBroadcastHashJoin(adaptivePlan)
assert(bhj.size == 1)
- checkNumLocalShuffleReads(adaptivePlan)
+ // Comet shuffle changes shuffle metrics
+ // checkNumLocalShuffleReads(adaptivePlan)
}
}
@@ -260,7 +284,8 @@ class AdaptiveQueryExecSuite
}
}
- test("Reuse the parallelism of coalesced shuffle in local shuffle read") {
+ test("Reuse the parallelism of coalesced shuffle in local shuffle read",
+ IgnoreComet("Comet shuffle changes shuffle partition size")) {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80",
@@ -292,7 +317,8 @@ class AdaptiveQueryExecSuite
}
}
- test("Reuse the default parallelism in local shuffle read") {
+ test("Reuse the default parallelism in local shuffle read",
+ IgnoreComet("Comet shuffle changes shuffle partition size")) {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80",
@@ -331,7 +357,9 @@ class AdaptiveQueryExecSuite
.groupBy($"a").count()
checkAnswer(testDf, Seq())
val plan = testDf.queryExecution.executedPlan
- assert(find(plan)(_.isInstanceOf[SortMergeJoinExec]).isDefined)
+ assert(find(plan) { case p =>
+ p.isInstanceOf[SortMergeJoinExec] || p.isInstanceOf[CometSortMergeJoinExec]
+ }.isDefined)
val coalescedReads = collect(plan) {
case r: AQEShuffleReadExec => r
}
@@ -345,7 +373,9 @@ class AdaptiveQueryExecSuite
.groupBy($"a").count()
checkAnswer(testDf, Seq())
val plan = testDf.queryExecution.executedPlan
- assert(find(plan)(_.isInstanceOf[BroadcastHashJoinExec]).isDefined)
+ assert(find(plan) { case p =>
+ p.isInstanceOf[BroadcastHashJoinExec] || p.isInstanceOf[CometBroadcastHashJoinExec]
+ }.isDefined)
val coalescedReads = collect(plan) {
case r: AQEShuffleReadExec => r
}
@@ -355,7 +385,7 @@ class AdaptiveQueryExecSuite
}
}
- test("Scalar subquery") {
+ test("Scalar subquery", IgnoreComet("Comet shuffle changes shuffle metrics")) {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
@@ -370,7 +400,7 @@ class AdaptiveQueryExecSuite
}
}
- test("Scalar subquery in later stages") {
+ test("Scalar subquery in later stages", IgnoreComet("Comet shuffle changes shuffle metrics")) {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
@@ -386,7 +416,7 @@ class AdaptiveQueryExecSuite
}
}
- test("multiple joins") {
+ test("multiple joins", IgnoreComet("Comet shuffle changes shuffle metrics")) {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
@@ -431,7 +461,7 @@ class AdaptiveQueryExecSuite
}
}
- test("multiple joins with aggregate") {
+ test("multiple joins with aggregate", IgnoreComet("Comet shuffle changes shuffle metrics")) {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
@@ -476,7 +506,7 @@ class AdaptiveQueryExecSuite
}
}
- test("multiple joins with aggregate 2") {
+ test("multiple joins with aggregate 2", IgnoreComet("Comet shuffle changes shuffle metrics")) {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") {
@@ -522,7 +552,7 @@ class AdaptiveQueryExecSuite
}
}
- test("Exchange reuse") {
+ test("Exchange reuse", IgnoreComet("Comet shuffle changes shuffle metrics")) {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
@@ -541,7 +571,7 @@ class AdaptiveQueryExecSuite
}
}
- test("Exchange reuse with subqueries") {
+ test("Exchange reuse with subqueries", IgnoreComet("Comet shuffle changes shuffle metrics")) {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80") {
@@ -572,7 +602,9 @@ class AdaptiveQueryExecSuite
assert(smj.size == 1)
val bhj = findTopLevelBroadcastHashJoin(adaptivePlan)
assert(bhj.size == 1)
- checkNumLocalShuffleReads(adaptivePlan)
+ // Comet shuffle changes shuffle metrics,
+ // so we can't check the number of local shuffle reads.
+ // checkNumLocalShuffleReads(adaptivePlan)
// Even with local shuffle read, the query stage reuse can also work.
val ex = findReusedExchange(adaptivePlan)
assert(ex.nonEmpty)
@@ -593,7 +625,9 @@ class AdaptiveQueryExecSuite
assert(smj.size == 1)
val bhj = findTopLevelBroadcastHashJoin(adaptivePlan)
assert(bhj.size == 1)
- checkNumLocalShuffleReads(adaptivePlan)
+ // Comet shuffle changes shuffle metrics,
+ // so we can't check the number of local shuffle reads.
+ // checkNumLocalShuffleReads(adaptivePlan)
// Even with local shuffle read, the query stage reuse can also work.
val ex = findReusedExchange(adaptivePlan)
assert(ex.isEmpty)
@@ -602,7 +636,8 @@ class AdaptiveQueryExecSuite
}
}
- test("Broadcast exchange reuse across subqueries") {
+ test("Broadcast exchange reuse across subqueries",
+ IgnoreComet("Comet shuffle changes shuffle metrics")) {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "20000000",
@@ -697,7 +732,8 @@ class AdaptiveQueryExecSuite
val bhj = findTopLevelBroadcastHashJoin(adaptivePlan)
assert(bhj.size == 1)
// There is still a SMJ, and its two shuffles can't apply local read.
- checkNumLocalShuffleReads(adaptivePlan, 2)
+ // Comet shuffle changes shuffle metrics
+ // checkNumLocalShuffleReads(adaptivePlan, 2)
}
}
@@ -825,7 +861,8 @@ class AdaptiveQueryExecSuite
}
}
- test("SPARK-29544: adaptive skew join with different join types") {
+ test("SPARK-29544: adaptive skew join with different join types",
+ IgnoreComet("Comet shuffle has different partition metrics")) {
Seq("SHUFFLE_MERGE", "SHUFFLE_HASH").foreach { joinHint =>
def getJoinNode(plan: SparkPlan): Seq[ShuffledJoin] = if (joinHint == "SHUFFLE_MERGE") {
findTopLevelSortMergeJoin(plan)
@@ -1269,7 +1306,8 @@ class AdaptiveQueryExecSuite
}
}
- test("metrics of the shuffle read") {
+ test("metrics of the shuffle read",
+ IgnoreComet("Comet shuffle changes the metrics")) {
withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
val (_, adaptivePlan) = runAdaptiveAndVerifyResult(
"SELECT key FROM testData GROUP BY key")
@@ -1876,13 +1914,13 @@ class AdaptiveQueryExecSuite
val (plan1, adaptivePlan1) = runAdaptiveAndVerifyResult(
"SELECT key FROM testData WHERE key = 0 ORDER BY key, value")
assert(findTopLevelSort(plan1).size == 1)
- assert(stripAQEPlan(adaptivePlan1).isInstanceOf[EmptyRelationExec])
+ assert(isEmptyRelation(adaptivePlan1), adaptivePlan1)
val (plan2, adaptivePlan2) = runAdaptiveAndVerifyResult(
"SELECT key FROM (SELECT * FROM testData WHERE value = 'no_match' ORDER BY key)" +
" WHERE key > rand()")
assert(findTopLevelSort(plan2).size == 1)
- assert(stripAQEPlan(adaptivePlan2).isInstanceOf[EmptyRelationExec])
+ assert(isEmptyRelation(adaptivePlan2), adaptivePlan2)
}
}
@@ -1890,18 +1928,18 @@ class AdaptiveQueryExecSuite
withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") {
val (plan1, adaptivePlan1) = runAdaptiveAndVerifyResult(
"SELECT key, count(*) FROM testData WHERE value = 'no_match' GROUP BY key")
- assert(!plan1.isInstanceOf[EmptyRelationExec])
- assert(stripAQEPlan(adaptivePlan1).isInstanceOf[EmptyRelationExec])
+ assert(!isEmptyRelation(plan1), plan1)
+ assert(isEmptyRelation(adaptivePlan1), adaptivePlan1)
val (plan2, adaptivePlan2) = runAdaptiveAndVerifyResult(
"SELECT key, count(*) FROM testData WHERE value = 'no_match' GROUP BY key limit 1")
- assert(!plan2.isInstanceOf[EmptyRelationExec])
- assert(stripAQEPlan(adaptivePlan2).isInstanceOf[EmptyRelationExec])
+ assert(!isEmptyRelation(plan2), plan2)
+ assert(isEmptyRelation(adaptivePlan2), adaptivePlan2)
val (plan3, adaptivePlan3) = runAdaptiveAndVerifyResult(
"SELECT count(*) FROM testData WHERE value = 'no_match'")
- assert(!plan3.isInstanceOf[EmptyRelationExec])
- assert(!stripAQEPlan(adaptivePlan3).isInstanceOf[EmptyRelationExec])
+ assert(!isEmptyRelation(plan3), plan3)
+ assert(!isEmptyRelation(adaptivePlan3), adaptivePlan3)
}
}
@@ -1922,7 +1960,7 @@ class AdaptiveQueryExecSuite
|""".stripMargin)
checkNumUnion(plan1, 1)
checkNumUnion(adaptivePlan1, 0)
- assert(!stripAQEPlan(adaptivePlan1).isInstanceOf[EmptyRelationExec])
+ assert(!isEmptyRelation(adaptivePlan1), adaptivePlan1)
val (plan2, adaptivePlan2) = runAdaptiveAndVerifyResult(
"""
@@ -1932,7 +1970,7 @@ class AdaptiveQueryExecSuite
|""".stripMargin)
checkNumUnion(plan2, 1)
checkNumUnion(adaptivePlan2, 0)
- assert(stripAQEPlan(adaptivePlan2).isInstanceOf[EmptyRelationExec])
+ assert(isEmptyRelation(adaptivePlan2), adaptivePlan2)
}
}
@@ -1944,7 +1982,7 @@ class AdaptiveQueryExecSuite
val (_, adaptivePlan) = runAdaptiveAndVerifyResult(
"SELECT id FROM v1 GROUP BY id DISTRIBUTE BY id")
assert(collect(adaptivePlan) {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
}.length == 1)
}
}
@@ -2025,7 +2063,8 @@ class AdaptiveQueryExecSuite
}
}
- test("SPARK-33551: Do not use AQE shuffle read for repartition") {
+ test("SPARK-33551: Do not use AQE shuffle read for repartition",
+ IgnoreComet("Comet shuffle changes partition size")) {
def hasRepartitionShuffle(plan: SparkPlan): Boolean = {
find(plan) {
case s: ShuffleExchangeLike =>
@@ -2212,6 +2251,9 @@ class AdaptiveQueryExecSuite
def checkNoCoalescePartitions(ds: Dataset[Row], origin: ShuffleOrigin): Unit = {
assert(collect(ds.queryExecution.executedPlan) {
case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s
+ case c: CometShuffleExchangeExec
+ if c.originalPlan.shuffleOrigin == origin &&
+ c.originalPlan.numPartitions == 2 => c
}.size == 1)
ds.collect()
val plan = ds.queryExecution.executedPlan
@@ -2220,6 +2262,9 @@ class AdaptiveQueryExecSuite
}.isEmpty)
assert(collect(plan) {
case s: ShuffleExchangeExec if s.shuffleOrigin == origin && s.numPartitions == 2 => s
+ case c: CometShuffleExchangeExec
+ if c.originalPlan.shuffleOrigin == origin &&
+ c.originalPlan.numPartitions == 2 => c
}.size == 1)
checkAnswer(ds, testData)
}
@@ -2384,7 +2429,8 @@ class AdaptiveQueryExecSuite
}
}
- test("SPARK-35264: Support AQE side shuffled hash join formula") {
+ test("SPARK-35264: Support AQE side shuffled hash join formula",
+ IgnoreComet("Comet shuffle changes the partition size")) {
withTempView("t1", "t2") {
def checkJoinStrategy(shouldShuffleHashJoin: Boolean): Unit = {
Seq("100", "100000").foreach { size =>
@@ -2470,7 +2516,8 @@ class AdaptiveQueryExecSuite
}
}
- test("SPARK-35725: Support optimize skewed partitions in RebalancePartitions") {
+ test("SPARK-35725: Support optimize skewed partitions in RebalancePartitions",
+ IgnoreComet("Comet shuffle changes shuffle metrics")) {
withTempView("v") {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
@@ -2569,7 +2616,7 @@ class AdaptiveQueryExecSuite
runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " +
s"JOIN skewData2 ON key1 = key2 GROUP BY key1")
val shuffles1 = collect(adaptive1) {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
}
assert(shuffles1.size == 3)
// shuffles1.head is the top-level shuffle under the Aggregate operator
@@ -2582,7 +2629,7 @@ class AdaptiveQueryExecSuite
runAdaptiveAndVerifyResult(s"SELECT $repartition key1 FROM skewData1 " +
s"JOIN skewData2 ON key1 = key2")
val shuffles2 = collect(adaptive2) {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
}
if (hasRequiredDistribution) {
assert(shuffles2.size == 3)
@@ -2616,7 +2663,8 @@ class AdaptiveQueryExecSuite
}
}
- test("SPARK-35794: Allow custom plugin for cost evaluator") {
+ test("SPARK-35794: Allow custom plugin for cost evaluator",
+ IgnoreComet("Comet shuffle changes shuffle metrics")) {
CostEvaluator.instantiate(
classOf[SimpleShuffleSortCostEvaluator].getCanonicalName, spark.sparkContext.getConf)
intercept[IllegalArgumentException] {
@@ -2787,6 +2835,7 @@ class AdaptiveQueryExecSuite
val (_, adaptive) = runAdaptiveAndVerifyResult(query)
assert(adaptive.collect {
case sort: SortExec => sort
+ case sort: CometSortExec => sort
}.size == 1)
val read = collect(adaptive) {
case read: AQEShuffleReadExec => read
@@ -2804,7 +2853,8 @@ class AdaptiveQueryExecSuite
}
}
- test("SPARK-37357: Add small partition factor for rebalance partitions") {
+ test("SPARK-37357: Add small partition factor for rebalance partitions",
+ IgnoreComet("Comet shuffle changes shuffle metrics")) {
withTempView("v") {
withSQLConf(
SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key -> "true",
@@ -2916,7 +2966,7 @@ class AdaptiveQueryExecSuite
runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
"JOIN skewData3 ON value2 = value3")
val shuffles1 = collect(adaptive1) {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
}
assert(shuffles1.size == 4)
val smj1 = findTopLevelSortMergeJoin(adaptive1)
@@ -2927,7 +2977,7 @@ class AdaptiveQueryExecSuite
runAdaptiveAndVerifyResult("SELECT key1 FROM skewData1 JOIN skewData2 ON key1 = key2 " +
"JOIN skewData3 ON value1 = value3")
val shuffles2 = collect(adaptive2) {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
}
assert(shuffles2.size == 4)
val smj2 = findTopLevelSortMergeJoin(adaptive2)
@@ -3185,6 +3235,7 @@ class AdaptiveQueryExecSuite
}.size == (if (firstAccess) 1 else 0))
assert(collect(initialExecutedPlan) {
case s: SortExec => s
+ case s: CometSortExec => s.originalPlan.asInstanceOf[SortExec]
}.size == (if (firstAccess) 2 else 0))
assert(collect(initialExecutedPlan) {
case i: InMemoryTableScanLike => i
@@ -3197,6 +3248,7 @@ class AdaptiveQueryExecSuite
}.isEmpty)
assert(collect(finalExecutedPlan) {
case s: SortExec => s
+ case s: CometSortExec => s.originalPlan.asInstanceOf[SortExec]
}.isEmpty)
assert(collect(initialExecutedPlan) {
case i: InMemoryTableScanLike => i
@@ -3278,7 +3330,8 @@ class AdaptiveQueryExecSuite
}
}
- test("SPARK-44040: Fix compute stats when AggregateExec nodes above QueryStageExec") {
+ test("SPARK-44040: Fix compute stats when AggregateExec nodes above QueryStageExec",
+ IgnoreComet("https://github.com/apache/datafusion-comet/issues/4412")) {
val emptyDf = spark.range(1).where("false")
val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1"))
val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2"))
@@ -3368,7 +3421,8 @@ class AdaptiveQueryExecSuite
val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec]
assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec])
- assert(plan.finalPhysicalPlan.isInstanceOf[WindowExec])
+ assert(plan.finalPhysicalPlan.isInstanceOf[WindowExec] ||
+ plan.finalPhysicalPlan.exists(_.isInstanceOf[CometWindowExec]))
plan.inputPlan.output.zip(plan.finalPhysicalPlan.output).foreach { case (o1, o2) =>
assert(o1.semanticEquals(o2), "Different output column order after AQE optimization")
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala
index 88be4adb6a4..f8fe831744e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala
@@ -228,9 +228,21 @@ class CachedBatchSerializerNoUnwrapSuite extends SharedSparkSession with Adaptiv
assert(cachedPlans.length == 2)
cachedPlans.foreach {
cachedPlan =>
- assert(cachedPlan.isInstanceOf[WholeStageCodegenExec])
- assert(cachedPlan.asInstanceOf[WholeStageCodegenExec]
- .child.isInstanceOf[ColumnarToRowExec])
+ // Comet replaces ColumnarToRowExec with its own columnar-to-row operator
+ // (CometNativeColumnarToRow / CometColumnarToRow). Accept either the
+ // Spark shape or the equivalent Comet shape, since the important property
+ // under this serializer is that the row conversion is not unwrapped.
+ // CometColumnarToRowExec supports codegen, so it appears as the child of a
+ // WholeStageCodegenExec, while CometNativeColumnarToRowExec does not and is
+ // the cached plan itself.
+ val transition = cachedPlan match {
+ case w: WholeStageCodegenExec => w.child
+ case other => other
+ }
+ val isSparkShape = cachedPlan.isInstanceOf[WholeStageCodegenExec] &&
+ transition.isInstanceOf[ColumnarToRowExec]
+ val isCometShape = transition.getClass.getName.startsWith("org.apache.spark.sql.comet.")
+ assert(isSparkShape || isCometShape, s"unexpected cached plan:\n$cachedPlan")
}
}
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
index fd8d1308e99..0e1f80045a3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala
@@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.Concat
import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
import org.apache.spark.sql.catalyst.plans.logical.Expand
import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.comet.{CometNativeScanExec, CometScanExec}
import org.apache.spark.sql.execution.FileSourceScanExec
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.functions._
@@ -884,6 +885,8 @@ abstract class SchemaPruningSuite
val fileSourceScanSchemata =
collect(df.queryExecution.executedPlan) {
case scan: FileSourceScanExec => scan.requiredSchema
+ case scan: CometScanExec => scan.requiredSchema
+ case scan: CometNativeScanExec => scan.requiredSchema
}
assert(fileSourceScanSchemata.size === expectedSchemaCatalogStrings.size,
s"Found ${fileSourceScanSchemata.size} file sources in dataframe, " +
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala
index bac19425558..4c93660a974 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala
@@ -17,9 +17,10 @@
package org.apache.spark.sql.execution.datasources
-import org.apache.spark.sql.{QueryTest, Row}
+import org.apache.spark.sql.{IgnoreComet, QueryTest, Row}
import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, NullsFirst, SortOrder}
import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Sort}
+import org.apache.spark.sql.comet.CometSortExec
import org.apache.spark.sql.execution.{QueryExecution, SortExec}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
@@ -244,6 +245,7 @@ class V1WriteCommandSuite extends SharedSparkSession with V1WriteCommandSuiteBas
// assert the outer most sort in the executed plan
assert(plan.collectFirst {
case s: SortExec => s
+ case s: CometSortExec => s.originalPlan.asInstanceOf[SortExec]
}.exists {
case SortExec(Seq(
SortOrder(AttributeReference("key", IntegerType, _, _), Ascending, NullsFirst, _),
@@ -291,6 +293,7 @@ class V1WriteCommandSuite extends SharedSparkSession with V1WriteCommandSuiteBas
// assert the outer most sort in the executed plan
assert(plan.collectFirst {
case s: SortExec => s
+ case s: CometSortExec => s.originalPlan.asInstanceOf[SortExec]
}.exists {
case SortExec(Seq(
SortOrder(AttributeReference("value", StringType, _, _), Ascending, NullsFirst, _),
@@ -324,7 +327,8 @@ class V1WriteCommandSuite extends SharedSparkSession with V1WriteCommandSuiteBas
}
}
- test("v1 write with AQE changing SMJ to BHJ") {
+ test("v1 write with AQE changing SMJ to BHJ",
+ IgnoreComet("TODO: Comet SMJ to BHJ by AQE")) {
withPlannedWrite { enabled =>
withTable("t") {
sql(
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormatSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormatSuite.scala
index 892c96e77c2..9936263a91e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormatSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormatSuite.scala
@@ -28,7 +28,7 @@ import org.apache.hadoop.fs.{FileStatus, FileSystem, GlobFilter, Path}
import org.mockito.Mockito.{mock, when}
import org.apache.spark.{SparkException, SparkUnsupportedOperationException}
-import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.{DataFrame, IgnoreCometSuite, Row}
import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
import org.apache.spark.sql.execution.datasources.PartitionedFile
import org.apache.spark.sql.functions.col
@@ -38,7 +38,9 @@ import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types._
import org.apache.spark.util.Utils
-class BinaryFileFormatSuite extends SharedSparkSession {
+// For some reason this suite is flaky w/ or w/o Comet when running in Github workflow.
+// Since it isn't related to Comet, we disable it for now.
+class BinaryFileFormatSuite extends SharedSparkSession with IgnoreCometSuite {
import BinaryFileFormat._
private var testDir: String = _
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetEncodingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetEncodingSuite.scala
index a2391a3ec22..7dfb84aff7d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetEncodingSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetEncodingSuite.scala
@@ -29,7 +29,7 @@ import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName
import org.apache.spark.TestUtils
import org.apache.spark.memory.MemoryMode
-import org.apache.spark.sql.Row
+import org.apache.spark.sql.{IgnoreComet, Row}
import org.apache.spark.sql.catalyst.util.DateTimeUtils
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
@@ -202,7 +202,8 @@ class ParquetEncodingSuite extends ParquetCompatibilityTest with SharedSparkSess
}
}
- test("parquet v2 pages - rle encoding for boolean value columns") {
+ test("parquet v2 pages - rle encoding for boolean value columns",
+ IgnoreComet("Comet doesn't support RLE encoding yet")) {
val extraOptions = Map[String, String](
ParquetOutputFormat.WRITER_VERSION -> ParquetProperties.WriterVersion.PARQUET_2_0.toString
)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
index ae977e17755..1d29f424490 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
@@ -38,6 +38,7 @@ import org.apache.parquet.schema.MessageType
import org.apache.spark.{SparkConf, SparkException, SparkRuntimeException}
import org.apache.spark.sql._
+import org.apache.spark.sql.IgnoreComet
import org.apache.spark.sql.catalyst.dsl.expressions._
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.optimizer.InferFiltersFromConstraints
@@ -1101,7 +1102,11 @@ abstract class ParquetFilterSuite extends ParquetTest with SharedSparkSession {
// When a filter is pushed to Parquet, Parquet can apply it to every row.
// So, we can check the number of rows returned from the Parquet
// to make sure our filter pushdown work.
- assert(stripSparkFilter(df).count() == 1)
+ // Similar to Spark's vectorized reader, Comet doesn't do row-level filtering but relies
+ // on Spark to apply the data filters after columnar batches are returned
+ if (!isCometEnabled) {
+ assert(stripSparkFilter(df).count() == 1)
+ }
}
}
}
@@ -1504,7 +1509,8 @@ abstract class ParquetFilterSuite extends ParquetTest with SharedSparkSession {
}
}
- test("Filters should be pushed down for vectorized Parquet reader at row group level") {
+ test("Filters should be pushed down for vectorized Parquet reader at row group level",
+ IgnoreComet("Native scans do not support the tested accumulator")) {
import testImplicits._
withSQLConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "true",
@@ -1586,7 +1592,11 @@ abstract class ParquetFilterSuite extends ParquetTest with SharedSparkSession {
// than the total length but should not be a single record.
// Note that, if record level filtering is enabled, it should be a single record.
// If no filter is pushed down to Parquet, it should be the total length of data.
- assert(actual > 1 && actual < data.length)
+ // Skip when Comet is enabled since with native execution
+ // `stripSparkFilter` can't remove the native filter
+ if (!isCometEnabled) {
+ assert(actual > 1 && actual < data.length)
+ }
}
}
}
@@ -1613,7 +1623,11 @@ abstract class ParquetFilterSuite extends ParquetTest with SharedSparkSession {
// than the total length but should not be a single record.
// Note that, if record level filtering is enabled, it should be a single record.
// If no filter is pushed down to Parquet, it should be the total length of data.
- assert(actual > 1 && actual < data.length)
+ // Skip when Comet is enabled since with native execution
+ // `stripSparkFilter` can't remove the native filter
+ if (!isCometEnabled) {
+ assert(actual > 1 && actual < data.length)
+ }
}
}
}
@@ -1705,7 +1719,7 @@ abstract class ParquetFilterSuite extends ParquetTest with SharedSparkSession {
(attr, value) => sources.StringContains(attr, value))
}
- test("filter pushdown - StringPredicate") {
+ test("filter pushdown - StringPredicate", IgnoreComet("cannot be pushed down")) {
import testImplicits._
// keep() should take effect on StartsWith/EndsWith/Contains
Seq(
@@ -1749,7 +1763,8 @@ abstract class ParquetFilterSuite extends ParquetTest with SharedSparkSession {
}
}
- test("SPARK-17091: Convert IN predicate to Parquet filter push-down") {
+ test("SPARK-17091: Convert IN predicate to Parquet filter push-down",
+ IgnoreComet("Comet has different push-down behavior")) {
val schema = StructType(Seq(
StructField("a", IntegerType, nullable = false)
))
@@ -1955,13 +1970,21 @@ abstract class ParquetFilterSuite extends ParquetTest with SharedSparkSession {
""".stripMargin)
withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") {
- val ex = intercept[SparkException] {
+ // Spark native readers wrap the error in SparkException(FAILED_READ_FILE).
+ // Comet native readers throw SparkRuntimeException directly.
+ try {
sql(s"select a from $tableName where b > 0").collect()
+ fail("Expected an exception")
+ } catch {
+ case ex: SparkException =>
+ assert(ex.getCondition.startsWith("FAILED_READ_FILE"))
+ assert(ex.getCause.isInstanceOf[SparkRuntimeException])
+ assert(ex.getCause.getMessage.contains(
+ """Found duplicate field(s) "B": [B, b] in case-insensitive mode"""))
+ case ex: SparkRuntimeException =>
+ assert(ex.getMessage.contains(
+ """Found duplicate field(s) "B": [B, b] in case-insensitive mode"""))
}
- assert(ex.getCondition.startsWith("FAILED_READ_FILE"))
- assert(ex.getCause.isInstanceOf[SparkRuntimeException])
- assert(ex.getCause.getMessage.contains(
- """Found duplicate field(s) "B": [B, b] in case-insensitive mode"""))
}
withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") {
@@ -1992,7 +2015,8 @@ abstract class ParquetFilterSuite extends ParquetTest with SharedSparkSession {
}
}
- test("Support Parquet column index") {
+ test("Support Parquet column index",
+ IgnoreComet("Comet doesn't support Parquet column index yet")) {
// block 1:
// null count min max
// page-0 0 0 99
@@ -2052,7 +2076,8 @@ abstract class ParquetFilterSuite extends ParquetTest with SharedSparkSession {
}
}
- test("SPARK-34562: Bloom filter push down") {
+ test("SPARK-34562: Bloom filter push down",
+ IgnoreComet("Native scans do not support the tested accumulator")) {
withTempPath { dir =>
val path = dir.getCanonicalPath
spark.range(100).selectExpr("id * 2 AS id")
@@ -2350,7 +2375,11 @@ class ParquetV1FilterSuite extends ParquetFilterSuite {
assert(pushedParquetFilters.exists(_.getClass === filterClass),
s"${pushedParquetFilters.map(_.getClass).toList} did not contain ${filterClass}.")
- checker(stripSparkFilter(query), expected)
+ // Similar to Spark's vectorized reader, Comet doesn't do row-level filtering but relies
+ // on Spark to apply the data filters after columnar batches are returned
+ if (!isCometEnabled) {
+ checker(stripSparkFilter(query), expected)
+ }
} else {
assert(selectedFilters.isEmpty, "There is filter pushed down")
}
@@ -2413,7 +2442,11 @@ class ParquetV2FilterSuite extends ParquetFilterSuite {
assert(pushedParquetFilters.exists(_.getClass === filterClass),
s"${pushedParquetFilters.map(_.getClass).toList} did not contain ${filterClass}.")
- checker(stripSparkFilter(query), expected)
+ // Similar to Spark's vectorized reader, Comet doesn't do row-level filtering but relies
+ // on Spark to apply the data filters after columnar batches are returned
+ if (!isCometEnabled) {
+ checker(stripSparkFilter(query), expected)
+ }
case _ => assert(false, "Can not match ParquetTable in the query.")
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
index e91562d097f..54e28f351e0 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
@@ -41,6 +41,7 @@ import org.apache.parquet.schema.{MessageType, MessageTypeParser}
import org.apache.spark.{SPARK_VERSION_SHORT, SparkException, TestUtils}
import org.apache.spark.rdd.RDD
import org.apache.spark.sql._
+import org.apache.spark.sql.IgnoreComet
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Attribute, GenericInternalRow, UnsafeRow}
import org.apache.spark.sql.catalyst.util.{DateTimeConstants, DateTimeUtils}
@@ -972,7 +973,8 @@ class ParquetIOSuite extends ParquetTest with SharedSparkSession {
}
}
- test("SPARK-54220: vectorized reader: missing all struct fields, struct with NullType only") {
+ test("SPARK-54220: vectorized reader: missing all struct fields, struct with NullType only",
+ IgnoreComet("https://github.com/apache/datafusion-comet/issues/4199")) {
val data = Seq(
Tuple1((null, null)),
Tuple1((null, null)),
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala
index 38582236c73..028ab7b7905 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala
@@ -1100,7 +1100,11 @@ abstract class ParquetQuerySuite extends ParquetTest with SharedSparkSession {
Seq(Some("A"), Some("A"), None).toDF().repartition(1)
.write.parquet(path.getAbsolutePath)
val df = spark.read.parquet(path.getAbsolutePath)
- checkAnswer(stripSparkFilter(df.where("NOT (value <=> 'A')")), df)
+ // Similar to Spark's vectorized reader, Comet doesn't do row-level filtering but relies
+ // on Spark to apply the data filters after columnar batches are returned
+ if (!isCometEnabled) {
+ checkAnswer(stripSparkFilter(df.where("NOT (value <=> 'A')")), df)
+ }
}
}
}
@@ -1164,7 +1168,8 @@ abstract class ParquetQuerySuite extends ParquetTest with SharedSparkSession {
checkAnswer(readParquet(schema2, path), df)
}
- withSQLConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "false") {
+ withSQLConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "false",
+ "spark.comet.enabled" -> "false") {
val schema1 = "a DECIMAL(3, 2), b DECIMAL(18, 3), c DECIMAL(37, 3)"
checkAnswer(readParquet(schema1, path), df)
val schema2 = "a DECIMAL(3, 0), b DECIMAL(18, 1), c DECIMAL(37, 1)"
@@ -1188,7 +1193,8 @@ abstract class ParquetQuerySuite extends ParquetTest with SharedSparkSession {
val df = sql(s"SELECT 1 a, 123456 b, ${Int.MaxValue.toLong * 10} c, CAST('1.2' AS BINARY) d")
df.write.parquet(path.toString)
- withSQLConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "false") {
+ withSQLConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "false",
+ "spark.comet.enabled" -> "false") {
checkAnswer(readParquet("a DECIMAL(3, 2)", path), sql("SELECT 1.00"))
checkAnswer(readParquet("a DECIMAL(11, 2)", path), sql("SELECT 1.00"))
checkAnswer(readParquet("b DECIMAL(3, 2)", path), Row(null))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRebaseDatetimeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRebaseDatetimeSuite.scala
index 9ad2437012a..a26ed15d0cf 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRebaseDatetimeSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRebaseDatetimeSuite.scala
@@ -21,7 +21,7 @@ import java.nio.file.{Files, Paths, StandardCopyOption}
import java.sql.{Date, Timestamp}
import org.apache.spark.{SPARK_VERSION_SHORT, SparkConf, SparkException, SparkUpgradeException}
-import org.apache.spark.sql.{Row, SPARK_LEGACY_DATETIME_METADATA_KEY, SPARK_LEGACY_INT96_METADATA_KEY, SPARK_TIMEZONE_METADATA_KEY}
+import org.apache.spark.sql.{IgnoreCometSuite, Row, SPARK_LEGACY_DATETIME_METADATA_KEY, SPARK_LEGACY_INT96_METADATA_KEY, SPARK_TIMEZONE_METADATA_KEY}
import org.apache.spark.sql.catalyst.util.DateTimeTestUtils
import org.apache.spark.sql.internal.{LegacyBehaviorPolicy, SQLConf}
import org.apache.spark.sql.internal.LegacyBehaviorPolicy.{CORRECTED, EXCEPTION, LEGACY}
@@ -30,8 +30,10 @@ import org.apache.spark.sql.internal.SQLConf.ParquetOutputTimestampType.{INT96,
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.tags.SlowSQLTest
+// Comet is disabled for this suite because it doesn't support datetime rebase mode
abstract class ParquetRebaseDatetimeSuite
extends ParquetTest
+ with IgnoreCometSuite
with SharedSparkSession {
import testImplicits._
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala
index a3f72f2a9cc..9f4adb13b78 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRowIndexSuite.scala
@@ -28,6 +28,7 @@ import org.apache.parquet.hadoop.ParquetWriter.DEFAULT_BLOCK_SIZE
import org.apache.parquet.hadoop.util.HadoopInputFile
import org.apache.spark.SparkException
+import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeScanExec}
import org.apache.spark.sql.execution.FileSourceScanExec
import org.apache.spark.sql.execution.datasources.FileFormat
import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
@@ -247,6 +248,14 @@ class ParquetRowIndexSuite extends SharedSparkSession {
case f: FileSourceScanExec =>
numPartitions += f.inputRDD.partitions.length
numOutputRows += f.metrics("numOutputRows").value
+ case b: CometBatchScanExec =>
+ numPartitions += b.inputRDD.partitions.length
+ numOutputRows += b.metrics("numOutputRows").value
+ case b: CometNativeScanExec =>
+ numPartitions +=
+ b.originalPlan.inputRDD.partitions.length
+ numOutputRows +=
+ b.metrics("output_rows").value
case _ =>
}
assert(numPartitions > 0)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaPruningSuite.scala
index 5c0b7def039..151184bc98c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaPruningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaPruningSuite.scala
@@ -20,6 +20,7 @@ package org.apache.spark.sql.execution.datasources.parquet
import org.apache.spark.SparkConf
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
+import org.apache.spark.sql.comet.CometBatchScanExec
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.datasources.SchemaPruningSuite
import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
@@ -56,6 +57,7 @@ class ParquetV2SchemaPruningSuite extends ParquetSchemaPruningSuite {
val fileSourceScanSchemata =
collect(df.queryExecution.executedPlan) {
case scan: BatchScanExec => scan.scan.asInstanceOf[ParquetScan].readDataSchema
+ case scan: CometBatchScanExec => scan.scan.asInstanceOf[ParquetScan].readDataSchema
}
assert(fileSourceScanSchemata.size === expectedSchemaCatalogStrings.size,
s"Found ${fileSourceScanSchemata.size} file sources in dataframe, " +
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaSuite.scala
index c8d866d8b2d..aacbd48ca1d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaSuite.scala
@@ -27,7 +27,7 @@ import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName
import org.apache.parquet.schema.Type._
import org.apache.spark.SparkException
-import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.{AnalysisException, IgnoreComet, Row}
import org.apache.spark.sql.catalyst.expressions.Cast.toSQLType
import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException
import org.apache.spark.sql.functions.desc
@@ -1037,7 +1037,8 @@ class ParquetSchemaSuite extends ParquetSchemaTest {
e
}
- test("schema mismatch failure error message for parquet reader") {
+ test("schema mismatch failure error message for parquet reader",
+ IgnoreComet("Comet doesn't work with vectorizedReaderEnabled = false")) {
withTempPath { dir =>
val e = testSchemaMismatch(dir.getCanonicalPath, vectorizedReaderEnabled = false)
val expectedMessage = "Encountered error while reading file"
@@ -1046,7 +1047,8 @@ class ParquetSchemaSuite extends ParquetSchemaTest {
}
}
- test("schema mismatch failure error message for parquet vectorized reader") {
+ test("schema mismatch failure error message for parquet vectorized reader",
+ IgnoreComet("https://github.com/apache/datafusion-comet/issues/4316")) {
withTempPath { dir =>
val e = testSchemaMismatch(dir.getCanonicalPath, vectorizedReaderEnabled = true)
assert(e.getCause.isInstanceOf[SchemaColumnConvertNotSupportedException])
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTypeWideningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTypeWideningSuite.scala
index 7ccd664f6c7..0f16c2f2fad 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTypeWideningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTypeWideningSuite.scala
@@ -190,7 +190,8 @@ class ParquetTypeWideningSuite
(Seq("1", "2", Short.MinValue.toString), ShortType, DoubleType),
(Seq("1", "2", Int.MinValue.toString), IntegerType, DoubleType),
(Seq("1.23", "10.34"), FloatType, DoubleType),
- (Seq("2020-01-01", "2020-01-02", "1312-02-27"), DateType, TimestampNTZType)
+ // TODO: Comet cannot handle older than "1582-10-15"
+ (Seq("2020-01-01", "2020-01-02"/* , "1312-02-27" */), DateType, TimestampNTZType)
)
}
test(s"parquet widening conversion $fromType -> $toType") {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetVariantShreddingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetVariantShreddingSuite.scala
index 5c0739b6276..ab12b78b7db 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetVariantShreddingSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetVariantShreddingSuite.scala
@@ -29,7 +29,7 @@ import org.apache.parquet.schema.{LogicalTypeAnnotation, PrimitiveType, Type}
import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName
import org.apache.spark.SparkException
-import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.{AnalysisException, IgnoreComet, Row}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.internal.SQLConf.ParquetOutputTimestampType
import org.apache.spark.sql.test.SharedSparkSession
@@ -230,7 +230,8 @@ class ParquetVariantShreddingSuite extends ParquetTest with SharedSparkSession {
}
}
- test("variant logical type annotation - ignore variant annotation") {
+ test("variant logical type annotation - ignore variant annotation",
+ IgnoreComet("https://github.com/apache/datafusion-comet/issues/5741")) {
Seq(true, false).foreach { ignoreVariantAnnotation =>
withSQLConf(SQLConf.PARQUET_ANNOTATE_VARIANT_LOGICAL_TYPE.key -> "true",
SQLConf.PARQUET_IGNORE_VARIANT_ANNOTATION.key -> ignoreVariantAnnotation.toString,
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/debug/DebuggingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/debug/DebuggingSuite.scala
index b8f3ea3c6f3..bbd44221288 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/debug/DebuggingSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/debug/DebuggingSuite.scala
@@ -20,6 +20,7 @@ package org.apache.spark.sql.execution.debug
import java.io.ByteArrayOutputStream
import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.IgnoreComet
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.Attribute
import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext
@@ -125,7 +126,8 @@ class DebuggingSuite extends DebuggingSuiteBase with DisableAdaptiveExecutionSui
| id LongType: {}""".stripMargin))
}
- test("SPARK-28537: DebugExec cannot debug columnar related queries") {
+ test("SPARK-28537: DebugExec cannot debug columnar related queries",
+ IgnoreComet("Comet does not use FileScan")) {
withTempPath { workDir =>
val workDirPath = workDir.getAbsolutePath
val input = spark.range(5).toDF("id")
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala
index 52746720eba..15df5e78440 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala
@@ -21,7 +21,7 @@ import scala.reflect.ClassTag
import org.apache.spark.AccumulatorSuite
import org.apache.spark.internal.config.EXECUTOR_MEMORY
-import org.apache.spark.sql.{QueryTest, Row}
+import org.apache.spark.sql.{IgnoreComet, QueryTest, Row}
import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, Cast, Expression, Literal, ShiftLeft}
import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide}
import org.apache.spark.sql.catalyst.plans.Inner
@@ -487,7 +487,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest
}
}
- test("broadcast join where streamed side's output partitioning is PartitioningCollection") {
+ test("broadcast join where streamed side's output partitioning is PartitioningCollection",
+ IgnoreComet("Comet replaces the join and shuffle operators this test inspects")) {
withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") {
val t1 = (0 until 100).map(i => (i % 5, i % 13)).toDF("i1", "j1")
val t2 = (0 until 100).map(i => (i % 5, i % 14)).toDF("i2", "j2")
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala
index 2ae2ea6339f..30c8959a057 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.execution.metric
import scala.util.Random
import org.apache.spark.internal.config
-import org.apache.spark.sql.{Column, Dataset}
+import org.apache.spark.sql.{Column, Dataset, IgnoreCometSuite}
import org.apache.spark.sql.execution.adaptive.{AQETestHelper, DisableAdaptiveExecutionSuite}
import org.apache.spark.sql.functions.udf
import org.apache.spark.sql.internal.SQLConf
@@ -29,7 +29,10 @@ class MetricsFailureInjectionSuite
extends SharedSparkSession
with SQLMetricsTestUtils
// Need to control AQE per-test to ensure expected plan shapes.
- with DisableAdaptiveExecutionSuite {
+ with DisableAdaptiveExecutionSuite
+ // Comet replaces Spark operators with native equivalents, so the SQL Last Attempt Metric
+ // is not propagated through the plan. See https://github.com/apache/datafusion-comet/issues/4963
+ with IgnoreCometSuite {
import testImplicits._
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricIntegrationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricIntegrationSuite.scala
index 2e7af075a3e..b8fa7d35783 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricIntegrationSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricIntegrationSuite.scala
@@ -17,7 +17,7 @@
package org.apache.spark.sql.execution.metric
import org.apache.spark.internal.config
-import org.apache.spark.sql.Column
+import org.apache.spark.sql.{Column, IgnoreComet}
import org.apache.spark.sql.catalyst.expressions.Literal
import org.apache.spark.sql.catalyst.optimizer.BuildRight
import org.apache.spark.sql.catalyst.plans.RightOuter
@@ -255,7 +255,10 @@ class SQLLastAttemptMetricIntegrationSuite
logInfo(slam.logAccumulatorState)
}
- test("dataset updates") {
+ test("dataset updates",
+ // Comet replaces Spark operators with native equivalents, so the SQL Last Attempt Metric
+ // is not propagated through the plan.
+ IgnoreComet("https://github.com/apache/datafusion-comet/issues/4963")) {
val slam1 = SQLLastAttemptMetrics.createMetric(spark.sparkContext, "SLAM1")
val df1 = spark.range(10).filter(Column(incrementMetric(slam1)))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricPlanShapesSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricPlanShapesSuite.scala
index ea8d9568f7e..34b9c2eef01 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricPlanShapesSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricPlanShapesSuite.scala
@@ -22,6 +22,7 @@ import scala.util.Random
import org.scalatest.Tag
import org.apache.spark.internal.config
+import org.apache.spark.sql.IgnoreCometSuite
import org.apache.spark.sql.execution.{CollectLimitExec, RDDScanExec, SparkPlan}
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQETestHelper, DisableAdaptiveExecutionSuite}
import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
@@ -34,7 +35,10 @@ class SQLLastAttemptMetricPlanShapesSuite
extends SharedSparkSession
with SQLMetricsTestUtils
// Need to control AQE per-test to ensure expected plan shapes.
- with DisableAdaptiveExecutionSuite {
+ with DisableAdaptiveExecutionSuite
+ // Comet replaces Spark operators with native equivalents, so the SQL Last Attempt Metric
+ // is not propagated through the plan. See https://github.com/apache/datafusion-comet/issues/4963
+ with IgnoreCometSuite {
import testImplicits._
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
index 16cde3ef2c5..cfc26410733 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
@@ -46,8 +46,10 @@ import org.apache.spark.sql.util.QueryExecutionListener
import org.apache.spark.util.{AccumulatorContext, JsonProtocol}
// Disable AQE because metric info is different with AQE on/off
+// This test suite runs tests against the metrics of physical operators.
+// Disabling it for Comet because the metrics are different with Comet enabled.
class SQLMetricsSuite extends SharedSparkSession with SQLMetricsTestUtils
- with DisableAdaptiveExecutionSuite {
+ with DisableAdaptiveExecutionSuite with IgnoreCometSuite {
import testImplicits._
/**
@@ -857,7 +859,8 @@ class SQLMetricsSuite extends SharedSparkSession with SQLMetricsTestUtils
}
}
- test("SPARK-26327: FileSourceScanExec metrics") {
+ test("SPARK-26327: FileSourceScanExec metrics",
+ IgnoreComet("Spark uses row-based Parquet reader while Comet is vectorized")) {
withTable("testDataForScan") {
spark.range(10).selectExpr("id", "id % 3 as p")
.write.partitionBy("p").saveAsTable("testDataForScan")
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFsSuite.scala
index 4c4273006b9..3386a70345c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFsSuite.scala
@@ -18,6 +18,7 @@
package org.apache.spark.sql.execution.python
import org.apache.spark.sql.catalyst.plans.logical.{ArrowEvalPython, BatchEvalPython, Limit, LocalLimit}
+import org.apache.spark.sql.comet._
import org.apache.spark.sql.execution.{FileSourceScanExec, SparkPlan}
import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan
@@ -108,6 +109,8 @@ class ExtractPythonUDFsSuite extends SharedSparkSession {
val scanNodes = query.queryExecution.executedPlan.collect {
case scan: FileSourceScanExec => scan
+ case scan: CometScanExec => scan
+ case scan: CometNativeScanExec => scan
}
assert(scanNodes.length == 1)
assert(scanNodes.head.output.map(_.name) == Seq("a"))
@@ -120,11 +123,18 @@ class ExtractPythonUDFsSuite extends SharedSparkSession {
val scanNodes = query.queryExecution.executedPlan.collect {
case scan: FileSourceScanExec => scan
+ case scan: CometScanExec => scan
+ case scan: CometNativeScanExec => scan
}
assert(scanNodes.length == 1)
// $"a" is not null and $"a" > 1
- assert(scanNodes.head.dataFilters.length == 2)
- assert(scanNodes.head.dataFilters.flatMap(_.references.map(_.name)).distinct == Seq("a"))
+ val dataFilters = scanNodes.head match {
+ case scan: FileSourceScanExec => scan.dataFilters
+ case scan: CometScanExec => scan.dataFilters
+ case scan: CometNativeScanExec => scan.dataFilters
+ }
+ assert(dataFilters.length == 2)
+ assert(dataFilters.flatMap(_.references.map(_.name)).distinct == Seq("a"))
}
}
}
@@ -145,6 +155,7 @@ class ExtractPythonUDFsSuite extends SharedSparkSession {
val scanNodes = query.queryExecution.executedPlan.collect {
case scan: BatchScanExec => scan
+ case scan: CometBatchScanExec => scan
}
assert(scanNodes.length == 1)
assert(scanNodes.head.output.map(_.name) == Seq("a"))
@@ -157,6 +168,7 @@ class ExtractPythonUDFsSuite extends SharedSparkSession {
val scanNodes = query.queryExecution.executedPlan.collect {
case scan: BatchScanExec => scan
+ case scan: CometBatchScanExec => scan
}
assert(scanNodes.length == 1)
// $"a" is not null and $"a" > 1
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/AsyncProgressTrackingMicroBatchExecutionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/AsyncProgressTrackingMicroBatchExecutionSuite.scala
index 91f6237d896..e9e63f594d4 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/AsyncProgressTrackingMicroBatchExecutionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/AsyncProgressTrackingMicroBatchExecutionSuite.scala
@@ -38,8 +38,10 @@ import org.apache.spark.sql.streaming.{StreamingQuery, StreamingQueryException,
import org.apache.spark.sql.streaming.util.StreamManualClock
import org.apache.spark.util.Utils
+// For some reason this suite is flaky w/ or w/o Comet when running in Github workflow.
+// Since it isn't related to Comet, we disable it for now.
class AsyncProgressTrackingMicroBatchExecutionSuite
- extends StreamTest with BeforeAndAfter with Matchers {
+ extends StreamTest with BeforeAndAfter with Matchers with IgnoreCometSuite {
import testImplicits._
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBCheckpointFailureInjectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBCheckpointFailureInjectionSuite.scala
index b0220c33962..ba057b8b753 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBCheckpointFailureInjectionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBCheckpointFailureInjectionSuite.scala
@@ -24,6 +24,7 @@ import scala.language.implicitConversions
import org.apache.hadoop.conf.Configuration
import org.apache.spark.{SparkConf, SparkException}
+import org.apache.spark.sql.IgnoreCometSuite
import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
import org.apache.spark.sql.functions.{count, udf}
import org.apache.spark.sql.internal.SQLConf
@@ -37,8 +38,11 @@ import org.apache.spark.util.Utils
@SlowSQLTest
/** Test suite to inject some failures in RocksDB checkpoint */
+// Comet: skip this suite under Comet. Comet falls back to Spark for every streaming plan
+// (see CometExecRule and CometScanRule), so these tests exercise no Comet code path.
+// See https://github.com/apache/datafusion-comet/issues/5986
class RocksDBCheckpointFailureInjectionSuite extends StreamTest
- with SharedSparkSession {
+ with SharedSparkSession with IgnoreCometSuite {
private val fileManagerClassName = classOf[FailureInjectionCheckpointFileManager].getName
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBLineageSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBLineageSuite.scala
index 48ef4158266..589c675773f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBLineageSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBLineageSuite.scala
@@ -20,10 +20,14 @@ package org.apache.spark.sql.execution.streaming.state
import org.apache.hadoop.conf.Configuration
import org.apache.spark.SparkException
+import org.apache.spark.sql.IgnoreCometSuite
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.util.Utils
-class RocksDBLineageSuite extends SharedSparkSession {
+// Comet: skip this suite under Comet. The tests target RocksDB lineage internals, not SQL
+// execution paths, so nothing here runs through Comet.
+// See https://github.com/apache/datafusion-comet/issues/5986
+class RocksDBLineageSuite extends SharedSparkSession with IgnoreCometSuite {
private def newDB(root: String, enableCheckpointIds: Boolean): RocksDB = {
val conf = RocksDBConf().copy(enableChangelogCheckpointing = true)
new RocksDB(
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala
index fe09506023c..15034b3f0f6 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreCheckpointFormatV2Suite.scala
@@ -25,7 +25,7 @@ import org.apache.hadoop.fs.Path
import org.scalatest.Tag
import org.apache.spark.{SparkContext, SparkException, TaskContext}
-import org.apache.spark.sql.{DataFrame, ForeachWriter}
+import org.apache.spark.sql.{DataFrame, ForeachWriter, IgnoreCometSuite}
import org.apache.spark.sql.catalyst.expressions.UnsafeRow
import org.apache.spark.sql.execution.streaming.checkpointing.CommitLog
import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream, StreamExecution}
@@ -275,8 +275,14 @@ class CkptIdCollectingStateStoreProviderWrapper extends StateStoreProvider {
innerProvider.supportedInstanceMetrics
}
+// Comet: skip this suite under Comet. Comet falls back to Spark for every streaming plan
+// (see CometExecRule and CometScanRule), so these tests exercise no Comet code path.
+// IgnoreCometSuite is listed before AlsoTestWithRocksDBFeatures on purpose: that trait
+// registers most tests through `super.test`, which only reaches the override when the
+// override sits after it in the linearization.
+// See https://github.com/apache/datafusion-comet/issues/5986
class RocksDBStateStoreCheckpointFormatV2Suite extends StreamTest
- with AlsoTestWithRocksDBFeatures {
+ with IgnoreCometSuite with AlsoTestWithRocksDBFeatures {
import testImplicits._
val providerClassName = classOf[CkptIdCollectingStateStoreProviderWrapper].getCanonicalName
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala
index 4060e44ef1a..df9285dd966 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala
@@ -24,6 +24,7 @@ import scala.jdk.CollectionConverters.SetHasAsScala
import org.scalatest.time.{Millis, Minute, Seconds, Span}
import org.apache.spark.memory.UnifiedMemoryManager
+import org.apache.spark.sql.IgnoreCometSuite
import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream, StreamingQueryWrapper}
import org.apache.spark.sql.functions.{count, max}
import org.apache.spark.sql.internal.SQLConf
@@ -31,8 +32,14 @@ import org.apache.spark.sql.streaming._
import org.apache.spark.sql.streaming.OutputMode.Update
import org.apache.spark.util.Utils
+// Comet: skip this suite under Comet. Comet falls back to Spark for every streaming plan
+// (see CometExecRule and CometScanRule), so these tests exercise no Comet code path.
+// IgnoreCometSuite is listed before AlsoTestWithRocksDBFeatures on purpose: that trait
+// registers most tests through `super.test`, which only reaches the override when the
+// override sits after it in the linearization.
+// See https://github.com/apache/datafusion-comet/issues/5986
class RocksDBStateStoreIntegrationSuite extends StreamTest
- with AlsoTestWithRocksDBFeatures {
+ with IgnoreCometSuite with AlsoTestWithRocksDBFeatures {
import testImplicits._
testWithColumnFamilies("RocksDBStateStore",
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
index ed8c22740b3..3307ef97b47 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreSuite.scala
@@ -35,6 +35,7 @@ import org.scalatest.time.SpanSugar._
import org.apache.spark.{SparkConf, SparkException, SparkRuntimeException, SparkUnsupportedOperationException, TaskContext}
import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.IgnoreCometSuite
import org.apache.spark.sql.LocalSparkSession.withSparkSession
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.InternalRow
@@ -52,7 +53,15 @@ import org.apache.spark.unsafe.types.UTF8String
import org.apache.spark.util.{ThreadUtils, Utils}
@ExtendedSQLTest
+// Comet: skip this suite under Comet. The tests target RocksDB state-store internals
+// (StateStore get/put/commit), not SQL execution paths. StateStoreSuite, the HDFS-backed
+// twin of this suite, is skipped for the same reason.
+// IgnoreCometSuite is listed before the AlsoTestWith* traits on purpose: those traits
+// register most tests through `super.test`, which only reaches the override when the
+// override sits after them in the linearization.
+// See https://github.com/apache/datafusion-comet/issues/5986
class RocksDBStateStoreSuite extends StateStoreSuiteBase[RocksDBStateStoreProvider]
+ with IgnoreCometSuite
with AlsoTestWithEncodingTypes
with AlsoTestWithRocksDBFeatures
with PrivateMethodTester
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala
index dc697f5b99d..2b68513ca0f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala
@@ -38,6 +38,7 @@ import org.scalatest.Tag
import org.apache.spark.{SparkConf, SparkException, SparkFunSuite, SparkIllegalArgumentException, TaskContext}
import org.apache.spark.internal.Logging
import org.apache.spark.io.CompressionCodec
+import org.apache.spark.sql.IgnoreCometSuite
import org.apache.spark.sql.QueryTest
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.UnsafeProjection
@@ -733,8 +734,14 @@ class RocksDBStateEncoderSuite extends SparkFunSuite {
}
@SlowSQLTest
-class RocksDBSuite extends AlsoTestWithRocksDBFeatures with SharedSparkSession
- with PrivateMethodTester {
+// Comet: skip this suite under Comet. The tests target RocksDB internals (open, checkpoint,
+// changelog replay), not SQL execution paths, so nothing here runs through Comet.
+// IgnoreCometSuite is listed before AlsoTestWithRocksDBFeatures on purpose: that trait
+// registers most tests through `super.test`, which only reaches the override when the
+// override sits after it in the linearization.
+// See https://github.com/apache/datafusion-comet/issues/5986
+class RocksDBSuite extends IgnoreCometSuite with AlsoTestWithRocksDBFeatures
+ with SharedSparkSession with PrivateMethodTester {
override protected def sparkConf: SparkConf = {
super.sparkConf
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala
index 65009a2c1c6..09c53621daf 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala
@@ -32,7 +32,8 @@ import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.fs._
import org.json4s.DefaultFormats
import org.json4s.jackson.JsonMethods
-import org.scalatest.{BeforeAndAfter, PrivateMethodTester}
+import org.scalactic.source.Position
+import org.scalatest.{BeforeAndAfter, PrivateMethodTester, Tag}
import org.scalatest.concurrent.Eventually._
import org.scalatest.time.SpanSugar._
@@ -41,6 +42,7 @@ import org.apache.spark.LocalSparkContext._
import org.apache.spark.internal.Logging
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeProjection, UnsafeRow}
+import org.apache.spark.sql.classic.{SparkSession => ClassicSparkSession}
import org.apache.spark.sql.catalyst.util.quietly
import org.apache.spark.sql.execution.streaming._
import org.apache.spark.sql.execution.streaming.checkpointing.CheckpointFileManager
@@ -291,6 +293,18 @@ class StateStoreSuite extends StateStoreSuiteBase[HDFSBackedStateStoreProvider]
import StateStoreTestsHelper._
import StateStoreCoordinatorSuite._
+ // Comet: skip this suite under Comet. The tests target streaming state-store internals
+ // (StateStore.get/put/commit), not SQL execution paths, and the `maintenance` test is
+ // flaky in CI. See https://github.com/apache/datafusion-comet/issues/4221
+ override protected def test(testName: String, testTags: Tag*)(testFun: => Any)
+ (implicit pos: Position): Unit = {
+ if (ClassicSparkSession.isCometEnabled) {
+ ignore(testName + " (disabled when Comet is on)", testTags: _*)(testFun)
+ } else {
+ super.test(testName, testTags: _*)(testFun)
+ }
+ }
+
before {
StateStore.stop()
require(!StateStore.isMaintenanceRunning)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowMetricsSuite.scala
index 5c3edb7e06c..f6accc654cd 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowMetricsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowMetricsSuite.scala
@@ -17,7 +17,7 @@
package org.apache.spark.sql.execution.window
-import org.apache.spark.sql.QueryTest
+import org.apache.spark.sql.{IgnoreCometSuite, QueryTest}
import org.apache.spark.sql.execution.SparkPlanInfo
import org.apache.spark.sql.execution.metric.SQLMetricsTestUtils
import org.apache.spark.sql.execution.ui.SparkPlanGraph
@@ -31,8 +31,11 @@ import org.apache.spark.sql.test.SharedSparkSession
* segtree path bumps `numSegmentTreeFrames`; fallback path bumps
* `numSegmentTreeFallbackFrames`; feature-flag off leaves both at 0.
*/
+// Comet replaces Spark's WindowExec with CometWindowExec, which has no segment-tree window
+// optimization and therefore does not expose the segment-tree frame metrics this suite reads.
+// See https://github.com/apache/datafusion-comet/issues/4964
class SegmentTreeWindowMetricsSuite
- extends QueryTest with SharedSparkSession with SQLMetricsTestUtils {
+ extends QueryTest with SharedSparkSession with SQLMetricsTestUtils with IgnoreCometSuite {
import testImplicits._
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreeAllowlistSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreeAllowlistSuite.scala
index d3da8e40c85..505f48092ea 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreeAllowlistSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreeAllowlistSuite.scala
@@ -17,7 +17,7 @@
package org.apache.spark.sql.execution.window
-import org.apache.spark.sql.{DataFrame, QueryTest}
+import org.apache.spark.sql.{DataFrame, IgnoreCometSuite, QueryTest}
import org.apache.spark.sql.execution.SparkPlanInfo
import org.apache.spark.sql.execution.metric.SQLMetricsTestUtils
import org.apache.spark.sql.execution.ui.SparkPlanGraph
@@ -35,8 +35,11 @@ import org.apache.spark.sql.test.SharedSparkSession
* Eligibility gating only; exhaustive equivalence lives in
* [[SegmentTreeWindowFunctionSuite]].
*/
+// Comet replaces Spark's WindowExec with CometWindowExec, which has no segment-tree window
+// optimization and therefore does not expose the segment-tree frame metrics this suite reads.
+// See https://github.com/apache/datafusion-comet/issues/4964
class WindowSegmentTreeAllowlistSuite
- extends QueryTest with SharedSparkSession with SQLMetricsTestUtils {
+ extends QueryTest with SharedSparkSession with SQLMetricsTestUtils with IgnoreCometSuite {
import testImplicits._
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala
index 452b7b9dedb..b2a73a8b772 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala
@@ -26,10 +26,11 @@ import org.apache.spark.sql.catalyst.expressions
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning
import org.apache.spark.sql.catalyst.types.DataTypeUtils
-import org.apache.spark.sql.execution.{FileSourceScanExec, SortExec, SparkPlan}
+import org.apache.spark.sql.comet._
+import org.apache.spark.sql.execution.{ColumnarToRowExec, FileSourceScanExec, SortExec, SparkPlan}
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper}
import org.apache.spark.sql.execution.datasources.BucketingUtils
-import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
+import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec, ShuffleExchangeLike}
import org.apache.spark.sql.execution.joins.SortMergeJoinExec
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
@@ -103,12 +104,22 @@ abstract class BucketedReadSuite extends QueryTest with AdaptiveSparkPlanHelper
}
}
- private def getFileScan(plan: SparkPlan): FileSourceScanExec = {
- val fileScan = collect(plan) { case f: FileSourceScanExec => f }
+ private def getFileScan(plan: SparkPlan): SparkPlan = {
+ val fileScan = collect(plan) {
+ case f: FileSourceScanExec => f
+ case f: CometScanExec => f
+ case f: CometNativeScanExec => f
+ }
assert(fileScan.nonEmpty, plan)
fileScan.head
}
+ private def getBucketScan(plan: SparkPlan): Boolean = getFileScan(plan) match {
+ case fs: FileSourceScanExec => fs.bucketedScan
+ case bs: CometScanExec => bs.bucketedScan
+ case ns: CometNativeScanExec => ns.bucketedScan
+ }
+
// To verify if the bucket pruning works, this function checks two conditions:
// 1) Check if the pruned buckets (before filtering) are empty.
// 2) Verify the final result is the same as the expected one
@@ -157,7 +168,8 @@ abstract class BucketedReadSuite extends QueryTest with AdaptiveSparkPlanHelper
val planWithoutBucketedScan = bucketedDataFrame.filter(filterCondition)
.queryExecution.executedPlan
val fileScan = getFileScan(planWithoutBucketedScan)
- assert(!fileScan.bucketedScan, s"except no bucketed scan but found\n$fileScan")
+ val bucketedScan = getBucketScan(planWithoutBucketedScan)
+ assert(!bucketedScan, s"except no bucketed scan but found\n$fileScan")
val bucketColumnType = bucketedDataFrame.schema.apply(bucketColumnIndex).dataType
val rowsWithInvalidBuckets = fileScan.execute().filter(row => {
@@ -454,28 +466,54 @@ abstract class BucketedReadSuite extends QueryTest with AdaptiveSparkPlanHelper
val joinOperator = if (joined.sparkSession.sessionState.conf.adaptiveExecutionEnabled) {
val executedPlan =
joined.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan
- assert(executedPlan.isInstanceOf[SortMergeJoinExec])
- executedPlan.asInstanceOf[SortMergeJoinExec]
+ executedPlan match {
+ case s: SortMergeJoinExec => s
+ case b: CometSortMergeJoinExec =>
+ b.originalPlan match {
+ case s: SortMergeJoinExec => s
+ case o => fail(s"expected SortMergeJoinExec, but found\n$o")
+ }
+ case o => fail(s"expected SortMergeJoinExec, but found\n$o")
+ }
} else {
val executedPlan = joined.queryExecution.executedPlan
- assert(executedPlan.isInstanceOf[SortMergeJoinExec])
- executedPlan.asInstanceOf[SortMergeJoinExec]
+ executedPlan match {
+ case s: SortMergeJoinExec => s
+ case ColumnarToRowExec(child) =>
+ child.asInstanceOf[CometSortMergeJoinExec].originalPlan match {
+ case s: SortMergeJoinExec => s
+ case o => fail(s"expected SortMergeJoinExec, but found\n$o")
+ }
+ case CometColumnarToRowExec(child) =>
+ child.asInstanceOf[CometSortMergeJoinExec].originalPlan match {
+ case s: SortMergeJoinExec => s
+ case o => fail(s"expected SortMergeJoinExec, but found\n$o")
+ }
+ case CometNativeColumnarToRowExec(child) =>
+ child.asInstanceOf[CometSortMergeJoinExec].originalPlan match {
+ case s: SortMergeJoinExec => s
+ case o => fail(s"expected SortMergeJoinExec, but found\n$o")
+ }
+ case o => fail(s"expected SortMergeJoinExec, but found\n$o")
+ }
}
// check existence of shuffle
assert(
- joinOperator.left.exists(_.isInstanceOf[ShuffleExchangeExec]) == shuffleLeft,
+ joinOperator.left.exists(op => op.isInstanceOf[ShuffleExchangeLike]) == shuffleLeft,
s"expected shuffle in plan to be $shuffleLeft but found\n${joinOperator.left}")
assert(
- joinOperator.right.exists(_.isInstanceOf[ShuffleExchangeExec]) == shuffleRight,
+ joinOperator.right.exists(op => op.isInstanceOf[ShuffleExchangeLike]) == shuffleRight,
s"expected shuffle in plan to be $shuffleRight but found\n${joinOperator.right}")
// check existence of sort
assert(
- joinOperator.left.exists(_.isInstanceOf[SortExec]) == sortLeft,
+ joinOperator.left.exists(op => op.isInstanceOf[SortExec] || op.isInstanceOf[CometExec] &&
+ op.asInstanceOf[CometExec].originalPlan.isInstanceOf[SortExec]) == sortLeft,
s"expected sort in the left child to be $sortLeft but found\n${joinOperator.left}")
assert(
- joinOperator.right.exists(_.isInstanceOf[SortExec]) == sortRight,
+ joinOperator.right.exists(op => op.isInstanceOf[SortExec] || op.isInstanceOf[CometExec] &&
+ op.asInstanceOf[CometExec].originalPlan.isInstanceOf[SortExec]) == sortRight,
s"expected sort in the right child to be $sortRight but found\n${joinOperator.right}")
// check the output partitioning
@@ -838,11 +876,11 @@ abstract class BucketedReadSuite extends QueryTest with AdaptiveSparkPlanHelper
df1.write.format("parquet").bucketBy(8, "i").saveAsTable("bucketed_table")
val scanDF = spark.table("bucketed_table").select("j")
- assert(!getFileScan(scanDF.queryExecution.executedPlan).bucketedScan)
+ assert(!getBucketScan(scanDF.queryExecution.executedPlan))
checkAnswer(scanDF, df1.select("j"))
val aggDF = spark.table("bucketed_table").groupBy("j").agg(max("k"))
- assert(!getFileScan(aggDF.queryExecution.executedPlan).bucketedScan)
+ assert(!getBucketScan(aggDF.queryExecution.executedPlan))
checkAnswer(aggDF, df1.groupBy("j").agg(max("k")))
}
}
@@ -1031,15 +1069,24 @@ abstract class BucketedReadSuite extends QueryTest with AdaptiveSparkPlanHelper
Seq(true, false).foreach { aqeEnabled =>
withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString) {
val plan = sql(query).queryExecution.executedPlan
- val shuffles = collect(plan) { case s: ShuffleExchangeExec => s }
+ val shuffles = collect(plan) { case s: ShuffleExchangeLike => s }
assert(shuffles.length == expectedNumShuffles)
val scans = collect(plan) {
case f: FileSourceScanExec if f.optionalNumCoalescedBuckets.isDefined => f
+ case b: CometScanExec if b.optionalNumCoalescedBuckets.isDefined => b
+ case b: CometNativeScanExec if b.optionalNumCoalescedBuckets.isDefined => b
}
if (expectedCoalescedNumBuckets.isDefined) {
assert(scans.length == 1)
- assert(scans.head.optionalNumCoalescedBuckets == expectedCoalescedNumBuckets)
+ scans.head match {
+ case f: FileSourceScanExec =>
+ assert(f.optionalNumCoalescedBuckets == expectedCoalescedNumBuckets)
+ case b: CometScanExec =>
+ assert(b.optionalNumCoalescedBuckets == expectedCoalescedNumBuckets)
+ case b: CometNativeScanExec =>
+ assert(b.optionalNumCoalescedBuckets == expectedCoalescedNumBuckets)
+ }
} else {
assert(scans.isEmpty)
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/CreateTableAsSelectSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/CreateTableAsSelectSuite.scala
index 95c2fcbd7b5..e2d4a20c5d9 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/sources/CreateTableAsSelectSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/CreateTableAsSelectSuite.scala
@@ -20,6 +20,7 @@ package org.apache.spark.sql.sources
import java.io.File
import org.apache.spark.SparkException
+import org.apache.spark.sql.IgnoreCometSuite
import org.apache.spark.sql.catalyst.TableIdentifier
import org.apache.spark.sql.catalyst.catalog.{BucketSpec, CatalogTableType}
import org.apache.spark.sql.catalyst.parser.ParseException
@@ -27,7 +28,10 @@ import org.apache.spark.sql.internal.SQLConf.BUCKETING_MAX_BUCKETS
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.util.Utils
-class CreateTableAsSelectSuite extends DataSourceTest with SharedSparkSession {
+// For some reason this suite is flaky w/ or w/o Comet when running in Github workflow.
+// Since it isn't related to Comet, we disable it for now.
+class CreateTableAsSelectSuite extends DataSourceTest with SharedSparkSession
+ with IgnoreCometSuite {
import testImplicits._
protected override lazy val sql = spark.sql _
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/DisableUnnecessaryBucketedScanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/DisableUnnecessaryBucketedScanSuite.scala
index fddcab525e3..29a8629943d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/sources/DisableUnnecessaryBucketedScanSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/DisableUnnecessaryBucketedScanSuite.scala
@@ -18,6 +18,7 @@
package org.apache.spark.sql.sources
import org.apache.spark.sql.QueryTest
+import org.apache.spark.sql.comet.{CometNativeScanExec, CometScanExec}
import org.apache.spark.sql.execution.FileSourceScanExec
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
import org.apache.spark.sql.internal.SQLConf
@@ -64,7 +65,11 @@ abstract class DisableUnnecessaryBucketedScanSuite
val df = sql(query)
df.collect()
val plan = df.queryExecution.executedPlan
- val bucketedScan = collect(plan) { case s: FileSourceScanExec if s.bucketedScan => s }
+ val bucketedScan = collect(plan) {
+ case s: FileSourceScanExec if s.bucketedScan => s
+ case s: CometScanExec if s.bucketedScan => s
+ case s: CometNativeScanExec if s.bucketedScan => s
+ }
assert(bucketedScan.length == expectedNumBucketedScan)
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala
index 87ec48c3743..54c0dfbae87 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala
@@ -35,6 +35,7 @@ import org.apache.spark.paths.SparkPath
import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd}
import org.apache.spark.sql.{AnalysisException, DataFrame}
import org.apache.spark.sql.catalyst.util.stringToFile
+import org.apache.spark.sql.comet.CometBatchScanExec
import org.apache.spark.sql.execution.DataSourceScanExec
import org.apache.spark.sql.execution.datasources._
import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
@@ -842,6 +843,8 @@ class FileStreamSinkV2Suite extends FileStreamSinkSuite {
val fileScan = df.queryExecution.executedPlan.collect {
case batch: BatchScanExec if batch.scan.isInstanceOf[FileScan] =>
batch.scan.asInstanceOf[FileScan]
+ case batch: CometBatchScanExec if batch.scan.isInstanceOf[FileScan] =>
+ batch.scan.asInstanceOf[FileScan]
}.headOption.getOrElse {
fail(s"No FileScan in query\n${df.queryExecution}")
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala
index a9b8d54a6c6..01b5a13d735 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala
@@ -20,11 +20,13 @@ package org.apache.spark.sql.streaming
import scala.concurrent.duration._
import org.apache.spark.SparkIllegalArgumentException
+import org.apache.spark.sql.IgnoreCometSuite
import org.apache.spark.sql.execution.streaming.sources.LowLatencyMemoryStream
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
-class StreamRealTimeModeAllowlistSuite extends StreamRealTimeModeE2ESuiteBase {
+class StreamRealTimeModeAllowlistSuite extends StreamRealTimeModeE2ESuiteBase
+ with IgnoreCometSuite {
import testImplicits._
test("rtm source allowlist") {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamSuite.scala
index b4fd41a6b55..07b59eae016 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamSuite.scala
@@ -39,6 +39,7 @@ import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
import org.apache.spark.sql.catalyst.util.DateTimeUtils
import org.apache.spark.sql.classic.{DataFrame, Dataset}
import org.apache.spark.sql.classic.ClassicConversions._
+import org.apache.spark.sql.comet.CometLocalLimitExec
import org.apache.spark.sql.execution.{LocalLimitExec, SimpleMode, SparkPlan}
import org.apache.spark.sql.execution.command.ExplainCommand
import org.apache.spark.sql.execution.streaming._
@@ -1119,11 +1120,12 @@ class StreamSuite extends StreamTest {
val localLimits = StreamingQueryPlanTraverseHelper.collectFromUnfoldedPlan(execPlan) {
case l: LocalLimitExec => l
case l: StreamingLocalLimitExec => l
+ case l: CometLocalLimitExec => l
}
require(
localLimits.size == 1,
- s"Cant verify local limit optimization with this plan:\n$execPlan")
+ s"Cant verify local limit optimization ${localLimits.size} with this plan:\n$execPlan")
if (expectStreamingLimit) {
assert(
@@ -1131,7 +1133,8 @@ class StreamSuite extends StreamTest {
s"Local limit was not StreamingLocalLimitExec:\n$execPlan")
} else {
assert(
- localLimits.head.isInstanceOf[LocalLimitExec],
+ localLimits.head.isInstanceOf[LocalLimitExec] ||
+ localLimits.head.isInstanceOf[CometLocalLimitExec],
s"Local limit was not LocalLimitExec:\n$execPlan")
}
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamTest.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamTest.scala
index a6067aaf189..76e73f1c554 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamTest.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamTest.scala
@@ -73,7 +73,8 @@ import org.apache.spark.util.{Clock, SystemClock, Utils}
* avoid hanging forever in the case of failures. However, individual suites can change this
* by overriding `streamingTimeout`.
*/
-trait StreamTest extends SharedSparkSession with TimeLimits {
+trait StreamTest extends SharedSparkSession with TimeLimits
+ with org.apache.spark.sql.IgnoreCometSuite {
// Necessary to make ScalaTest 3.x interrupt a thread on the JVM like ScalaTest 2.2.x
implicit val defaultSignaler: Signaler = ThreadSignaler
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
index 465da3cd469..2e6bcb6ca0d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationDistributionSuite.scala
@@ -22,6 +22,7 @@ import java.io.File
import org.scalatest.Assertions
import org.apache.spark.sql.catalyst.plans.physical.UnspecifiedDistribution
+import org.apache.spark.sql.comet.CometBaseAggregateExec
import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
import org.apache.spark.sql.execution.streaming.operators.stateful.{StateStoreRestoreExec, StateStoreSaveExec}
import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
@@ -67,6 +68,7 @@ class StreamingAggregationDistributionSuite extends StreamTest
// verify aggregations in between, except partial aggregation
val allAggregateExecs = query.lastExecution.executedPlan.collect {
case a: BaseAggregateExec => a
+ case c: CometBaseAggregateExec => c.originalPlan
}
val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter {
@@ -201,6 +203,7 @@ class StreamingAggregationDistributionSuite extends StreamTest
// verify aggregations in between, except partial aggregation
val allAggregateExecs = executedPlan.collect {
case a: BaseAggregateExec => a
+ case c: CometBaseAggregateExec => c.originalPlan
}
val aggregateExecsWithoutPartialAgg = allAggregateExecs.filter {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala
index 1e1aa451a0a..dd18fe8cae9 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala
@@ -34,7 +34,7 @@ import org.apache.spark.sql.{AnalysisException, DataFrame, Row, SparkSession}
import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression}
import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning
import org.apache.spark.sql.execution.datasources.v2.state.StateSourceOptions
-import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
import org.apache.spark.sql.execution.streaming.checkpointing.CheckpointFileManager
import org.apache.spark.sql.execution.streaming.operators.stateful.StatefulOperatorStateInfo
import org.apache.spark.sql.execution.streaming.operators.stateful.join.{StreamingSymmetricHashJoinExec, StreamingSymmetricHashJoinHelper}
@@ -675,14 +675,28 @@ abstract class StreamingInnerJoinBase extends StreamingJoinSuite {
val numPartitions = spark.sessionState.conf.getConf(SQLConf.SHUFFLE_PARTITIONS)
- assert(query.lastExecution.executedPlan.collect {
- case j @ StreamingSymmetricHashJoinExec(_, _, _, _, _, _, _, _, _,
- ShuffleExchangeExec(opA: HashPartitioning, _, _, _),
- ShuffleExchangeExec(opB: HashPartitioning, _, _, _), _)
- if partitionExpressionsColumns(opA.expressions) === Seq("a", "b")
- && partitionExpressionsColumns(opB.expressions) === Seq("a", "b")
- && opA.numPartitions == numPartitions && opB.numPartitions == numPartitions => j
- }.size == 1)
+ val join = query.lastExecution.executedPlan.collect {
+ case j: StreamingSymmetricHashJoinExec => j
+ }.head
+ val opA = join.left.collect {
+ case s: ShuffleExchangeLike
+ if s.outputPartitioning.isInstanceOf[HashPartitioning] &&
+ partitionExpressionsColumns(
+ s.outputPartitioning
+ .asInstanceOf[HashPartitioning].expressions) === Seq("a", "b") =>
+ s.outputPartitioning
+ .asInstanceOf[HashPartitioning]
+ }.head
+ val opB = join.right.collect {
+ case s: ShuffleExchangeLike
+ if s.outputPartitioning.isInstanceOf[HashPartitioning] &&
+ partitionExpressionsColumns(
+ s.outputPartitioning
+ .asInstanceOf[HashPartitioning].expressions) === Seq("a", "b") =>
+ s.outputPartitioning
+ .asInstanceOf[HashPartitioning]
+ }.head
+ assert(opA.numPartitions == numPartitions && opB.numPartitions == numPartitions)
})
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQuerySuite.scala
index 7adf98b7920..ac46942b55d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQuerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQuerySuite.scala
@@ -43,7 +43,7 @@ import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
import org.apache.spark.sql.classic.{DataFrame, Dataset}
import org.apache.spark.sql.connector.read.InputPartition
import org.apache.spark.sql.connector.read.streaming.{Offset => OffsetV2, ReadLimit}
-import org.apache.spark.sql.execution.exchange.{REQUIRED_BY_STATEFUL_OPERATOR, ReusedExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.exchange.{REQUIRED_BY_STATEFUL_OPERATOR, ReusedExchangeExec, ShuffleExchangeLike}
import org.apache.spark.sql.execution.streaming._
import org.apache.spark.sql.execution.streaming.checkpointing.{CheckpointFileManager, OffsetSeqMetadata}
import org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreSaveExec
@@ -1471,7 +1471,7 @@ class StreamingQuerySuite extends StreamTest with BeforeAndAfter with Logging wi
CheckAnswer((1, 2), (2, 2), (3, 2)),
Execute { qe =>
val shuffleOpt = qe.lastExecution.executedPlan.collect {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
}
assert(shuffleOpt.nonEmpty, "No shuffle exchange found in the query plan")
@@ -1572,7 +1572,7 @@ class StreamingQuerySuite extends StreamTest with BeforeAndAfter with Logging wi
.get(SQLConf.SHUFFLE_PARTITIONS.key) === Some("1"))
val shuffles = q.lastExecution.executedPlan.collect {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
}
val shuffle = shuffles.head
@@ -1603,7 +1603,7 @@ class StreamingQuerySuite extends StreamTest with BeforeAndAfter with Logging wi
.get(SQLConf.SHUFFLE_PARTITIONS.key) === Some("1"))
val shuffles = q.lastExecution.executedPlan.collect {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
}
val shuffle = shuffles.head
@@ -1647,7 +1647,7 @@ class StreamingQuerySuite extends StreamTest with BeforeAndAfter with Logging wi
CheckAnswer((5, 10, 15)),
Execute { q =>
val shuffles = q.lastExecution.executedPlan.collect {
- case s: ShuffleExchangeExec => s
+ case s: ShuffleExchangeLike => s
}
val shuffle = shuffles.head
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/DataStreamTableAPISuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/DataStreamTableAPISuite.scala
index e2c74533e7f..a12d55848ea 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/DataStreamTableAPISuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/DataStreamTableAPISuite.scala
@@ -22,7 +22,7 @@ import java.util
import org.scalatest.BeforeAndAfter
-import org.apache.spark.sql.{AnalysisException, Row, SaveMode}
+import org.apache.spark.sql.{AnalysisException, IgnoreComet, Row, SaveMode}
import org.apache.spark.sql.catalyst.TableIdentifier
import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException
import org.apache.spark.sql.catalyst.catalog.{CatalogStorageFormat, CatalogTable, CatalogTableType}
@@ -363,7 +363,8 @@ class DataStreamTableAPISuite extends StreamTest with BeforeAndAfter {
}
}
- test("explain with table on DSv1 data source") {
+ test("explain with table on DSv1 data source",
+ IgnoreComet("Comet explain output is different")) {
val tblSourceName = "tbl_src"
val tblTargetName = "tbl_target"
val tblSourceQualified = s"default.$tblSourceName"
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
index fb26d3311eb..13dc8b89a91 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
@@ -102,6 +102,20 @@ trait SharedSparkSessionBase
// this rule may potentially block testing of other optimization rules such as
// ConstantPropagation etc.
.set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName)
+ // Enable Comet if `ENABLE_COMET` environment variable is set
+ if (isCometEnabled) {
+ conf
+ .set("spark.sql.extensions", "org.apache.comet.CometSparkSessionExtensions")
+ .set("spark.comet.enabled", "true")
+ .set("spark.comet.parquet.rowFilterPushdown.enabled", "true")
+
+ conf
+ .set("spark.comet.exec.enabled", "true")
+ .set("spark.shuffle.manager",
+ "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
+ .set("spark.comet.shuffle.enabled", "true")
+
+ }
conf.set(
StaticSQLConf.WAREHOUSE_PATH,
conf.get(StaticSQLConf.WAREHOUSE_PATH) + "/" + getClass.getCanonicalName)
diff --git a/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceWithActualMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceWithActualMetricsSuite.scala
index d6a5008f574..9cdc1d95347 100644
--- a/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceWithActualMetricsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceWithActualMetricsSuite.scala
@@ -48,7 +48,7 @@ class SqlResourceWithActualMetricsSuite
import testImplicits._
// Exclude nodes which may not have the metrics
- val excludedNodes = List("WholeStageCodegen", "Project", "SerializeFromObject")
+ val excludedNodes = List("WholeStageCodegen", "Project", "SerializeFromObject", "RowToColumnar")
implicit val formats: DefaultFormats = new DefaultFormats {
override def dateFormatter = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss")
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/DynamicPartitionPruningHiveScanSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/DynamicPartitionPruningHiveScanSuite.scala
index 1a02a68a069..ee104c6e193 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/DynamicPartitionPruningHiveScanSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/DynamicPartitionPruningHiveScanSuite.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql.hive
import org.apache.spark.sql._
import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression}
+import org.apache.spark.sql.comet._
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.adaptive.{DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite}
import org.apache.spark.sql.hive.execution.HiveTableScanExec
@@ -35,6 +36,12 @@ abstract class DynamicPartitionPruningHiveScanSuiteBase
case s: FileSourceScanExec => s.partitionFilters.collect {
case d: DynamicPruningExpression => d.child
}
+ case s: CometScanExec => s.partitionFilters.collect {
+ case d: DynamicPruningExpression => d.child
+ }
+ case s: CometNativeScanExec => s.partitionFilters.collect {
+ case d: DynamicPruningExpression => d.child
+ }
case h: HiveTableScanExec => h.partitionPruningPred.collect {
case d: DynamicPruningExpression => d.child
}
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/InsertSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/InsertSuite.scala
index bd2b256e492..9a3f5b9e160 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/InsertSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/InsertSuite.scala
@@ -871,7 +871,8 @@ class InsertSuite extends QueryTest with TestHiveSingleton with BeforeAndAfter {
}
}
- test("SPARK-30201 HiveOutputWriter standardOI should use ObjectInspectorCopyOption.DEFAULT") {
+ test("SPARK-30201 HiveOutputWriter standardOI should use ObjectInspectorCopyOption.DEFAULT",
+ IgnoreComet("Comet does not support reading non UTF-8 strings")) {
withTable("t1", "t2") {
withTempDir { dir =>
val file = new File(dir, "test.hex")
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/PartitionedTablePerfStatsSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/PartitionedTablePerfStatsSuite.scala
index 59022deaed7..f9aeacb5a9b 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/PartitionedTablePerfStatsSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/PartitionedTablePerfStatsSuite.scala
@@ -21,13 +21,14 @@ import java.io.File
import java.util.concurrent.{Executors, TimeUnit}
import org.apache.spark.metrics.source.HiveCatalogMetrics
-import org.apache.spark.sql.QueryTest
+import org.apache.spark.sql.{IgnoreCometSuite, QueryTest}
import org.apache.spark.sql.execution.datasources.FileStatusCache
import org.apache.spark.sql.hive.test.TestHiveSingleton
import org.apache.spark.sql.internal.SQLConf
class PartitionedTablePerfStatsSuite
- extends QueryTest with TestHiveSingleton {
+ extends QueryTest with TestHiveSingleton
+ with IgnoreCometSuite {
override def beforeEach(): Unit = {
super.beforeEach()
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala
index 8e7ff526a95..ea1072a7195 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala
@@ -54,24 +54,34 @@ object TestHive
new SparkContext(
System.getProperty("spark.sql.test.master", "local[1]"),
"TestSQLContext",
- new SparkConf()
- .set("spark.sql.test", "")
- .set(SQLConf.CODEGEN_FALLBACK.key, "false")
- .set(SQLConf.CODEGEN_FACTORY_MODE.key, CodegenObjectFactoryMode.CODEGEN_ONLY.toString)
- .set(HiveUtils.HIVE_METASTORE_BARRIER_PREFIXES.key,
- "org.apache.spark.sql.hive.execution.PairSerDe")
- .set(WAREHOUSE_PATH.key, TestHiveContext.makeWarehouseDir().toURI.getPath)
- // SPARK-8910
- .set(UI_ENABLED, false)
- .set(config.UNSAFE_EXCEPTION_ON_MEMORY_LEAK, true)
- // Hive changed the default of hive.metastore.disallow.incompatible.col.type.changes
- // from false to true. For details, see the JIRA HIVE-12320 and HIVE-17764.
- .set("spark.hadoop.hive.metastore.disallow.incompatible.col.type.changes", "false")
- // Disable ConvertToLocalRelation for better test coverage. Test cases built on
- // LocalRelation will exercise the optimization rules better by disabling it as
- // this rule may potentially block testing of other optimization rules such as
- // ConstantPropagation etc.
- .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName)
+ {
+ val conf = new SparkConf()
+ .set("spark.sql.test", "")
+ .set(SQLConf.CODEGEN_FALLBACK.key, "false")
+ .set(SQLConf.CODEGEN_FACTORY_MODE.key, CodegenObjectFactoryMode.CODEGEN_ONLY.toString)
+ .set(HiveUtils.HIVE_METASTORE_BARRIER_PREFIXES.key,
+ "org.apache.spark.sql.hive.execution.PairSerDe")
+ .set(WAREHOUSE_PATH.key, TestHiveContext.makeWarehouseDir().toURI.getPath)
+ .set(UI_ENABLED, false)
+ .set(config.UNSAFE_EXCEPTION_ON_MEMORY_LEAK, true)
+ .set("spark.hadoop.hive.metastore.disallow.incompatible.col.type.changes", "false")
+ .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName)
+
+ if (SparkSession.isCometEnabled) {
+ conf
+ .set("spark.sql.extensions", "org.apache.comet.CometSparkSessionExtensions")
+ .set("spark.comet.enabled", "true")
+
+ conf
+ .set("spark.comet.exec.enabled", "true")
+ .set("spark.shuffle.manager",
+ "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
+ .set("spark.comet.shuffle.enabled", "true")
+ }
+
+ conf
+ }
+
.set(SHUFFLE_EXCHANGE_MAX_THREAD_THRESHOLD,
sys.env.getOrElse("SPARK_TEST_HIVE_SHUFFLE_EXCHANGE_MAX_THREAD_THRESHOLD",
SHUFFLE_EXCHANGE_MAX_THREAD_THRESHOLD.defaultValueString).toInt)