diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index d85eb14d2c..f6b85808f2 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -515,6 +515,8 @@ jobs: org.apache.comet.exec.CometInMemoryCacheSuite org.apache.comet.exec.CometInMemoryCachePruningSuite org.apache.comet.exec.CometInMemoryCacheKryoSuite + org.apache.comet.exec.CometInMemoryCacheKryoUnregisteredSuite + org.apache.comet.exec.CometInMemoryCacheKryoClassesToRegisterSuite org.apache.comet.exec.CometGenerateExecSuite org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 02d3077e4f..7d6b1ac9fc 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -219,6 +219,8 @@ jobs: org.apache.comet.exec.CometInMemoryCacheSuite org.apache.comet.exec.CometInMemoryCachePruningSuite org.apache.comet.exec.CometInMemoryCacheKryoSuite + org.apache.comet.exec.CometInMemoryCacheKryoUnregisteredSuite + org.apache.comet.exec.CometInMemoryCacheKryoClassesToRegisterSuite org.apache.comet.exec.CometGenerateExecSuite org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff index 8fc6451ef1..537cef243c 100644 --- a/dev/diffs/3.4.3.diff +++ b/dev/diffs/3.4.3.diff @@ -237,10 +237,14 @@ index 0efe0877e9b..423d3b3d76d 100644 -- SELECT_HAVING -- https://github.com/postgres/postgres/blob/REL_12_BETA2/src/test/regress/sql/select_having.sql 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 cf40e944c09..bdd5be4f462 100644 +index cf40e944c09..3dc5574f819 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 -@@ -38,7 +38,7 @@ import org.apache.spark.sql.catalyst.util.DateTimeConstants +@@ -35,10 +35,11 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException + import org.apache.spark.sql.catalyst.expressions.SubqueryExpression + import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, Join, JoinStrategyHint, SHUFFLE_HASH} + import org.apache.spark.sql.catalyst.util.DateTimeConstants ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec import org.apache.spark.sql.execution.{ColumnarToRowExec, ExecSubqueryExpression, RDDScanExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.columnar._ @@ -249,7 +253,28 @@ index cf40e944c09..bdd5be4f462 100644 import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -516,7 +516,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -113,6 +114,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + case inMemoryTable @ InMemoryTableScanExec(_, _, relation) => + getNumInMemoryTablesRecursively(relation.cachedPlan) + + getNumInMemoryTablesInSubquery(inMemoryTable) + 1 ++ // Comet's cache scan keeps the predicates pushed into it on its original plan, out of reach ++ // of subquery reuse. The filter above it evaluates the same subqueries, and counts them. ++ case cometScan: CometInMemoryTableScanExec => ++ getNumInMemoryTablesRecursively(cometScan.originalPlan.relation.cachedPlan) + 1 + case p => + getNumInMemoryTablesInSubquery(p) + }.sum +@@ -393,7 +398,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + assert(isExpectStorageLevel(rddId, Disk)) + } + +- test("InMemoryRelation statistics") { ++ test("InMemoryRelation statistics", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + sql("CACHE TABLE testData") + spark.table("testData").queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +@@ -516,7 +522,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils */ private def verifyNumExchanges(df: DataFrame, expected: Int): Unit = { assert( @@ -259,6 +284,16 @@ index cf40e944c09..bdd5be4f462 100644 } test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { +@@ -1559,7 +1566,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + } + } + +- test("SPARK-36120: Support cache/uncache table with TimestampNTZ type") { ++ test("SPARK-36120: Support cache/uncache table with TimestampNTZ type", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + val tableName = "ntzCache" + withTable(tableName) { + sql(s"CACHE TABLE $tableName AS SELECT TIMESTAMP_NTZ'2021-01-01 00:00:00'") 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 1cc09c3d7fc..9e1e883d450 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -315,6 +350,20 @@ index 56e9520fdab..917932336df 100644 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 61724a39dfa..8aa517c1575 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 +@@ -1376,7 +1376,8 @@ class DataFrameSetOperationsSuite extends QueryTest with SharedSparkSession { + Row(Row(Seq(Seq(Row(null, "ba"))))) :: Nil) + } + +- test("SPARK-37371: UnionExec should support columnar if all children support columnar") { ++ test("SPARK-37371: UnionExec should support columnar if all children support columnar", ++ IgnoreComet("Comet replaces the cache scans and the union with its own operators")) { + def checkIfColumnar( + plan: SparkPlan, + targetPlan: (SparkPlan) => Boolean, 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 a9f69ab28a1..760ea0e9565 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala @@ -438,6 +487,66 @@ index 433b4741979..e13e69deb79 100644 case _ => false } case _ => false +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +index a657c6212aa..c90f10c8fa5 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +@@ -20,6 +20,8 @@ package org.apache.spark.sql + import org.scalatest.concurrent.TimeLimits + import org.scalatest.time.SpanSugar._ + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec ++import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec} + import org.apache.spark.sql.functions._ +@@ -34,6 +36,10 @@ class DatasetCacheSuite extends QueryTest + with AdaptiveSparkPlanHelper { + import testImplicits._ + ++ // A scan of a cached relation, Spark's or Comet's. ++ private def isCacheScan(plan: SparkPlan): Boolean = ++ plan.isInstanceOf[InMemoryTableScanExec] || plan.isInstanceOf[CometInMemoryTableScanExec] ++ + /** + * Asserts that a cached [[Dataset]] will be built using the given number of other cached results. + */ +@@ -41,7 +47,7 @@ class DatasetCacheSuite extends QueryTest + val plan = df.queryExecution.withCachedData + assert(plan.isInstanceOf[InMemoryRelation]) + val internalPlan = plan.asInstanceOf[InMemoryRelation].cacheBuilder.cachedPlan +- assert(find(internalPlan)(_.isInstanceOf[InMemoryTableScanExec]).size ++ assert(find(internalPlan)(isCacheScan).size + == numOfCachesDependedUpon) + } + +@@ -251,7 +257,7 @@ class DatasetCacheSuite extends QueryTest + case i: InMemoryRelation => i.cacheBuilder.cachedPlan + } + assert(df2LimitInnerPlan.isDefined && +- !df2LimitInnerPlan.get.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ !df2LimitInnerPlan.get.exists(isCacheScan)) + } + + test("SPARK-27739 Save stats from optimized plan") { +@@ -283,7 +289,7 @@ class DatasetCacheSuite extends QueryTest + val unionDf = df1.union(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.union(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative") { +@@ -291,7 +297,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.union(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + } + } 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 daef11ae4d6..9f3cc9181f2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -986,10 +1095,18 @@ index b5b34922694..a72403780c4 100644 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/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -index 525d97e4998..aded8906d75 100644 +index 525d97e4998..c15eaf98ddc 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 -@@ -1508,7 +1508,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -36,6 +36,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.{Complete, Partial} + import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, NestedColumnAliasingSuite} + import org.apache.spark.sql.catalyst.plans.logical.{LocalLimit, Project, RepartitionByExpression, Sort} + import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.{CommandResultExec, UnionExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.aggregate._ +@@ -1508,7 +1509,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001"))) } @@ -999,7 +1116,7 @@ index 525d97e4998..aded8906d75 100644 AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") { sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect() } -@@ -1960,8 +1961,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -1960,8 +1962,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark countAcc.add(1) x }) @@ -1016,7 +1133,15 @@ index 525d97e4998..aded8906d75 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) -@@ -3730,7 +3738,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -3269,6 +3278,7 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark + + val inMemoryTableScan = collect(queryDf.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c + } + assert(inMemoryTableScan.size == 2) + checkAnswer(queryDf, Row(0, 1) :: Row(1, 2) :: Row(2, 3) :: Row(3, 4) :: Row(4, 5) :: Nil) +@@ -3730,7 +3740,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark } } @@ -1056,6 +1181,32 @@ index 2dabcf01be7..8fcec0d1ce4 100644 } } } +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 48ad10992c5..23bf476d4dc 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 +@@ -208,6 +208,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper { + df.select("i").filter($"i" > 1).cache() + assert(df.filter($"i" > 1).select("i").queryExecution.executedPlan.find { + case _: org.apache.spark.sql.execution.columnar.InMemoryTableScanExec => true ++ case _: org.apache.spark.sql.comet.CometInMemoryTableScanExec => true + case _ => false + }.isDefined) + } +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +index e6b74a328e5..d4aa93eaebb 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +@@ -811,7 +811,8 @@ class StatisticsCollectionSuite extends StatisticsCollectionTestBase with Shared + } + } + +- test("SPARK-33687: analyze all tables in a specific database") { ++ test("SPARK-33687: analyze all tables in a specific database", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + withTempDatabase { database => + spark.catalog.setCurrentDatabase(database) + withTempDir { dir => 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 18123a4d6ec..0fe185baa33 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala @@ -1964,6 +2115,151 @@ index 593bd7bb4ba..b327d84d5cc 100644 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")) +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +index d15fabd9403..9fc89bbd2a0 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +@@ -22,11 +22,12 @@ import java.sql.{Date, Timestamp} + import java.util.concurrent.atomic.AtomicInteger + + import org.apache.spark.rdd.RDD +-import org.apache.spark.sql.{DataFrame, QueryTest, Row} ++import org.apache.spark.sql.{DataFrame, IgnoreComet, QueryTest, Row} + import org.apache.spark.sql.catalyst.InternalRow + import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, In} + import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning + import org.apache.spark.sql.columnar.CachedBatch ++import org.apache.spark.sql.comet.{CometFilterExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} + import org.apache.spark.sql.functions._ + import org.apache.spark.sql.internal.SQLConf +@@ -507,11 +508,17 @@ class InMemoryColumnarQuerySuite extends QueryTest with SharedSparkSession { + val planBeforeFilter = df2.queryExecution.executedPlan.collect { + case f: FilterExec => f.child + case WholeStageCodegenExec(FilterExec(_, i: InputAdapter)) => i.child ++ case f: CometFilterExec => f.child + } +- assert(planBeforeFilter.head.isInstanceOf[InMemoryTableScanExec]) +- + val execPlan = planBeforeFilter.head +- assert(execPlan.executeCollectPublic().length == 0) ++ execPlan match { ++ // Comet's cache scan is columnar only, so count the rows of the batches it emits. ++ case c: CometInMemoryTableScanExec => ++ assert(c.executeColumnar().map(_.numRows().toLong).collect().sum == 0) ++ case _ => ++ assert(execPlan.isInstanceOf[InMemoryTableScanExec]) ++ assert(execPlan.executeCollectPublic().length == 0) ++ } + } + + test("SPARK-25727 - otherCopyArgs in InMemoryRelation does not include outputOrdering") { +@@ -520,7 +527,8 @@ class InMemoryColumnarQuerySuite extends QueryTest with SharedSparkSession { + assert(json.contains("outputOrdering")) + } + +- test("SPARK-22673: InMemoryRelation should utilize existing stats of the plan to be cached") { ++ test("SPARK-22673: InMemoryRelation should utilize existing stats of the plan to be cached", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + Seq("orc", "").foreach { useV1SourceReaderList => + // This test case depends on the size of ORC in statistics. + withSQLConf( +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +index e032e0c2b27..89d7867c250 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution.columnar + + import org.apache.spark.SparkFunSuite ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions.expr +@@ -38,11 +39,15 @@ class InMemoryRelationSuite extends SparkFunSuite + test("SPARK-47177: Cached SQL plan do not display final AQE plan in explain string") { + def findIMRInnerChild(p: SparkPlan): SparkPlan = { + val tableCache = find(p) { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + case _ => false + } + assert(tableCache.isDefined) +- tableCache.get.asInstanceOf[InMemoryTableScanExec].relation.innerChildren.head ++ val scan = tableCache.get match { ++ case c: CometInMemoryTableScanExec => c.originalPlan ++ case s => s.asInstanceOf[InMemoryTableScanExec] ++ } ++ scan.relation.innerChildren.head + } + + withSQLConf(SQLConf.CAN_CHANGE_CACHED_PLAN_OUTPUT_PARTITIONING.key -> "true") { +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +index a22cb664744..3831e428785 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +@@ -17,6 +17,7 @@ + + package org.apache.spark.sql.execution.columnar + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.internal.SQLConf + import org.apache.spark.sql.test.SharedSparkSession + import org.apache.spark.sql.test.SQLTestData._ +@@ -180,11 +181,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession { + val result = df.collect().map(_(0)).toArray + assert(result.length === 1) + +- val (readPartitions, readBatches) = df.queryExecution.executedPlan.collect { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- assert(readPartitions === 5) +- assert(readBatches === 10) ++ val scans = df.queryExecution.executedPlan.collect { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. ++ if (scans.isEmpty) { ++ assert(df.queryExecution.executedPlan.collect { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ assert(scans.head.readPartitions.value === 5) ++ assert(scans.head.readBatches.value === 10) ++ } + } + + def checkBatchPruning( +@@ -201,14 +207,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession { + df.collect().map(_(0)).toArray + } + +- val (readPartitions, readBatches) = df.queryExecution.executedPlan.collect { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- +- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") +- assert( +- readPartitions === expectedReadPartitions, +- s"Wrong number of read partitions: $queryExecution") ++ val scans = df.queryExecution.executedPlan.collect { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan prunes on the same statistics, but has none of these test-only ++ // accumulators to read, so only the answer above is checked for it. ++ if (scans.isEmpty) { ++ assert(df.queryExecution.executedPlan.collect { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ val readPartitions = scans.head.readPartitions.value ++ val readBatches = scans.head.readBatches.value ++ assert( ++ readBatches === expectedReadBatches, ++ s"Wrong number of read batches: $queryExecution") ++ assert( ++ readPartitions === expectedReadPartitions, ++ s"Wrong number of read partitions: $queryExecution") ++ } + } + } + } 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 bd9c79e5b96..2ada8c28842 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala @@ -3053,10 +3349,10 @@ index dd55fcfe42c..d9a3f2df535 100644 spark.internalCreateDataFrame(withoutFilters.execute(), schema) 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 ed2e309fa07..54d417624ff 100644 +index ed2e309fa07..8e3aaa888c7 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 -@@ -74,6 +74,19 @@ trait SharedSparkSessionBase +@@ -74,6 +74,24 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -3072,6 +3368,11 @@ index ed2e309fa07..54d417624ff 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ // CometDriverPlugin installs Comet's cache serializer when ++ // spark.comet.exec.inMemoryCache.enabled is on, as it is by default. These sessions do not ++ // load the plugin, so install it here. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } conf.set( StaticSQLConf.WAREHOUSE_PATH, @@ -3157,10 +3458,10 @@ index a902cb3a69e..e652edd9f81 100644 test("SPARK-4963 DataFrame sample on mutable row return wrong result") { 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 07361cfdce9..af6dcfc2302 100644 +index 07361cfdce9..f9002ce0d98 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 -@@ -55,25 +55,41 @@ object TestHive +@@ -55,25 +55,45 @@ object TestHive new SparkContext( System.getProperty("spark.sql.test.master", "local[1]"), "TestSQLContext", @@ -3212,6 +3513,10 @@ index 07361cfdce9..af6dcfc2302 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ ++ // As in SharedSparkSession: what CometDriverPlugin would install. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } + conf diff --git a/dev/diffs/3.5.9.diff b/dev/diffs/3.5.9.diff index 462554a104..653566b701 100644 --- a/dev/diffs/3.5.9.diff +++ b/dev/diffs/3.5.9.diff @@ -218,10 +218,14 @@ index 0efe0877e9b..423d3b3d76d 100644 -- SELECT_HAVING -- https://github.com/postgres/postgres/blob/REL_12_BETA2/src/test/regress/sql/select_having.sql 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 e5494726695..00937f025c2 100644 +index e5494726695..7a2a2d8b721 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 -@@ -38,7 +38,7 @@ import org.apache.spark.sql.catalyst.util.DateTimeConstants +@@ -35,10 +35,11 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException + import org.apache.spark.sql.catalyst.expressions.SubqueryExpression + import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, Join, JoinStrategyHint, SHUFFLE_HASH} + import org.apache.spark.sql.catalyst.util.DateTimeConstants ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec import org.apache.spark.sql.execution.{ColumnarToRowExec, ExecSubqueryExpression, RDDScanExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEPropagateEmptyRelation} import org.apache.spark.sql.execution.columnar._ @@ -230,7 +234,28 @@ index e5494726695..00937f025c2 100644 import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -519,7 +519,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -113,6 +114,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + case inMemoryTable @ InMemoryTableScanExec(_, _, relation) => + getNumInMemoryTablesRecursively(relation.cachedPlan) + + getNumInMemoryTablesInSubquery(inMemoryTable) + 1 ++ // Comet's cache scan keeps the predicates pushed into it on its original plan, out of reach ++ // of subquery reuse. The filter above it evaluates the same subqueries, and counts them. ++ case cometScan: CometInMemoryTableScanExec => ++ getNumInMemoryTablesRecursively(cometScan.originalPlan.relation.cachedPlan) + 1 + case p => + getNumInMemoryTablesInSubquery(p) + }.sum +@@ -393,7 +398,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + assert(isExpectStorageLevel(rddId, Disk)) + } + +- test("InMemoryRelation statistics") { ++ test("InMemoryRelation statistics", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + sql("CACHE TABLE testData") + spark.table("testData").queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +@@ -519,7 +525,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils df.collect() } assert( @@ -240,6 +265,16 @@ index e5494726695..00937f025c2 100644 } test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { +@@ -1574,7 +1581,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + } + } + +- test("SPARK-36120: Support cache/uncache table with TimestampNTZ type") { ++ test("SPARK-36120: Support cache/uncache table with TimestampNTZ type", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + val tableName = "ntzCache" + withTable(tableName) { + sql(s"CACHE TABLE $tableName AS SELECT TIMESTAMP_NTZ'2021-01-01 00:00:00'") 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 6f3090d8908..4774aad5019 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -296,6 +331,20 @@ index 56e9520fdab..917932336df 100644 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 747f43fa2a7..2e52a4acb79 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 +@@ -1398,7 +1398,8 @@ class DataFrameSetOperationsSuite extends QueryTest + Row(Row(Seq(Seq(Row(null, "ba"))))) :: Nil) + } + +- test("SPARK-37371: UnionExec should support columnar if all children support columnar") { ++ test("SPARK-37371: UnionExec should support columnar if all children support columnar", ++ IgnoreComet("Comet replaces the cache scans and the union with its own operators")) { + def checkIfColumnar( + plan: SparkPlan, + targetPlan: (SparkPlan) => Boolean, 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 7ee18df3756..d09f70e5d99 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala @@ -419,6 +468,83 @@ index a1d5d579338..8825683ebcd 100644 case _ => false } case _ => false +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +index bda8c7f2608..10a70745d07 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +@@ -20,6 +20,8 @@ package org.apache.spark.sql + import org.scalatest.concurrent.TimeLimits + import org.scalatest.time.SpanSugar._ + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec ++import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec} + import org.apache.spark.sql.functions._ +@@ -35,6 +37,10 @@ class DatasetCacheSuite extends QueryTest + with AdaptiveSparkPlanHelper { + import testImplicits._ + ++ // A scan of a cached relation, Spark's or Comet's. ++ private def isCacheScan(plan: SparkPlan): Boolean = ++ plan.isInstanceOf[InMemoryTableScanExec] || plan.isInstanceOf[CometInMemoryTableScanExec] ++ + /** + * Asserts that a cached [[Dataset]] will be built using the given number of other cached results. + */ +@@ -42,7 +48,7 @@ class DatasetCacheSuite extends QueryTest + val plan = df.queryExecution.withCachedData + assert(plan.isInstanceOf[InMemoryRelation]) + val internalPlan = plan.asInstanceOf[InMemoryRelation].cacheBuilder.cachedPlan +- assert(find(internalPlan)(_.isInstanceOf[InMemoryTableScanExec]).size ++ assert(find(internalPlan)(isCacheScan).size + == numOfCachesDependedUpon) + } + +@@ -252,7 +258,7 @@ class DatasetCacheSuite extends QueryTest + case i: InMemoryRelation => i.cacheBuilder.cachedPlan + } + assert(df2LimitInnerPlan.isDefined && +- !df2LimitInnerPlan.get.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ !df2LimitInnerPlan.get.exists(isCacheScan)) + } + + test("SPARK-27739 Save stats from optimized plan") { +@@ -285,14 +291,14 @@ class DatasetCacheSuite extends QueryTest + val unionDf = df1.union(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.union(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("positive: union by name") { + val unionDf = df1.unionByName(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.unionByName(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by position") { +@@ -300,7 +306,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.union(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by name") { +@@ -308,7 +314,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.unionByName(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + } + } 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 c4fb4fa943c..a04b23870a8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -993,10 +1119,18 @@ index c26757c9cff..d55775f09d7 100644 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/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -index 3cf2bfd17ab..5bcf9478e9b 100644 +index 3cf2bfd17ab..11d3ca4d5a9 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 -@@ -1521,7 +1521,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -38,6 +38,7 @@ import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, NestedCo + import org.apache.spark.sql.catalyst.parser.ParseException + import org.apache.spark.sql.catalyst.plans.logical.{LocalLimit, Project, RepartitionByExpression, Sort} + import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.{CommandResultExec, UnionExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.aggregate._ +@@ -1521,7 +1522,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001"))) } @@ -1006,7 +1140,7 @@ index 3cf2bfd17ab..5bcf9478e9b 100644 AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") { sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect() } -@@ -1979,8 +1980,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -1979,8 +1981,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark countAcc.add(1) x }) @@ -1023,7 +1157,15 @@ index 3cf2bfd17ab..5bcf9478e9b 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) -@@ -3750,7 +3758,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -3289,6 +3298,7 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark + + val inMemoryTableScan = collect(queryDf.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c + } + assert(inMemoryTableScan.size == 2) + checkAnswer(queryDf, Row(0, 1) :: Row(1, 2) :: Row(2, 3) :: Row(3, 4) :: Row(4, 5) :: Nil) +@@ -3750,7 +3760,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark } } @@ -1064,6 +1206,32 @@ index 71af1fd69c3..81a04c93c9c 100644 } } } +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 8b4ac474f87..096146cd69a 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 +@@ -210,6 +210,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt + df.select("i").filter($"i" > 1).cache() + assert(find(df.filter($"i" > 1).select("i").queryExecution.executedPlan) { + case _: org.apache.spark.sql.execution.columnar.InMemoryTableScanExec => true ++ case _: org.apache.spark.sql.comet.CometInMemoryTableScanExec => true + case _ => false + }.isDefined) + } +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +index e827396009d..066d1a6bf04 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +@@ -811,7 +811,8 @@ class StatisticsCollectionSuite extends StatisticsCollectionTestBase with Shared + } + } + +- test("SPARK-33687: analyze all tables in a specific database") { ++ test("SPARK-33687: analyze all tables in a specific database", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + withTempDatabase { database => + spark.catalog.setCurrentDatabase(database) + withTempDir { dir => 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 04702201f82..4d38d8d6e51 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala @@ -1561,7 +1729,7 @@ index 5a413c77754..207b66e1d7b 100644 import testImplicits._ 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 2f8e401e743..7849c685b19 100644 +index 2f8e401e743..1b3a6ff17b5 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,9 +27,11 @@ import org.scalatest.time.SpanSugar._ @@ -1967,7 +2135,27 @@ index 2f8e401e743..7849c685b19 100644 }.size == (if (firstAccess) 2 else 0)) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -2938,7 +2981,8 @@ class AdaptiveQueryExecSuite +@@ -2898,7 +2941,8 @@ class AdaptiveQueryExecSuite + } + } + +- test("SPARK-42101: Coalesce shuffle partition with union even if exists TableCacheQueryStage") { ++ test("SPARK-42101: Coalesce shuffle partition with union even if exists TableCacheQueryStage", ++ IgnoreComet("https://github.com/apache/datafusion-comet/issues/6454")) { + withSQLConf(SQLConf.CAN_CHANGE_CACHED_PLAN_OUTPUT_PARTITIONING.key -> "true", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1") { + val cached = Seq(1).toDF("c").cache() +@@ -2923,7 +2967,8 @@ class AdaptiveQueryExecSuite + } + } + +- test("SPARK-43376: Improve reuse subquery with table cache") { ++ test("SPARK-43376: Improve reuse subquery with table cache", ++ IgnoreComet("Comet's cache scan does not plan the subqueries in its pruning predicates")) { + withSQLConf(SQLConf.CAN_CHANGE_CACHED_PLAN_OUTPUT_PARTITIONING.key -> "true") { + withTable("t1", "t2") { + withCache("t1") { +@@ -2938,7 +2983,8 @@ class AdaptiveQueryExecSuite } } @@ -1977,7 +2165,7 @@ index 2f8e401e743..7849c685b19 100644 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")) -@@ -2980,7 +3024,9 @@ class AdaptiveQueryExecSuite +@@ -2980,7 +3026,9 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -1988,6 +2176,151 @@ index 2f8e401e743..7849c685b19 100644 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/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +index de04938f247..fa019925b6f 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +@@ -22,11 +22,12 @@ import java.sql.{Date, Timestamp} + import java.util.concurrent.atomic.AtomicInteger + + import org.apache.spark.rdd.RDD +-import org.apache.spark.sql.{DataFrame, QueryTest, Row} ++import org.apache.spark.sql.{DataFrame, IgnoreComet, QueryTest, Row} + import org.apache.spark.sql.catalyst.InternalRow + import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, In} + import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning + import org.apache.spark.sql.columnar.CachedBatch ++import org.apache.spark.sql.comet.{CometFilterExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions._ +@@ -509,11 +510,17 @@ class InMemoryColumnarQuerySuite extends QueryTest + val planBeforeFilter = collect(df2.queryExecution.executedPlan) { + case f: FilterExec => f.child + case WholeStageCodegenExec(FilterExec(_, i: InputAdapter)) => i.child ++ case f: CometFilterExec => f.child + } +- assert(planBeforeFilter.head.isInstanceOf[InMemoryTableScanExec]) +- + val execPlan = planBeforeFilter.head +- assert(execPlan.executeCollectPublic().length == 0) ++ execPlan match { ++ // Comet's cache scan is columnar only, so count the rows of the batches it emits. ++ case c: CometInMemoryTableScanExec => ++ assert(c.executeColumnar().map(_.numRows().toLong).collect().sum == 0) ++ case _ => ++ assert(execPlan.isInstanceOf[InMemoryTableScanExec]) ++ assert(execPlan.executeCollectPublic().length == 0) ++ } + } + + test("SPARK-25727 - otherCopyArgs in InMemoryRelation does not include outputOrdering") { +@@ -522,7 +529,8 @@ class InMemoryColumnarQuerySuite extends QueryTest + assert(json.contains("outputOrdering")) + } + +- test("SPARK-22673: InMemoryRelation should utilize existing stats of the plan to be cached") { ++ test("SPARK-22673: InMemoryRelation should utilize existing stats of the plan to be cached", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + Seq("orc", "").foreach { useV1SourceReaderList => + // This test case depends on the size of ORC in statistics. + withSQLConf( +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +index 2c73622739a..5d0efeb263a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution.columnar + + import org.apache.spark.SparkFunSuite ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions.expr +@@ -37,11 +38,15 @@ class InMemoryRelationSuite extends SparkFunSuite + test("SPARK-47177: Cached SQL plan do not display final AQE plan in explain string") { + def findIMRInnerChild(p: SparkPlan): SparkPlan = { + val tableCache = find(p) { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + case _ => false + } + assert(tableCache.isDefined) +- tableCache.get.asInstanceOf[InMemoryTableScanExec].relation.innerChildren.head ++ val scan = tableCache.get match { ++ case c: CometInMemoryTableScanExec => c.originalPlan ++ case s => s.asInstanceOf[InMemoryTableScanExec] ++ } ++ scan.relation.innerChildren.head + } + + val d1 = spark.range(1).withColumn("key", expr("id % 100")) +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +index 885286843a1..f0f805a6cd9 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +@@ -17,6 +17,7 @@ + + package org.apache.spark.sql.execution.columnar + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.internal.SQLConf + import org.apache.spark.sql.test.SharedSparkSession +@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + val result = df.collect().map(_(0)).toArray + assert(result.length === 1) + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- assert(readPartitions === 5) +- assert(readBatches === 10) ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ assert(scans.head.readPartitions.value === 5) ++ assert(scans.head.readBatches.value === 10) ++ } + } + + def checkBatchPruning( +@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + df.collect().map(_(0)).toArray + } + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- +- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") +- assert( +- readPartitions === expectedReadPartitions, +- s"Wrong number of read partitions: $queryExecution") ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan prunes on the same statistics, but has none of these test-only ++ // accumulators to read, so only the answer above is checked for it. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ val readPartitions = scans.head.readPartitions.value ++ val readBatches = scans.head.readBatches.value ++ assert( ++ readBatches === expectedReadBatches, ++ s"Wrong number of read batches: $queryExecution") ++ assert( ++ readPartitions === expectedReadPartitions, ++ s"Wrong number of read partitions: $queryExecution") ++ } + } + } + } 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 fd52d038ca6..154c800be67 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala @@ -3064,10 +3397,10 @@ index e937173a590..263934fbe7b 100644 spark.internalCreateDataFrame(withoutFilters.execute(), schema) 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 c23bf4204f7..07d215aad2b 100644 +index c23bf4204f7..d922f0b3d23 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 -@@ -97,6 +97,19 @@ trait SharedSparkSessionBase +@@ -97,6 +97,24 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -3083,6 +3416,11 @@ index c23bf4204f7..07d215aad2b 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ // CometDriverPlugin installs Comet's cache serializer when ++ // spark.comet.exec.inMemoryCache.enabled is on, as it is by default. These sessions do not ++ // load the plugin, so install it here. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } conf.set( StaticSQLConf.WAREHOUSE_PATH, @@ -3168,10 +3506,10 @@ index 6160c3e5f6c..bfc0c618a9b 100644 test("SPARK-4963 DataFrame sample on mutable row return wrong result") { 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 1d646f40b3e..c8192f52f98 100644 +index 1d646f40b3e..b8a043c6248 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 -@@ -53,25 +53,41 @@ object TestHive +@@ -53,25 +53,45 @@ object TestHive new SparkContext( System.getProperty("spark.sql.test.master", "local[1]"), "TestSQLContext", @@ -3223,6 +3561,10 @@ index 1d646f40b3e..c8192f52f98 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ ++ // As in SharedSparkSession: what CometDriverPlugin would install. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } + conf diff --git a/dev/diffs/4.0.4.diff b/dev/diffs/4.0.4.diff index 38ec78543a..37521ebcf9 100644 --- a/dev/diffs/4.0.4.diff +++ b/dev/diffs/4.0.4.diff @@ -333,10 +333,14 @@ index 21a3ce1e122..f4762ab98f0 100644 -- 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 0f42502f1d9..e9ff802141f 100644 +index 0f42502f1d9..81990d4a97f 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 -@@ -39,7 +39,7 @@ import org.apache.spark.sql.catalyst.util.DateTimeConstants +@@ -36,10 +36,11 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException + import org.apache.spark.sql.catalyst.expressions.SubqueryExpression + import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, Join, JoinStrategyHint, SHUFFLE_HASH} + import org.apache.spark.sql.catalyst.util.DateTimeConstants ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec import org.apache.spark.sql.execution.{ColumnarToRowExec, ExecSubqueryExpression, RDDScanExec, SparkPlan, SparkPlanInfo} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEPropagateEmptyRelation} import org.apache.spark.sql.execution.columnar._ @@ -345,7 +349,28 @@ index 0f42502f1d9..e9ff802141f 100644 import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -520,7 +520,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -114,6 +115,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + case inMemoryTable @ InMemoryTableScanExec(_, _, relation) => + getNumInMemoryTablesRecursively(relation.cachedPlan) + + getNumInMemoryTablesInSubquery(inMemoryTable) + 1 ++ // Comet's cache scan keeps the predicates pushed into it on its original plan, out of reach ++ // of subquery reuse. The filter above it evaluates the same subqueries, and counts them. ++ case cometScan: CometInMemoryTableScanExec => ++ getNumInMemoryTablesRecursively(cometScan.originalPlan.relation.cachedPlan) + 1 + case p => + getNumInMemoryTablesInSubquery(p) + }.sum +@@ -394,7 +399,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + assert(isExpectStorageLevel(rddId, Disk)) + } + +- test("InMemoryRelation statistics") { ++ test("InMemoryRelation statistics", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + sql("CACHE TABLE testData") + spark.table("testData").queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +@@ -520,7 +526,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils df.collect() } assert( @@ -355,7 +380,27 @@ index 0f42502f1d9..e9ff802141f 100644 } test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { -@@ -1659,9 +1660,18 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -1581,7 +1588,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + } + } + +- test("SPARK-36120: Support cache/uncache table with TimestampNTZ type") { ++ test("SPARK-36120: Support cache/uncache table with TimestampNTZ type", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + val tableName = "ntzCache" + withTable(tableName) { + sql(s"CACHE TABLE $tableName AS SELECT TIMESTAMP_NTZ'2021-01-01 00:00:00'") +@@ -1626,7 +1634,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + } + } + +- test("SPARK-35332: Make cache plan disable configs configurable - check AQE") { ++ test("SPARK-35332: Make cache plan disable configs configurable - check AQE", ++ IgnoreComet("Spark's SQL UI shows a cached plan only under Spark's own cache scan")) { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { +@@ -1659,9 +1668,18 @@ class CachedTableSuite extends QueryTest with SQLTestUtils _.nodeName.contains("TableCacheQueryStage")) val aqeNode = findNodeInSparkPlanInfo(inMemoryScanNode.get, _.nodeName.contains("AdaptiveSparkPlan")) @@ -377,6 +422,14 @@ index 0f42502f1d9..e9ff802141f 100644 } withTempView("t0", "t1", "t2") { +@@ -1750,6 +1768,7 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + val cached = spark.table("t") + val tableCache = collect(cached.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + if (expected == StorageLevel.NONE) { + assert(tableCache.isEmpty) 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 9db406ff12f..b3d55394d25 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -433,6 +486,20 @@ index ed182322aec..1ae6afa686a 100644 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 332be4c7bbc..02899683f81 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 +@@ -1474,7 +1474,8 @@ class DataFrameSetOperationsSuite extends QueryTest + Row(Row(Seq(Seq(Row(null, "ba"))))) :: Nil) + } + +- test("SPARK-37371: UnionExec should support columnar if all children support columnar") { ++ test("SPARK-37371: UnionExec should support columnar if all children support columnar", ++ IgnoreComet("Comet replaces the cache scans and the union with its own operators")) { + def checkIfColumnar( + plan: SparkPlan, + targetPlan: (SparkPlan) => Boolean, 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 d9ce3000a0c..f2d044ed6b8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala @@ -556,6 +623,97 @@ index 552e2b2e274..17a5ae20f0f 100644 case _ => false } case _ => false +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +index 9d8aaf8d90e..41afdbdafc1 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +@@ -20,6 +20,8 @@ package org.apache.spark.sql + import org.scalatest.concurrent.TimeLimits + import org.scalatest.time.SpanSugar._ + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec ++import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec} + import org.apache.spark.sql.functions._ +@@ -36,6 +38,10 @@ class DatasetCacheSuite extends QueryTest + with AdaptiveSparkPlanHelper { + import testImplicits._ + ++ // A scan of a cached relation, Spark's or Comet's. ++ private def isCacheScan(plan: SparkPlan): Boolean = ++ plan.isInstanceOf[InMemoryTableScanExec] || plan.isInstanceOf[CometInMemoryTableScanExec] ++ + /** + * Asserts that a cached [[Dataset]] will be built using the given number of other cached results. + */ +@@ -43,7 +49,7 @@ class DatasetCacheSuite extends QueryTest + val plan = df.queryExecution.withCachedData + assert(plan.isInstanceOf[InMemoryRelation]) + val internalPlan = plan.asInstanceOf[InMemoryRelation].cacheBuilder.cachedPlan +- assert(find(internalPlan)(_.isInstanceOf[InMemoryTableScanExec]).size ++ assert(find(internalPlan)(isCacheScan).size + == numOfCachesDependedUpon) + } + +@@ -253,7 +259,7 @@ class DatasetCacheSuite extends QueryTest + case i: InMemoryRelation => i.cacheBuilder.cachedPlan + } + assert(df2LimitInnerPlan.isDefined && +- !df2LimitInnerPlan.get.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ !df2LimitInnerPlan.get.exists(isCacheScan)) + } + + test("SPARK-27739 Save stats from optimized plan") { +@@ -286,14 +292,14 @@ class DatasetCacheSuite extends QueryTest + val unionDf = df1.union(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.union(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("positive: union by name") { + val unionDf = df1.unionByName(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.unionByName(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by position") { +@@ -301,7 +307,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.union(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by name") { +@@ -309,7 +315,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.unionByName(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + } + } +@@ -321,11 +327,11 @@ class DatasetCacheSuite extends QueryTest + df1.cache() + // This is exactly the same as df1. + val df2 = spark.range(5).select(struct($"id".as("name", metadata))) +- assert(df2.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(df2.queryExecution.executedPlan.exists(isCacheScan)) + + val metadata2 = Metadata.fromJson("""{"k2": "v2"}""") + // Same with df1 except for the Alias metadata + val df3 = spark.range(5).select(struct($"id".as("name", metadata2))) +- assert(!df3.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!df3.queryExecution.executedPlan.exists(isCacheScan)) + } + } 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 81713c777bc..b5f92ed9742 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -1140,10 +1298,18 @@ index ad424b3a7cc..4ece0117a34 100644 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/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -index f294ff81021..02d72be8d29 100644 +index f294ff81021..17d07ab6103 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 -@@ -1524,7 +1524,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -38,6 +38,7 @@ import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, NestedCo + import org.apache.spark.sql.catalyst.parser.ParseException + import org.apache.spark.sql.catalyst.plans.logical.{LocalLimit, Project, RepartitionByExpression, Sort} + import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.{CommandResultExec, UnionExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.aggregate._ +@@ -1524,7 +1525,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001"))) } @@ -1153,7 +1319,7 @@ index f294ff81021..02d72be8d29 100644 AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") { sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect() } -@@ -1985,8 +1986,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -1985,8 +1987,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark countAcc.add(1) x }) @@ -1170,6 +1336,14 @@ index f294ff81021..02d72be8d29 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) +@@ -3278,6 +3287,7 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark + + val inMemoryTableScan = collect(queryDf.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c + } + assert(inMemoryTableScan.size == 2) + checkAnswer(queryDf, Row(0, 1) :: Row(1, 2) :: Row(2, 3) :: Row(3, 4) :: Row(4, 5) :: Nil) 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 575a4ae69d1..129d9f27232 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala @@ -1201,6 +1375,18 @@ index 575a4ae69d1..129d9f27232 100644 } } } +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 c1c041509c3..580d320394d 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 +@@ -222,6 +222,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt + df.select("i").filter($"i" > 1).cache() + assert(find(df.filter($"i" > 1).select("i").queryExecution.executedPlan) { + case _: org.apache.spark.sql.execution.columnar.InMemoryTableScanExec => true ++ case _: org.apache.spark.sql.comet.CometInMemoryTableScanExec => true + case _ => false + }.isDefined) + } 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 5ba69c8f9d9..ac1256afe88 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala @@ -1216,6 +1402,20 @@ index 5ba69c8f9d9..ac1256afe88 100644 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/StatisticsCollectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +index 5222d5ce266..67f663c2764 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +@@ -829,7 +829,8 @@ class StatisticsCollectionSuite extends StatisticsCollectionTestBase with Shared + } + } + +- test("SPARK-33687: analyze all tables in a specific database") { ++ test("SPARK-33687: analyze all tables in a specific database", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + withTempDatabase { database => + spark.catalog.setCurrentDatabase(database) + withTempDir { dir => 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 0df7f806272..9cdfe8b8f46 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala @@ -1415,7 +1615,7 @@ index a40e34d94d0..abc1f035d15 100644 } } 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 11e9547dfc5..ba340c4ebcf 100644 +index 11e9547dfc5..327e0b1bc5b 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 @@ -24,6 +24,8 @@ import org.apache.spark.sql.{AnalysisException, Row} @@ -1423,7 +1623,7 @@ index 11e9547dfc5..ba340c4ebcf 100644 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.comet.{CometHashAggregateExec, CometInMemoryTableScanExec} import org.apache.spark.sql.connector.{DatasourceV2SQLBase, FakeV2ProviderWithCustomSchema} import org.apache.spark.sql.connector.catalog.{Identifier, InMemoryTable} import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper @@ -1482,6 +1682,15 @@ index 11e9547dfc5..ba340c4ebcf 100644 }.head.isInstanceOf[ArrayTransform]) } } +@@ -1897,7 +1909,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { + // Checks in-memory fetching code path. + val all = sql("SELECT col FROM tbl") + assert(all.queryExecution.executedPlan.collectFirst { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + }.nonEmpty) + checkAnswer(all, Row("a")) + // Checks column stats code path. 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 3eeed2e4175..9f21d547c1c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala @@ -2208,7 +2417,7 @@ index a3cfdc5a240..3793b6191bf 100644 }) 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 fb8fab6a80f..403eb411920 100644 +index fb8fab6a80f..8a8dcee18db 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 @@ -28,12 +28,14 @@ import org.apache.spark.SparkException @@ -2693,7 +2902,27 @@ index fb8fab6a80f..403eb411920 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3048,7 +3101,8 @@ class AdaptiveQueryExecSuite +@@ -3003,7 +3056,8 @@ class AdaptiveQueryExecSuite + } + } + +- test("SPARK-42101: Coalesce shuffle partition with union even if exists TableCacheQueryStage") { ++ test("SPARK-42101: Coalesce shuffle partition with union even if exists TableCacheQueryStage", ++ IgnoreComet("https://github.com/apache/datafusion-comet/issues/6454")) { + withSQLConf(SQLConf.CAN_CHANGE_CACHED_PLAN_OUTPUT_PARTITIONING.key -> "true", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1") { + val cached = Seq(1).toDF("c").cache() +@@ -3033,7 +3087,8 @@ class AdaptiveQueryExecSuite + } + } + +- test("SPARK-43376: Improve reuse subquery with table cache") { ++ test("SPARK-43376: Improve reuse subquery with table cache", ++ IgnoreComet("Comet's cache scan does not plan the subqueries in its pruning predicates")) { + withSQLConf(SQLConf.CAN_CHANGE_CACHED_PLAN_OUTPUT_PARTITIONING.key -> "true") { + withTable("t1", "t2") { + withCache("t1") { +@@ -3048,7 +3103,8 @@ class AdaptiveQueryExecSuite } } @@ -2703,7 +2932,7 @@ index fb8fab6a80f..403eb411920 100644 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")) -@@ -3138,7 +3192,8 @@ class AdaptiveQueryExecSuite +@@ -3138,7 +3194,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -2713,6 +2942,152 @@ index fb8fab6a80f..403eb411920 100644 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/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +index 4f07d3d1c03..c0e1c829bea 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +@@ -22,12 +22,13 @@ import java.sql.{Date, Timestamp} + import java.util.concurrent.atomic.AtomicInteger + + import org.apache.spark.rdd.RDD +-import org.apache.spark.sql.{QueryTest, Row} ++import org.apache.spark.sql.{IgnoreComet, QueryTest, Row} + import org.apache.spark.sql.catalyst.InternalRow + import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, In} + import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning + import org.apache.spark.sql.classic.DataFrame + import org.apache.spark.sql.columnar.CachedBatch ++import org.apache.spark.sql.comet.{CometFilterExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions._ +@@ -510,11 +511,17 @@ class InMemoryColumnarQuerySuite extends QueryTest + val planBeforeFilter = collect(df2.queryExecution.executedPlan) { + case f: FilterExec => f.child + case WholeStageCodegenExec(FilterExec(_, i: InputAdapter)) => i.child ++ case f: CometFilterExec => f.child + } +- assert(planBeforeFilter.head.isInstanceOf[InMemoryTableScanExec]) +- + val execPlan = planBeforeFilter.head +- assert(execPlan.executeCollectPublic().length == 0) ++ execPlan match { ++ // Comet's cache scan is columnar only, so count the rows of the batches it emits. ++ case c: CometInMemoryTableScanExec => ++ assert(c.executeColumnar().map(_.numRows().toLong).collect().sum == 0) ++ case _ => ++ assert(execPlan.isInstanceOf[InMemoryTableScanExec]) ++ assert(execPlan.executeCollectPublic().length == 0) ++ } + } + + test("SPARK-25727 - otherCopyArgs in InMemoryRelation does not include outputOrdering") { +@@ -523,7 +530,8 @@ class InMemoryColumnarQuerySuite extends QueryTest + assert(json.contains("outputOrdering")) + } + +- test("SPARK-22673: InMemoryRelation should utilize existing stats of the plan to be cached") { ++ test("SPARK-22673: InMemoryRelation should utilize existing stats of the plan to be cached", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + Seq("orc", "").foreach { useV1SourceReaderList => + // This test case depends on the size of ORC in statistics. + withSQLConf( +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +index 2c73622739a..5d0efeb263a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution.columnar + + import org.apache.spark.SparkFunSuite ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions.expr +@@ -37,11 +38,15 @@ class InMemoryRelationSuite extends SparkFunSuite + test("SPARK-47177: Cached SQL plan do not display final AQE plan in explain string") { + def findIMRInnerChild(p: SparkPlan): SparkPlan = { + val tableCache = find(p) { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + case _ => false + } + assert(tableCache.isDefined) +- tableCache.get.asInstanceOf[InMemoryTableScanExec].relation.innerChildren.head ++ val scan = tableCache.get match { ++ case c: CometInMemoryTableScanExec => c.originalPlan ++ case s => s.asInstanceOf[InMemoryTableScanExec] ++ } ++ scan.relation.innerChildren.head + } + + val d1 = spark.range(1).withColumn("key", expr("id % 100")) +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +index 88ff51d0ff4..7d00fca3b0a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +@@ -17,6 +17,7 @@ + + package org.apache.spark.sql.execution.columnar + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.internal.SQLConf + import org.apache.spark.sql.test.SharedSparkSession +@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + val result = df.collect().map(_(0)).toArray + assert(result.length === 1) + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- assert(readPartitions === 5) +- assert(readBatches === 10) ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ assert(scans.head.readPartitions.value === 5) ++ assert(scans.head.readBatches.value === 10) ++ } + } + + def checkBatchPruning( +@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + df.collect().map(_(0)).toArray + } + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- +- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") +- assert( +- readPartitions === expectedReadPartitions, +- s"Wrong number of read partitions: $queryExecution") ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan prunes on the same statistics, but has none of these test-only ++ // accumulators to read, so only the answer above is checked for it. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ val readPartitions = scans.head.readPartitions.value ++ val readBatches = scans.head.readBatches.value ++ assert( ++ readBatches === expectedReadBatches, ++ s"Wrong number of read batches: $queryExecution") ++ assert( ++ readPartitions === expectedReadPartitions, ++ s"Wrong number of read partitions: $queryExecution") ++ } + } + } + } 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 269990d7d14..140ee4112b1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala @@ -3856,10 +4231,10 @@ index f0f3f94b811..b7d18771314 100644 spark.internalCreateDataFrame(withoutFilters.execute(), schema) 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 720b13b812e..e3ac2cebc6e 100644 +index 720b13b812e..388dad6b4f3 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 -@@ -98,6 +98,20 @@ trait SharedSparkSessionBase +@@ -98,6 +98,25 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -3876,6 +4251,11 @@ index 720b13b812e..e3ac2cebc6e 100644 + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") + ++ // CometDriverPlugin installs Comet's cache serializer when ++ // spark.comet.exec.inMemoryCache.enabled is on, as it is by default. These sessions do not ++ // load the plugin, so install it here. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } conf.set( StaticSQLConf.WAREHOUSE_PATH, @@ -3998,10 +4378,10 @@ index b67370f6eb9..746b3974b29 100644 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 a394d0b7393..7056c74759c 100644 +index a394d0b7393..9c626b05ab2 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 -@@ -53,24 +53,34 @@ object TestHive +@@ -53,24 +53,38 @@ object TestHive new SparkContext( System.getProperty("spark.sql.test.master", "local[1]"), "TestSQLContext", @@ -4046,6 +4426,10 @@ index a394d0b7393..7056c74759c 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ ++ // As in SharedSparkSession: what CometDriverPlugin would install. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } + + conf diff --git a/dev/diffs/4.1.3.diff b/dev/diffs/4.1.3.diff index a8d9383ced..5de0a7f4f4 100644 --- a/dev/diffs/4.1.3.diff +++ b/dev/diffs/4.1.3.diff @@ -358,11 +358,41 @@ index 21a3ce1e122..f4762ab98f0 100644 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/CacheTableInKryoSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala +index 26d8f750f6e..c888f8e0844 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala +@@ -32,9 +32,16 @@ class CacheTableInKryoSuite extends QueryTest + with SharedSparkSession { + + override def sparkConf: SparkConf = { +- super.sparkConf ++ val conf = super.sparkConf + .set("spark.kryo.registrationRequired", "true") + .set("spark.serializer", "org.apache.spark.serializer.KryoSerializer") ++ // Comet's cache format needs its classes registered too, which is what Comet asks of any ++ // application that sets registrationRequired. ++ if (isCometEnabled) { ++ conf.set("spark.kryo.registrator", "org.apache.comet.CometKryoRegistrator") ++ } else { ++ conf ++ } + } + + test("SPARK-51777: sql.columnar.* classes registered in KryoSerializer") { 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 0d807aeae4d..6d7744e771b 100644 +index 0d807aeae4d..ce07fba149d 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 +@@ -37,6 +37,7 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException + import org.apache.spark.sql.catalyst.expressions.SubqueryExpression + import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, Join, JoinStrategyHint, SHUFFLE_HASH} + import org.apache.spark.sql.catalyst.util.DateTimeConstants ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.connector.catalog.BasicInMemoryTableCatalog + import org.apache.spark.sql.connector.catalog.CatalogPlugin + import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper +@@ -49,7 +50,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 @@ -371,7 +401,28 @@ index 0d807aeae4d..6d7744e771b 100644 import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -534,7 +534,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -128,6 +129,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + case inMemoryTable @ InMemoryTableScanExec(_, _, relation) => + getNumInMemoryTablesRecursively(relation.cachedPlan) + + getNumInMemoryTablesInSubquery(inMemoryTable) + 1 ++ // Comet's cache scan keeps the predicates pushed into it on its original plan, out of reach ++ // of subquery reuse. The filter above it evaluates the same subqueries, and counts them. ++ case cometScan: CometInMemoryTableScanExec => ++ getNumInMemoryTablesRecursively(cometScan.originalPlan.relation.cachedPlan) + 1 + case p => + getNumInMemoryTablesInSubquery(p) + }.sum +@@ -408,7 +413,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + assert(isExpectStorageLevel(rddId, Disk)) + } + +- test("InMemoryRelation statistics") { ++ test("InMemoryRelation statistics", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + sql("CACHE TABLE testData") + spark.table("testData").queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +@@ -534,7 +540,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils df.collect() } assert( @@ -381,7 +432,27 @@ index 0d807aeae4d..6d7744e771b 100644 } test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { -@@ -1673,9 +1674,18 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -1595,7 +1602,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + } + } + +- test("SPARK-36120: Support cache/uncache table with TimestampNTZ type") { ++ test("SPARK-36120: Support cache/uncache table with TimestampNTZ type", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + val tableName = "ntzCache" + withTable(tableName) { + sql(s"CACHE TABLE $tableName AS SELECT TIMESTAMP_NTZ'2021-01-01 00:00:00'") +@@ -1640,7 +1648,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + } + } + +- test("SPARK-35332: Make cache plan disable configs configurable - check AQE") { ++ test("SPARK-35332: Make cache plan disable configs configurable - check AQE", ++ IgnoreComet("Spark's SQL UI shows a cached plan only under Spark's own cache scan")) { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { +@@ -1673,9 +1682,18 @@ class CachedTableSuite extends QueryTest with SQLTestUtils _.nodeName.contains("TableCacheQueryStage")) val aqeNode = findNodeInSparkPlanInfo(inMemoryScanNode.get, _.nodeName.contains("AdaptiveSparkPlan")) @@ -403,6 +474,30 @@ index 0d807aeae4d..6d7744e771b 100644 } withTempView("t0", "t1", "t2") { +@@ -1764,6 +1782,7 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + val cached = spark.table("t") + val tableCache = collect(cached.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + if (expected == StorageLevel.NONE) { + assert(tableCache.isEmpty) +@@ -2630,6 +2649,7 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + + val inMemoryTableScan = collect(df.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + assert(inMemoryTableScan.size == 1) + checkAnswer(df, Row(5) :: Nil) +@@ -2657,6 +2677,7 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + + val subqueryInMemoryTableScan = collect(cteInSubquery.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + assert(subqueryInMemoryTableScan.size == 1) + checkAnswer(cteInSubquery, Row(1) :: Nil) 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 bfe15b33768..13aeb3f6610 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -460,7 +555,7 @@ index ed182322aec..1ae6afa686a 100644 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 93ff7becaec..87537a25b3b 100644 +index 93ff7becaec..27c366a9e4f 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 @@ -476,7 +571,17 @@ index 93ff7becaec..87537a25b3b 100644 import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.{ExamplePoint, ExamplePointUDT, SharedSparkSession, SQLTestData} -@@ -1519,11 +1520,12 @@ class DataFrameSetOperationsSuite extends QueryTest +@@ -1476,7 +1477,8 @@ class DataFrameSetOperationsSuite extends QueryTest + Row(Row(Seq(Seq(Row(null, "ba"))))) :: Nil) + } + +- test("SPARK-37371: UnionExec should support columnar if all children support columnar") { ++ test("SPARK-37371: UnionExec should support columnar if all children support columnar", ++ IgnoreComet("Comet replaces the cache scans and the union with its own operators")) { + def checkIfColumnar( + plan: SparkPlan, + targetPlan: (SparkPlan) => Boolean, +@@ -1519,11 +1521,12 @@ class DataFrameSetOperationsSuite extends QueryTest val union = df1.repartition($"a").union(df2.repartition($"a")) val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u @@ -490,7 +595,7 @@ index 93ff7becaec..87537a25b3b 100644 } assert(shuffle.size == 1) -@@ -1554,11 +1556,12 @@ class DataFrameSetOperationsSuite extends QueryTest +@@ -1554,11 +1557,12 @@ class DataFrameSetOperationsSuite extends QueryTest val union = df1.repartition($"a").union(df2.repartition($"d")) val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u @@ -504,7 +609,7 @@ index 93ff7becaec..87537a25b3b 100644 } assert(shuffle.size == 1) -@@ -1573,10 +1576,10 @@ class DataFrameSetOperationsSuite extends QueryTest +@@ -1573,10 +1577,10 @@ class DataFrameSetOperationsSuite extends QueryTest // Avoid unnecessary shuffle if union output partitioning is enabled val shuffledUnion = union.repartition($"a") val shuffleNumBefore = union.queryExecution.executedPlan.collect { @@ -517,7 +622,7 @@ index 93ff7becaec..87537a25b3b 100644 } if (enabled) { -@@ -1605,6 +1608,7 @@ class DataFrameSetOperationsSuite extends QueryTest +@@ -1605,6 +1609,7 @@ class DataFrameSetOperationsSuite extends QueryTest val union = df1.repartitionByRange($"a").union(df2.repartitionByRange($"d")) val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u @@ -658,6 +763,99 @@ index 4a070becfa6..61d515d127e 100644 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/DatasetCacheSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +index 627811eaecf..7269aeab949 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +@@ -22,7 +22,8 @@ import java.time.LocalTime + import org.scalatest.concurrent.TimeLimits + import org.scalatest.time.SpanSugar._ + +-import org.apache.spark.sql.execution.ColumnarToRowExec ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec ++import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec} + import org.apache.spark.sql.functions._ +@@ -39,6 +40,10 @@ class DatasetCacheSuite extends QueryTest + with AdaptiveSparkPlanHelper { + import testImplicits._ + ++ // A scan of a cached relation, Spark's or Comet's. ++ private def isCacheScan(plan: SparkPlan): Boolean = ++ plan.isInstanceOf[InMemoryTableScanExec] || plan.isInstanceOf[CometInMemoryTableScanExec] ++ + /** + * Asserts that a cached [[Dataset]] will be built using the given number of other cached results. + */ +@@ -46,7 +51,7 @@ class DatasetCacheSuite extends QueryTest + val plan = df.queryExecution.withCachedData + assert(plan.isInstanceOf[InMemoryRelation]) + val internalPlan = plan.asInstanceOf[InMemoryRelation].cacheBuilder.cachedPlan +- assert(find(internalPlan)(_.isInstanceOf[InMemoryTableScanExec]).size ++ assert(find(internalPlan)(isCacheScan).size + == numOfCachesDependedUpon) + } + +@@ -256,7 +261,7 @@ class DatasetCacheSuite extends QueryTest + case i: InMemoryRelation => i.cacheBuilder.cachedPlan + } + assert(df2LimitInnerPlan.isDefined && +- !df2LimitInnerPlan.get.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ !df2LimitInnerPlan.get.exists(isCacheScan)) + } + + test("SPARK-27739 Save stats from optimized plan") { +@@ -289,14 +294,14 @@ class DatasetCacheSuite extends QueryTest + val unionDf = df1.union(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.union(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("positive: union by name") { + val unionDf = df1.unionByName(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.unionByName(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by position") { +@@ -304,7 +309,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.union(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by name") { +@@ -312,7 +317,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.unionByName(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + } + } +@@ -324,12 +329,12 @@ class DatasetCacheSuite extends QueryTest + df1.cache() + // This is exactly the same as df1. + val df2 = spark.range(5).select(struct($"id".as("name", metadata))) +- assert(df2.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(df2.queryExecution.executedPlan.exists(isCacheScan)) + + val metadata2 = Metadata.fromJson("""{"k2": "v2"}""") + // Same with df1 except for the Alias metadata + val df3 = spark.range(5).select(struct($"id".as("name", metadata2))) +- assert(!df3.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!df3.queryExecution.executedPlan.exists(isCacheScan)) + } + + test("SPARK-53418: Handle TimeType in ColumnAccessor") { 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 6df8d66ee7f..35e270c7241 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -1247,10 +1445,18 @@ index cb9d0909554..084d6515e8b 100644 } 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 74cdee49e55..f7452c9abb7 100644 +index 74cdee49e55..6a544644c32 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 -@@ -1521,7 +1521,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -36,6 +36,7 @@ import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, NestedCo + import org.apache.spark.sql.catalyst.parser.ParseException + import org.apache.spark.sql.catalyst.plans.logical.{LocalLimit, Project, RepartitionByExpression, Sort} + import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.{CommandResultExec, OneRowRelationExec, UnionExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.aggregate._ +@@ -1521,7 +1522,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001"))) } @@ -1260,7 +1466,7 @@ index 74cdee49e55..f7452c9abb7 100644 AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") { sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect() } -@@ -1982,8 +1983,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -1982,8 +1984,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark countAcc.add(1) x }) @@ -1277,6 +1483,14 @@ index 74cdee49e55..f7452c9abb7 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) +@@ -3275,6 +3284,7 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark + + val inMemoryTableScan = collect(queryDf.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c + } + assert(inMemoryTableScan.size == 2) + checkAnswer(queryDf, Row(0, 1) :: Row(1, 2) :: Row(2, 3) :: Row(3, 4) :: Row(4, 5) :: Nil) 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 23f0144dcec..40d536bb23a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala @@ -1321,6 +1535,18 @@ index 23f0144dcec..40d536bb23a 100644 } } } +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 66826a9ca76..efeac3e4529 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 +@@ -239,6 +239,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt + df.select("i").filter($"i" > 1).cache() + assert(find(df.filter($"i" > 1).select("i").queryExecution.executedPlan) { + case _: org.apache.spark.sql.execution.columnar.InMemoryTableScanExec => true ++ case _: org.apache.spark.sql.comet.CometInMemoryTableScanExec => true + case _ => false + }.isDefined) + } 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 @@ -1336,6 +1562,20 @@ index d7b2511eac2..d5f5b940b94 100644 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/StatisticsCollectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +index 5222d5ce266..67f663c2764 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +@@ -829,7 +829,8 @@ class StatisticsCollectionSuite extends StatisticsCollectionTestBase with Shared + } + } + +- test("SPARK-33687: analyze all tables in a specific database") { ++ test("SPARK-33687: analyze all tables in a specific database", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + withTempDatabase { database => + spark.catalog.setCurrentDatabase(database) + withTempDir { dir => 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 7bfc8cf4fa6..4bd387801db 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala @@ -1535,7 +1775,7 @@ index 8a0e2c29653..d276a51cbc6 100644 } } 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 8f7a68bcbe6..88dbe1793c9 100644 +index 8f7a68bcbe6..c09c5d74309 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} @@ -1543,7 +1783,7 @@ index 8f7a68bcbe6..88dbe1793c9 100644 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.comet.{CometHashAggregateExec, CometInMemoryTableScanExec} 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 @@ -1602,6 +1842,15 @@ index 8f7a68bcbe6..88dbe1793c9 100644 }.head.isInstanceOf[ArrayTransform]) } } +@@ -1948,7 +1960,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { + // Checks in-memory fetching code path. + val all = sql("SELECT col FROM tbl") + assert(all.queryExecution.executedPlan.collectFirst { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + }.nonEmpty) + checkAnswer(all, Row("a")) + // Checks column stats code path. 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 a09b7e0827c..ffc29f764bc 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala @@ -2328,7 +2577,7 @@ index a3cfdc5a240..3793b6191bf 100644 }) 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 188a28ff1c0..8fdccf31749 100644 +index 188a28ff1c0..cfc84dcba01 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,12 +27,14 @@ import org.apache.spark.SparkException @@ -2813,7 +3062,27 @@ index 188a28ff1c0..8fdccf31749 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3229,7 +3282,8 @@ class AdaptiveQueryExecSuite +@@ -3184,7 +3237,8 @@ class AdaptiveQueryExecSuite + } + } + +- test("SPARK-42101: Coalesce shuffle partition with union even if exists TableCacheQueryStage") { ++ test("SPARK-42101: Coalesce shuffle partition with union even if exists TableCacheQueryStage", ++ IgnoreComet("https://github.com/apache/datafusion-comet/issues/6454")) { + withSQLConf(SQLConf.CAN_CHANGE_CACHED_PLAN_OUTPUT_PARTITIONING.key -> "true", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1") { + val cached = Seq(1).toDF("c").cache() +@@ -3214,7 +3268,8 @@ class AdaptiveQueryExecSuite + } + } + +- test("SPARK-43376: Improve reuse subquery with table cache") { ++ test("SPARK-43376: Improve reuse subquery with table cache", ++ IgnoreComet("Comet's cache scan does not plan the subqueries in its pruning predicates")) { + withSQLConf(SQLConf.CAN_CHANGE_CACHED_PLAN_OUTPUT_PARTITIONING.key -> "true") { + withTable("t1", "t2") { + withCache("t1") { +@@ -3229,7 +3284,8 @@ class AdaptiveQueryExecSuite } } @@ -2823,7 +3092,7 @@ index 188a28ff1c0..8fdccf31749 100644 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")) -@@ -3319,7 +3373,8 @@ class AdaptiveQueryExecSuite +@@ -3319,7 +3375,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -2834,10 +3103,29 @@ index 188a28ff1c0..8fdccf31749 100644 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 47b935a2880..3fdeab3113c 100644 +index 47b935a2880..65ee66c1975 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 -@@ -230,9 +230,21 @@ class CachedBatchSerializerNoUnwrapSuite extends QueryTest +@@ -212,6 +212,18 @@ class CachedBatchSerializerNoUnwrapSuite extends QueryTest + classOf[DefaultCachedBatchSerializerNoUnwrap].getName) + } + ++ // Spark keeps the first cache serializer it loads for the rest of the JVM, which the suites ++ // before this one set to Comet's. Clear it on both sides, as CachedBatchSerializerSuite does. ++ protected override def beforeAll(): Unit = { ++ super.beforeAll() ++ clearSerializer() ++ } ++ ++ protected override def afterAll(): Unit = { ++ clearSerializer() ++ super.afterAll() ++ } ++ + test("Do not unwrap ColumnarToRowExec") { + withTempPath { workDir => + val workDirPath = workDir.getAbsolutePath +@@ -230,9 +242,21 @@ class CachedBatchSerializerNoUnwrapSuite extends QueryTest assert(cachedPlans.length == 2) cachedPlans.foreach { cachedPlan => @@ -2862,6 +3150,152 @@ index 47b935a2880..3fdeab3113c 100644 } } } +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +index 4f07d3d1c03..c0e1c829bea 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +@@ -22,12 +22,13 @@ import java.sql.{Date, Timestamp} + import java.util.concurrent.atomic.AtomicInteger + + import org.apache.spark.rdd.RDD +-import org.apache.spark.sql.{QueryTest, Row} ++import org.apache.spark.sql.{IgnoreComet, QueryTest, Row} + import org.apache.spark.sql.catalyst.InternalRow + import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, In} + import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning + import org.apache.spark.sql.classic.DataFrame + import org.apache.spark.sql.columnar.CachedBatch ++import org.apache.spark.sql.comet.{CometFilterExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions._ +@@ -510,11 +511,17 @@ class InMemoryColumnarQuerySuite extends QueryTest + val planBeforeFilter = collect(df2.queryExecution.executedPlan) { + case f: FilterExec => f.child + case WholeStageCodegenExec(FilterExec(_, i: InputAdapter)) => i.child ++ case f: CometFilterExec => f.child + } +- assert(planBeforeFilter.head.isInstanceOf[InMemoryTableScanExec]) +- + val execPlan = planBeforeFilter.head +- assert(execPlan.executeCollectPublic().length == 0) ++ execPlan match { ++ // Comet's cache scan is columnar only, so count the rows of the batches it emits. ++ case c: CometInMemoryTableScanExec => ++ assert(c.executeColumnar().map(_.numRows().toLong).collect().sum == 0) ++ case _ => ++ assert(execPlan.isInstanceOf[InMemoryTableScanExec]) ++ assert(execPlan.executeCollectPublic().length == 0) ++ } + } + + test("SPARK-25727 - otherCopyArgs in InMemoryRelation does not include outputOrdering") { +@@ -523,7 +530,8 @@ class InMemoryColumnarQuerySuite extends QueryTest + assert(json.contains("outputOrdering")) + } + +- test("SPARK-22673: InMemoryRelation should utilize existing stats of the plan to be cached") { ++ test("SPARK-22673: InMemoryRelation should utilize existing stats of the plan to be cached", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + Seq("orc", "").foreach { useV1SourceReaderList => + // This test case depends on the size of ORC in statistics. + withSQLConf( +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +index 2c73622739a..5d0efeb263a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution.columnar + + import org.apache.spark.SparkFunSuite ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions.expr +@@ -37,11 +38,15 @@ class InMemoryRelationSuite extends SparkFunSuite + test("SPARK-47177: Cached SQL plan do not display final AQE plan in explain string") { + def findIMRInnerChild(p: SparkPlan): SparkPlan = { + val tableCache = find(p) { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + case _ => false + } + assert(tableCache.isDefined) +- tableCache.get.asInstanceOf[InMemoryTableScanExec].relation.innerChildren.head ++ val scan = tableCache.get match { ++ case c: CometInMemoryTableScanExec => c.originalPlan ++ case s => s.asInstanceOf[InMemoryTableScanExec] ++ } ++ scan.relation.innerChildren.head + } + + val d1 = spark.range(1).withColumn("key", expr("id % 100")) +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +index 88ff51d0ff4..7d00fca3b0a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +@@ -17,6 +17,7 @@ + + package org.apache.spark.sql.execution.columnar + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.internal.SQLConf + import org.apache.spark.sql.test.SharedSparkSession +@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + val result = df.collect().map(_(0)).toArray + assert(result.length === 1) + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- assert(readPartitions === 5) +- assert(readBatches === 10) ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ assert(scans.head.readPartitions.value === 5) ++ assert(scans.head.readBatches.value === 10) ++ } + } + + def checkBatchPruning( +@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + df.collect().map(_(0)).toArray + } + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- +- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") +- assert( +- readPartitions === expectedReadPartitions, +- s"Wrong number of read partitions: $queryExecution") ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan prunes on the same statistics, but has none of these test-only ++ // accumulators to read, so only the answer above is checked for it. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ val readPartitions = scans.head.readPartitions.value ++ val readBatches = scans.head.readBatches.value ++ assert( ++ readBatches === expectedReadBatches, ++ s"Wrong number of read batches: $queryExecution") ++ assert( ++ readPartitions === expectedReadPartitions, ++ s"Wrong number of read partitions: $queryExecution") ++ } + } + } + } 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 269990d7d14..140ee4112b1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala @@ -4155,10 +4589,10 @@ index f0f3f94b811..b7d18771314 100644 spark.internalCreateDataFrame(withoutFilters.execute(), schema) 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 720b13b812e..e3ac2cebc6e 100644 +index 720b13b812e..388dad6b4f3 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 -@@ -98,6 +98,20 @@ trait SharedSparkSessionBase +@@ -98,6 +98,25 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -4175,6 +4609,11 @@ index 720b13b812e..e3ac2cebc6e 100644 + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") + ++ // CometDriverPlugin installs Comet's cache serializer when ++ // spark.comet.exec.inMemoryCache.enabled is on, as it is by default. These sessions do not ++ // load the plugin, so install it here. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } conf.set( StaticSQLConf.WAREHOUSE_PATH, @@ -4297,10 +4736,10 @@ index b67370f6eb9..746b3974b29 100644 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 a394d0b7393..7056c74759c 100644 +index a394d0b7393..9c626b05ab2 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 -@@ -53,24 +53,34 @@ object TestHive +@@ -53,24 +53,38 @@ object TestHive new SparkContext( System.getProperty("spark.sql.test.master", "local[1]"), "TestSQLContext", @@ -4345,6 +4784,10 @@ index a394d0b7393..7056c74759c 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ ++ // As in SharedSparkSession: what CometDriverPlugin would install. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } + + conf diff --git a/dev/diffs/4.2.0.diff b/dev/diffs/4.2.0.diff index 60ad822d6d..ea06e8cbd4 100644 --- a/dev/diffs/4.2.0.diff +++ b/dev/diffs/4.2.0.diff @@ -376,11 +376,41 @@ index 21a3ce1e122..f4762ab98f0 100644 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/CacheTableInKryoSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala +index 72a2da16054..54146036f8a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala +@@ -30,9 +30,16 @@ import org.apache.spark.storage.StorageLevel + class CacheTableInKryoSuite extends SharedSparkSession { + + override def sparkConf: SparkConf = { +- super.sparkConf ++ val conf = super.sparkConf + .set("spark.kryo.registrationRequired", "true") + .set("spark.serializer", "org.apache.spark.serializer.KryoSerializer") ++ // Comet's cache format needs its classes registered too, which is what Comet asks of any ++ // application that sets registrationRequired. ++ if (isCometEnabled) { ++ conf.set("spark.kryo.registrator", "org.apache.comet.CometKryoRegistrator") ++ } else { ++ conf ++ } + } + + test("SPARK-51777: sql.columnar.* classes registered in KryoSerializer") { 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 +index 085dbcd8046..ac98c64ae92 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 +@@ -37,6 +37,7 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException + import org.apache.spark.sql.catalyst.expressions.SubqueryExpression + import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, Join, JoinStrategyHint, SHUFFLE_HASH} + import org.apache.spark.sql.catalyst.util.DateTimeConstants ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.connector.catalog.BasicInMemoryTableCatalog + import org.apache.spark.sql.connector.catalog.CatalogPlugin + import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper +@@ -49,7 +50,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 @@ -389,7 +419,28 @@ index 085dbcd8046..3090d321b6c 100644 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 +@@ -127,6 +128,10 @@ class CachedTableSuite extends SharedSparkSession + case inMemoryTable @ InMemoryTableScanExec(_, _, relation) => + getNumInMemoryTablesRecursively(relation.cachedPlan) + + getNumInMemoryTablesInSubquery(inMemoryTable) + 1 ++ // Comet's cache scan keeps the predicates pushed into it on its original plan, out of reach ++ // of subquery reuse. The filter above it evaluates the same subqueries, and counts them. ++ case cometScan: CometInMemoryTableScanExec => ++ getNumInMemoryTablesRecursively(cometScan.originalPlan.relation.cachedPlan) + 1 + case p => + getNumInMemoryTablesInSubquery(p) + }.sum +@@ -407,7 +412,8 @@ class CachedTableSuite extends SharedSparkSession + assert(isExpectStorageLevel(rddId, Disk)) + } + +- test("InMemoryRelation statistics") { ++ test("InMemoryRelation statistics", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + sql("CACHE TABLE testData") + spark.table("testData").queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +@@ -564,7 +570,8 @@ class CachedTableSuite extends SharedSparkSession df.collect() } assert( @@ -399,7 +450,27 @@ index 085dbcd8046..3090d321b6c 100644 } test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { -@@ -1703,9 +1704,18 @@ class CachedTableSuite extends SharedSparkSession +@@ -1625,7 +1632,8 @@ class CachedTableSuite extends SharedSparkSession + } + } + +- test("SPARK-36120: Support cache/uncache table with TimestampNTZ type") { ++ test("SPARK-36120: Support cache/uncache table with TimestampNTZ type", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + val tableName = "ntzCache" + withTable(tableName) { + sql(s"CACHE TABLE $tableName AS SELECT TIMESTAMP_NTZ'2021-01-01 00:00:00'") +@@ -1670,7 +1678,8 @@ class CachedTableSuite extends SharedSparkSession + } + } + +- test("SPARK-35332: Make cache plan disable configs configurable - check AQE") { ++ test("SPARK-35332: Make cache plan disable configs configurable - check AQE", ++ IgnoreComet("Spark's SQL UI shows a cached plan only under Spark's own cache scan")) { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { +@@ -1703,9 +1712,18 @@ class CachedTableSuite extends SharedSparkSession _.nodeName.contains("TableCacheQueryStage")) val aqeNode = findNodeInSparkPlanInfo(inMemoryScanNode.get, _.nodeName.contains("AdaptiveSparkPlan")) @@ -421,6 +492,30 @@ index 085dbcd8046..3090d321b6c 100644 } withTempView("t0", "t1", "t2") { +@@ -1794,6 +1812,7 @@ class CachedTableSuite extends SharedSparkSession + val cached = spark.table("t") + val tableCache = collect(cached.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + if (expected == StorageLevel.NONE) { + assert(tableCache.isEmpty) +@@ -2660,6 +2679,7 @@ class CachedTableSuite extends SharedSparkSession + + val inMemoryTableScan = collect(df.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + assert(inMemoryTableScan.size == 1) + checkAnswer(df, Row(5) :: Nil) +@@ -2687,6 +2707,7 @@ class CachedTableSuite extends SharedSparkSession + + val subqueryInMemoryTableScan = collect(cteInSubquery.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + assert(subqueryInMemoryTableScan.size == 1) + checkAnswer(cteInSubquery, Row(1) :: Nil) 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..f01366b66bc 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -489,7 +584,7 @@ index 9733d51a91c..395a108abc8 100644 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 +index d838ba4c234..4661d627a55 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 @@ -505,7 +600,17 @@ index d838ba4c234..cb0573d56d0 100644 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 +@@ -1475,7 +1476,8 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP + Row(Row(Seq(Seq(Row(null, "ba"))))) :: Nil) + } + +- test("SPARK-37371: UnionExec should support columnar if all children support columnar") { ++ test("SPARK-37371: UnionExec should support columnar if all children support columnar", ++ IgnoreComet("Comet replaces the cache scans and the union with its own operators")) { + def checkIfColumnar( + plan: SparkPlan, + targetPlan: (SparkPlan) => Boolean, +@@ -1518,11 +1520,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 @@ -519,7 +624,7 @@ index d838ba4c234..cb0573d56d0 100644 } assert(shuffle.size == 1) -@@ -1553,11 +1555,12 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP +@@ -1553,11 +1556,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 @@ -533,7 +638,7 @@ index d838ba4c234..cb0573d56d0 100644 } assert(shuffle.size == 1) -@@ -1572,10 +1575,10 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP +@@ -1572,10 +1576,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 { @@ -546,7 +651,7 @@ index d838ba4c234..cb0573d56d0 100644 } if (enabled) { -@@ -1604,6 +1607,7 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP +@@ -1604,6 +1608,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 @@ -687,6 +792,99 @@ index f79824de8ff..5432984960f 100644 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/DatasetCacheSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +index 0d1b0e1d981..e12c11e7b59 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +@@ -22,7 +22,8 @@ import java.time.LocalTime + import org.scalatest.concurrent.TimeLimits + import org.scalatest.time.SpanSugar._ + +-import org.apache.spark.sql.execution.ColumnarToRowExec ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec ++import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec} + import org.apache.spark.sql.functions._ +@@ -38,6 +39,10 @@ class DatasetCacheSuite extends SharedSparkSession + with AdaptiveSparkPlanHelper { + import testImplicits._ + ++ // A scan of a cached relation, Spark's or Comet's. ++ private def isCacheScan(plan: SparkPlan): Boolean = ++ plan.isInstanceOf[InMemoryTableScanExec] || plan.isInstanceOf[CometInMemoryTableScanExec] ++ + /** + * Asserts that a cached [[Dataset]] will be built using the given number of other cached results. + */ +@@ -45,7 +50,7 @@ class DatasetCacheSuite extends SharedSparkSession + val plan = df.queryExecution.withCachedData + assert(plan.isInstanceOf[InMemoryRelation]) + val internalPlan = plan.asInstanceOf[InMemoryRelation].cacheBuilder.cachedPlan +- assert(find(internalPlan)(_.isInstanceOf[InMemoryTableScanExec]).size ++ assert(find(internalPlan)(isCacheScan).size + == numOfCachesDependedUpon) + } + +@@ -255,7 +260,7 @@ class DatasetCacheSuite extends SharedSparkSession + case i: InMemoryRelation => i.cacheBuilder.cachedPlan + } + assert(df2LimitInnerPlan.isDefined && +- !df2LimitInnerPlan.get.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ !df2LimitInnerPlan.get.exists(isCacheScan)) + } + + test("SPARK-27739 Save stats from optimized plan") { +@@ -288,14 +293,14 @@ class DatasetCacheSuite extends SharedSparkSession + val unionDf = df1.union(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.union(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("positive: union by name") { + val unionDf = df1.unionByName(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.unionByName(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by position") { +@@ -303,7 +308,7 @@ class DatasetCacheSuite extends SharedSparkSession + unionDf.cache() + val finalDf = unionDf.union(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by name") { +@@ -311,7 +316,7 @@ class DatasetCacheSuite extends SharedSparkSession + unionDf.cache() + val finalDf = unionDf.unionByName(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + } + } +@@ -323,12 +328,12 @@ class DatasetCacheSuite extends SharedSparkSession + df1.cache() + // This is exactly the same as df1. + val df2 = spark.range(5).select(struct($"id".as("name", metadata))) +- assert(df2.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(df2.queryExecution.executedPlan.exists(isCacheScan)) + + val metadata2 = Metadata.fromJson("""{"k2": "v2"}""") + // Same with df1 except for the Alias metadata + val df3 = spark.range(5).select(struct($"id".as("name", metadata2))) +- assert(!df3.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!df3.queryExecution.executedPlan.exists(isCacheScan)) + } + + test("SPARK-53418: Handle TimeType in ColumnAccessor") { 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 @@ -1307,10 +1505,18 @@ index 291aa7cab72..7783c37683e 100644 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..c02b7c99490 100644 +index da6f6aca2ad..c62307049a2 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 +@@ -40,6 +40,7 @@ import org.apache.spark.sql.catalyst.parser.ParseException + import org.apache.spark.sql.catalyst.plans.logical.{LocalLimit, Project, RepartitionByExpression, Sort} + import org.apache.spark.sql.connector.catalog.CatalogManager + import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.{CommandResultExec, OneRowRelationExec, UnionExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.aggregate._ +@@ -1529,7 +1530,8 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001"))) } @@ -1320,7 +1526,7 @@ index da6f6aca2ad..c02b7c99490 100644 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 +@@ -1990,8 +1992,15 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper countAcc.add(1) x }) @@ -1337,6 +1543,14 @@ index da6f6aca2ad..c02b7c99490 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) +@@ -3283,6 +3292,7 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper + + val inMemoryTableScan = collect(queryDf.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c + } + assert(inMemoryTableScan.size == 2) + checkAnswer(queryDf, Row(0, 1) :: Row(1, 2) :: Row(2, 3) :: Row(3, 4) :: Row(4, 5) :: Nil) 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 @@ -1402,6 +1616,18 @@ index 395cb67f441..33ac6ed19af 100644 } } } +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..3228c78b603 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 +@@ -260,6 +260,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp + df.select("i").filter($"i" > 1).cache() + assert(find(df.filter($"i" > 1).select("i").queryExecution.executedPlan) { + case _: org.apache.spark.sql.execution.columnar.InMemoryTableScanExec => true ++ case _: org.apache.spark.sql.comet.CometInMemoryTableScanExec => true + case _ => false + }.isDefined) + } 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 @@ -1417,6 +1643,20 @@ index d7b2511eac2..d5f5b940b94 100644 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/StatisticsCollectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +index 5222d5ce266..67f663c2764 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +@@ -829,7 +829,8 @@ class StatisticsCollectionSuite extends StatisticsCollectionTestBase with Shared + } + } + +- test("SPARK-33687: analyze all tables in a specific database") { ++ test("SPARK-33687: analyze all tables in a specific database", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + withTempDatabase { database => + spark.catalog.setCurrentDatabase(database) + withTempDir { dir => 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 @@ -1616,7 +1856,7 @@ index 2d26356890d..2c5994f5fbc 100644 } } 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 +index 37684c7fce3..a0a2f72b6c9 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} @@ -1624,7 +1864,7 @@ index 37684c7fce3..f3574dec867 100644 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.comet.{CometHashAggregateExec, CometInMemoryTableScanExec} 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 @@ -1683,6 +1923,15 @@ index 37684c7fce3..f3574dec867 100644 }.head.isInstanceOf[ArrayTransform]) } } +@@ -2005,7 +2017,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { + // Checks in-memory fetching code path. + val all = sql("SELECT col FROM tbl") + assert(all.queryExecution.executedPlan.collectFirst { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + }.nonEmpty) + checkAnswer(all, Row("a")) + // Checks column stats code path. 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 @@ -2419,7 +2668,7 @@ index d70bd715879..074a9fa29d9 100644 }) 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..751ad50a569 100644 +index d6d19d21e65..702e77758fa 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 @@ -2895,7 +3144,27 @@ index d6d19d21e65..751ad50a569 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3278,7 +3330,8 @@ class AdaptiveQueryExecSuite +@@ -3233,7 +3285,8 @@ class AdaptiveQueryExecSuite + } + } + +- test("SPARK-42101: Coalesce shuffle partition with union even if exists TableCacheQueryStage") { ++ test("SPARK-42101: Coalesce shuffle partition with union even if exists TableCacheQueryStage", ++ IgnoreComet("https://github.com/apache/datafusion-comet/issues/6454")) { + withSQLConf(SQLConf.CAN_CHANGE_CACHED_PLAN_OUTPUT_PARTITIONING.key -> "true", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1") { + val cached = Seq(1).toDF("c").cache() +@@ -3263,7 +3316,8 @@ class AdaptiveQueryExecSuite + } + } + +- test("SPARK-43376: Improve reuse subquery with table cache") { ++ test("SPARK-43376: Improve reuse subquery with table cache", ++ IgnoreComet("Comet's cache scan does not plan the subqueries in its pruning predicates")) { + withSQLConf(SQLConf.CAN_CHANGE_CACHED_PLAN_OUTPUT_PARTITIONING.key -> "true") { + withTable("t1", "t2") { + withCache("t1") { +@@ -3278,7 +3332,8 @@ class AdaptiveQueryExecSuite } } @@ -2905,7 +3174,7 @@ index d6d19d21e65..751ad50a569 100644 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 +@@ -3368,7 +3423,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -2916,10 +3185,29 @@ index d6d19d21e65..751ad50a569 100644 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 +index 88be4adb6a4..23ec5374f8a 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 +@@ -210,6 +210,18 @@ class CachedBatchSerializerNoUnwrapSuite extends SharedSparkSession with Adaptiv + classOf[DefaultCachedBatchSerializerNoUnwrap].getName) + } + ++ // Spark keeps the first cache serializer it loads for the rest of the JVM, which the suites ++ // before this one set to Comet's. Clear it on both sides, as CachedBatchSerializerSuite does. ++ protected override def beforeAll(): Unit = { ++ super.beforeAll() ++ clearSerializer() ++ } ++ ++ protected override def afterAll(): Unit = { ++ clearSerializer() ++ super.afterAll() ++ } ++ + test("Do not unwrap ColumnarToRowExec") { + withTempPath { workDir => + val workDirPath = workDir.getAbsolutePath +@@ -228,9 +240,21 @@ class CachedBatchSerializerNoUnwrapSuite extends SharedSparkSession with Adaptiv assert(cachedPlans.length == 2) cachedPlans.foreach { cachedPlan => @@ -2944,6 +3232,152 @@ index 88be4adb6a4..f8fe831744e 100644 } } } +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +index 57da12e8797..413b5cf31a4 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +@@ -22,12 +22,13 @@ import java.sql.{Date, Timestamp} + import java.util.concurrent.atomic.AtomicInteger + + import org.apache.spark.rdd.RDD +-import org.apache.spark.sql.Row ++import org.apache.spark.sql.{IgnoreComet, Row} + import org.apache.spark.sql.catalyst.InternalRow + import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, In} + import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning + import org.apache.spark.sql.classic.DataFrame + import org.apache.spark.sql.columnar.CachedBatch ++import org.apache.spark.sql.comet.{CometFilterExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions._ +@@ -509,11 +510,17 @@ class InMemoryColumnarQuerySuite extends SharedSparkSession with AdaptiveSparkPl + val planBeforeFilter = collect(df2.queryExecution.executedPlan) { + case f: FilterExec => f.child + case WholeStageCodegenExec(FilterExec(_, i: InputAdapter)) => i.child ++ case f: CometFilterExec => f.child + } +- assert(planBeforeFilter.head.isInstanceOf[InMemoryTableScanExec]) +- + val execPlan = planBeforeFilter.head +- assert(execPlan.executeCollectPublic().length == 0) ++ execPlan match { ++ // Comet's cache scan is columnar only, so count the rows of the batches it emits. ++ case c: CometInMemoryTableScanExec => ++ assert(c.executeColumnar().map(_.numRows().toLong).collect().sum == 0) ++ case _ => ++ assert(execPlan.isInstanceOf[InMemoryTableScanExec]) ++ assert(execPlan.executeCollectPublic().length == 0) ++ } + } + + test("SPARK-25727 - otherCopyArgs in InMemoryRelation does not include outputOrdering") { +@@ -522,7 +529,8 @@ class InMemoryColumnarQuerySuite extends SharedSparkSession with AdaptiveSparkPl + assert(json.contains("outputOrdering")) + } + +- test("SPARK-22673: InMemoryRelation should utilize existing stats of the plan to be cached") { ++ test("SPARK-22673: InMemoryRelation should utilize existing stats of the plan to be cached", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + Seq("orc", "").foreach { useV1SourceReaderList => + // This test case depends on the size of ORC in statistics. + withSQLConf( +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +index 2c73622739a..5d0efeb263a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution.columnar + + import org.apache.spark.SparkFunSuite ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions.expr +@@ -37,11 +38,15 @@ class InMemoryRelationSuite extends SparkFunSuite + test("SPARK-47177: Cached SQL plan do not display final AQE plan in explain string") { + def findIMRInnerChild(p: SparkPlan): SparkPlan = { + val tableCache = find(p) { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + case _ => false + } + assert(tableCache.isDefined) +- tableCache.get.asInstanceOf[InMemoryTableScanExec].relation.innerChildren.head ++ val scan = tableCache.get match { ++ case c: CometInMemoryTableScanExec => c.originalPlan ++ case s => s.asInstanceOf[InMemoryTableScanExec] ++ } ++ scan.relation.innerChildren.head + } + + val d1 = spark.range(1).withColumn("key", expr("id % 100")) +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +index 88ff51d0ff4..7d00fca3b0a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +@@ -17,6 +17,7 @@ + + package org.apache.spark.sql.execution.columnar + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.internal.SQLConf + import org.apache.spark.sql.test.SharedSparkSession +@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + val result = df.collect().map(_(0)).toArray + assert(result.length === 1) + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- assert(readPartitions === 5) +- assert(readBatches === 10) ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ assert(scans.head.readPartitions.value === 5) ++ assert(scans.head.readBatches.value === 10) ++ } + } + + def checkBatchPruning( +@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + df.collect().map(_(0)).toArray + } + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- +- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") +- assert( +- readPartitions === expectedReadPartitions, +- s"Wrong number of read partitions: $queryExecution") ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan prunes on the same statistics, but has none of these test-only ++ // accumulators to read, so only the answer above is checked for it. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ val readPartitions = scans.head.readPartitions.value ++ val readBatches = scans.head.readBatches.value ++ assert( ++ readBatches === expectedReadBatches, ++ s"Wrong number of read batches: $queryExecution") ++ assert( ++ readPartitions === expectedReadPartitions, ++ s"Wrong number of read partitions: $queryExecution") ++ } + } + } + } 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 @@ -4302,10 +4736,10 @@ index e2c74533e7f..a12d55848ea 100644 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 +index fb26d3311eb..188b4d439a2 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 +@@ -102,6 +102,25 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -4322,6 +4756,11 @@ index fb26d3311eb..13dc8b89a91 100644 + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") + ++ // CometDriverPlugin installs Comet's cache serializer when ++ // spark.comet.exec.inMemoryCache.enabled is on, as it is by default. These sessions do not ++ // load the plugin, so install it here. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } conf.set( StaticSQLConf.WAREHOUSE_PATH, @@ -4400,10 +4839,10 @@ index 59022deaed7..f9aeacb5a9b 100644 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 +index 8e7ff526a95..518e1bae289 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 +@@ -54,24 +54,38 @@ object TestHive new SparkContext( System.getProperty("spark.sql.test.master", "local[1]"), "TestSQLContext", @@ -4448,6 +4887,10 @@ index 8e7ff526a95..ea1072a7195 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ ++ // As in SharedSparkSession: what CometDriverPlugin would install. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } + + conf diff --git a/docs/source/contributor-guide/plugin_overview.md b/docs/source/contributor-guide/plugin_overview.md index 89586bf7b6..072f59e01a 100644 --- a/docs/source/contributor-guide/plugin_overview.md +++ b/docs/source/contributor-guide/plugin_overview.md @@ -48,6 +48,8 @@ and skips the remaining steps. Otherwise it: - Appends `CometSparkSessionExtensions` to `spark.sql.extensions`, unless it is already listed. - Sets `spark.sql.cache.serializer` to Comet's `ArrowCachedBatchSerializer` when `spark.comet.exec.inMemoryCache.enabled=true`, unless the application has chosen a different serializer. + It leaves Spark's serializer in place where Comet could not scan its format natively or Kryo would + reject it; see [In-Memory Cache](../user-guide/latest/in-memory-cache.md). - Registers `CometSource` with Spark's metrics system and adds `CometMetricsListener` to `spark.sql.queryExecutionListeners` when `spark.comet.metrics.enabled=true`. - Logs a warning for settings that are likely to cause problems, such as an unset `spark.executor.memoryOverhead`. diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 53b290726a..0ea0ac5b3e 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -35,8 +35,8 @@ readable empty output files and their schema metadata. ## In-Memory Cache Comet can store cached relations (`df.cache()`, `CACHE TABLE`) in Arrow format and scan them -natively. This is experimental and disabled by default; see [In-Memory Cache](../in-memory-cache.md) -for how to enable it. Comet does not replace a `spark.sql.cache.serializer` that the application +natively. This is experimental and enabled by default; see [In-Memory Cache](../in-memory-cache.md) +for how to turn it off. Comet does not replace a `spark.sql.cache.serializer` that the application has already set. Relations whose schema Comet's Arrow writer does not support are cached in Spark's default format, and their scans fall back to Spark. Reads that feed Spark operators rather than Comet operators can be slower than Spark's cache. diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md index 8de4759957..01d975cc0e 100644 --- a/docs/source/user-guide/latest/in-memory-cache.md +++ b/docs/source/user-guide/latest/in-memory-cache.md @@ -21,24 +21,27 @@ Comet can store Spark's in-memory cache (`CACHE TABLE`, `df.cache()`, `df.persist()`) in an Arrow format that Comet operators read directly. Without it, a cached table is stored in Spark's own -format and every scan of it has to convert each batch before Comet can continue, which shows up in -the plan as a `CometSparkColumnarToColumnar` above the cache scan. +format, which Comet operators cannot read. Under Comet's default settings the operators above the +cache scan then run on Spark. With `spark.comet.sparkToColumnar.enabled`, a +`CometSparkColumnarToColumnar` above the scan converts each batch for Comet operators instead. -This feature is **experimental and disabled by default**. Turn it on at startup, alongside the rest -of Comet's configuration: +This feature is **experimental and enabled by default**. To turn it off, set the config at startup, +alongside the rest of Comet's configuration: ```shell $SPARK_HOME/bin/spark-shell \ ... \ - --conf spark.comet.exec.inMemoryCache.enabled=true + --conf spark.comet.exec.inMemoryCache.enabled=false ``` It has to be set before the `SparkContext` starts. Comet's driver plugin chooses -`spark.sql.cache.serializer` once, while the context is initializing, so a session that started -with the default goes on using Spark's cache format however the config is set afterwards. The -plugin installs Comet's serializer only if `spark.comet.enabled` and `spark.comet.exec.enabled` -are enabled at that point too, because an application that starts without native execution could -not scan Comet's format natively. +`spark.sql.cache.serializer` once, while the context is initializing, so an application keeps the +cache format it started with however the config is set afterwards. The plugin installs Comet's +serializer only if `spark.comet.enabled` and `spark.comet.exec.enabled` are enabled at that point +too, because an application that starts without native execution could not scan Comet's format +natively. It also keeps Spark's format when Comet shuffle is enabled but `spark.shuffle.manager` is +not one of Comet's shuffle managers, since Comet then disables itself, and when Kryo would reject +Comet's format; see [Kryo](#kryo). ## What changes when it is enabled @@ -109,7 +112,7 @@ nowhere to record either that a column is dictionary encoded or the dictionary i | Config | Default | Description | | ------------------------------------------------------- | ------- | ---------------------------------------------------------------------------------------------------------------------------------------------- | -| `spark.comet.exec.inMemoryCache.enabled` | `false` | Whether to store and scan Spark's in-memory cache in Comet's format. Read at startup. | +| `spark.comet.exec.inMemoryCache.enabled` | `true` | Whether to store and scan Spark's in-memory cache in Comet's format. Read at startup. | | `spark.comet.exec.inMemoryCache.compression.codec` | `zstd` | Arrow IPC compression codec for cached data: `zstd` or `none`. Affects newly cached data only — a batch records the codec it was written with. | | `spark.comet.exec.inMemoryCache.compression.zstd.level` | `1` | Compression level when the codec is `zstd`. Ignored otherwise. | @@ -158,10 +161,43 @@ back to Spark row execution above the scan and the two columns stop measuring th Read what this compares carefully. Comet execution is on in both columns, so the aggregation runs on Comet either way and only the cache-scan boundary moves: on the left, Spark's `InMemoryTableScanExec` feeds those same Comet operators through a `CometSparkColumnarToColumnar` -bridge; on the right, `CometInMemoryTableScan` feeds them directly. Both columns read the same +bridge, which the benchmark turns on with `spark.comet.sparkToColumnar.enabled`; on the right, +`CometInMemoryTableScan` feeds them directly. Both columns read the same Comet-written `CometCachedBatch`. These numbers are therefore "keep the cached scan native" against "fall back to a Spark cache scan and convert", not Comet against Spark execution, and not a -comparison with Spark's own cache format. That comparison is under [Limitations](#limitations). +comparison with Spark's own cache format, which follows. + +### Against Spark's cache format + +What turning the feature on changes for a query that Comet runs is measured against Spark's own +cache format by the benchmark's adaptive cases. Comet and AQE are on, Comet's other settings are at +their defaults, and the same 5M-row relation is cached in each format. The defaults leave +`spark.comet.sparkToColumnar.enabled` off, so Comet operators cannot read Spark's cache scan, and +with Spark's format the operators directly above the scan run on Spark. Measured on an AMD Ryzen 9 +7950X3D (JDK 17, Spark 4.1, release build): + +| Query shape | Spark's cache format | Comet's cache format | Relative | +| -------------------------- | -------------------: | -------------------: | -------: | +| Row count only (0 of 6) | 29 ms | 24 ms | 1.2x | +| Narrow projection (1 of 6) | 52 ms | 34 ms | 1.5x | +| 3 of 6 columns | 102 ms | 112 ms | 0.9x | +| Full projection (6 of 6) | 299 ms | 224 ms | 1.3x | + +A Spark operator above the cache scan, standing in for any operator Comet does not support, is +measured the same way, with Comet's aggregate turned off. With Comet's format, the native scan feeds +that operator through a columnar-to-row transition: + +| Query shape | Spark's cache format | Comet's cache format | Relative | +| -------------------------- | -------------------: | -------------------: | -------: | +| Row count only (0 of 6) | 39 ms | 16 ms | 2.4x | +| Narrow projection (1 of 6) | 50 ms | 27 ms | 1.8x | +| 3 of 6 columns | 97 ms | 112 ms | 0.9x | +| Full projection (6 of 6) | 303 ms | 299 ms | 1.0x | + +Comet's format is as fast or faster in every shape but one: the read of three of the six columns, +all of them longs, is about 10% slower under either kind of operator. That cost is `zstd` +decompression. With the `none` codec, the same read is 2.7x faster than Spark's format with Comet +operators above the scan, and 1.6x faster with a Spark operator above it. ## Kryo @@ -179,16 +215,27 @@ spark.kryo.registrator=org.apache.comet.CometKryoRegistrator Comet cannot set `spark.kryo.registrator` for you the way it sets `spark.sql.cache.serializer`: `KryoSerializer` reads it when `SparkEnv` builds the serializer, which happens before any plugin -runs. Without it, caching fails with a "Class is not registered" error that does not name this -feature. Comet's driver plugin warns at startup when it sees Kryo, `registrationRequired`, and no -registrator. Native broadcast needs the same registrator even when the cache is disabled; see +runs. Without it, Kryo would reject Comet's cached batch with a "Class is not registered" error +that does not name this feature. So when Kryo requires registration and has not registered +Comet's cached batch, Comet's driver plugin does not install Comet's serializer, and caches stay in +Spark's format. Registrations made another way, through a registrator of the application's own or +`spark.kryo.classesToRegister`, count as well. The plugin warns at startup when Kryo requires +registration and has not registered every class `CometKryoRegistrator` registers. An application +that sets `spark.sql.cache.serializer` to Comet's serializer itself gets the error instead. Native +broadcast needs the same registrator even when the cache is disabled; see [Kryo serialization](installation.md#kryo-serialization). +Spark registers its own cached batch with Kryo only from Spark 4.1, so on earlier versions caching +in either format under `registrationRequired` needs a registrator. `CometKryoRegistrator` registers +Spark's cached batch too. + ## Limitations -Reads that feed **Spark** operators rather than Comet ones are slower than Spark's own cache -format, and the narrower the read, the wider the gap. Measured by the same benchmark over the same -5M-row relation, with Comet off so that Spark operators consume the cached data: +Spark's own cache scan, `InMemoryTableScanExec`, reads Comet's format more slowly than Spark's, and +the narrower the read, the wider the gap. Spark's scan reads a cached relation when a session turns +Comet or its native execution off after caching, and when the relation's cached plan records +`Dataset.observe` metrics, and Comet records a fallback reason on the scan in either case. Measured +by the same benchmark over the same 5M-row relation, with Comet off: | Read shape | Spark's cache format | Comet's cache format | Slowdown | | ----------------------- | -------------------: | -------------------: | -------: | @@ -197,8 +244,9 @@ format, and the narrower the read, the wider the gap. Measured by the same bench | 3 of 6 columns | 98 ms | 331 ms | 3.4x | | 6 of 6 columns | 410 ms | 623 ms | 1.5x | -This is why the feature is off by default. The cause is not yet established; -[#5485](https://github.com/apache/datafusion-comet/issues/5485) tracks it. +A Spark operator above Comet's native cache scan does not pay this; see [Performance](#performance). +This gap is the main reason the feature is still described as experimental. The cause is not yet +established; [#5485](https://github.com/apache/datafusion-comet/issues/5485) tracks it. Comet's serializer exists because Spark's own Arrow cache format ([SPARK-57268](https://issues.apache.org/jira/browse/SPARK-57268)) is only available from Spark diff --git a/docs/source/user-guide/latest/installation.md b/docs/source/user-guide/latest/installation.md index b9d28b74b0..c6ca172ab0 100644 --- a/docs/source/user-guide/latest/installation.md +++ b/docs/source/user-guide/latest/installation.md @@ -269,8 +269,9 @@ If the application uses Kryo (`spark.serializer=org.apache.spark.serializer.Kryo Without it, any query that uses Comet's native broadcast exchange, which is enabled by default, fails with Kryo's "Class is not registered" error, for example on the first broadcast hash join. -The [in-memory cache](in-memory-cache.md#kryo) needs the same registrator. Set it before the -`SparkContext` is created: `KryoSerializer` reads it before Comet's plugin runs, so Comet cannot -add it for you. `spark.kryo.registrator` accepts a comma-separated list, so an application with -its own registrator can list both. Comet logs a warning at startup when Kryo requires registration -and this registrator is missing. +Comet's [in-memory cache](in-memory-cache.md#kryo) format needs the same registrations, and while +Kryo has not registered Comet's cached batch, Comet's plugin keeps caches in Spark's format. Set it +before the `SparkContext` is created: `KryoSerializer` reads it before Comet's plugin runs, so +Comet cannot add it for you. `spark.kryo.registrator` accepts a comma-separated list, so an +application with its own registrator can list both. Comet logs a warning at startup when Kryo +requires registration and has not registered the classes this registrator covers. diff --git a/docs/source/user-guide/latest/migration-guide.md b/docs/source/user-guide/latest/migration-guide.md index 342f8d2118..a28c600c8f 100644 --- a/docs/source/user-guide/latest/migration-guide.md +++ b/docs/source/user-guide/latest/migration-guide.md @@ -57,6 +57,35 @@ Treat setting one of these keys as a temporary measure. If you find you cannot s legacy behavior, please open an issue describing your use case so it can be considered before the key is removed. +## Upgrading to Comet 1.2.0 + +Comet `1.2.0` makes no behavior changes that need a `spark.comet.legacy.*` key. The changes below +need none either, but check whether any of them applies to your deployment. + +### In-Memory Cache Enabled by Default + +`spark.comet.exec.inMemoryCache.enabled` now defaults to `true`. An application that loads +`CometPlugin` now stores what it caches with `CACHE TABLE`, `df.cache()` or `df.persist()` in +Comet's Arrow format instead of Spark's, and Comet scans it natively. The format does not change +query results, but it can change performance. Spark's own cache scan reads Comet's format more +slowly than Spark's, which matters when a session turns Comet or its native execution off after +caching, and Comet records a fallback reason on such a scan. See +[In-Memory Cache](in-memory-cache.md#limitations). + +The format is chosen once, when the application starts. To keep Spark's format, set +`spark.comet.exec.inMemoryCache.enabled=false` then. Comet also keeps Spark's format without that +setting when the application: + +- starts with `spark.comet.enabled` or `spark.comet.exec.enabled` set to `false`. +- leaves Comet shuffle enabled without one of Comet's shuffle managers, so that Comet disables + itself. +- uses Kryo with `spark.kryo.registrationRequired=true` and has not registered Comet's cached + batch, because Kryo would reject it. To use Comet's format, register Comet's classes with + `spark.kryo.registrator=org.apache.comet.CometKryoRegistrator`; see + [Kryo](in-memory-cache.md#kryo). + +An application that sets `spark.sql.cache.serializer` itself keeps the serializer it chose. + ## Upgrading to Comet 1.1.0 Comet `1.1.0` makes no behavior changes that need a `spark.comet.legacy.*` key. The changes below diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index a47055df7e..db29377ff4 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -55,7 +55,7 @@ omitted from the tables below and may be reconsidered based on demand: | `BatchScanExec` | ✅ | Apache Iceberg Parquet scans run natively. Native CSV scans are experimental and disabled by default. DataSource V2 Parquet scans are not accelerated. See [Parquet Scan Compatibility](compatibility/scans.md) and the [Iceberg Guide](iceberg.md). | | `LocalTableScanExec` | ⚠️ | Disabled by default; there is no acceleration advantage and this operator is typically only used in test code. Can be opted into via config ([#4393](https://github.com/apache/datafusion-comet/pull/4393)). | | `EmptyRelationExec` | ✅ | Spark 4.0 and later. See [Empty Relations](compatibility/operators.md#empty-relations) for native-input support and writer fallback. | -| `InMemoryTableScanExec` | ⚠️ | Experimental, disabled by default. Set `spark.comet.exec.inMemoryCache.enabled=true` before the application starts so Comet installs its Arrow cache serializer. Relations with unsupported column types stay in Spark's cache format and fall back. See [In-Memory Cache](in-memory-cache.md). | +| `InMemoryTableScanExec` | ⚠️ | Experimental, enabled by default. Comet installs its Arrow cache serializer as the application starts, unless `spark.comet.exec.inMemoryCache.enabled` is false. Relations with unsupported column types stay in Spark's cache format and fall back. See [In-Memory Cache](in-memory-cache.md). | ## Projection and filtering diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6cfe5bc92b..12cbc18bff 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -275,7 +275,8 @@ object CometConf extends ShimCometConf { .doc("Whether to enable Comet native execution for in-memory cached tables. Its value at " + "startup also decides whether CometDriverPlugin installs Comet's cache serializer, " + "which stores cached data in Arrow format. The plugin installs it only if " + - "spark.comet.enabled and spark.comet.exec.enabled are also enabled at startup. " + + "spark.comet.enabled and spark.comet.exec.enabled are also enabled at startup, and " + + "only with one of Comet's shuffle managers while Comet shuffle is enabled. " + "Because spark.sql.cache.serializer is a " + "static config, the cached format is fixed for the application, and disabling this " + "at runtime only sends cached scans back to Spark's execution path. Relations whose " + @@ -284,12 +285,14 @@ object CometConf extends ShimCometConf { "zstd compression, and a scan copies out only the buffers of the columns it projected, " + "so the unselected ones are never decompressed. Reads that feed Spark operators rather " + "than Comet ones still pay a row conversion the default format avoids, and can be " + - "slower than Spark's cache. With spark.kryo.registrationRequired=true, also set " + - "spark.kryo.registrator=org.apache.comet.CometKryoRegistrator before creating the " + - "SparkContext, otherwise caching fails as soon as a block is serialized, including " + - "the disk half of the default MEMORY_AND_DISK storage level.") + "slower than Spark's cache. With spark.kryo.registrationRequired=true, the plugin " + + "installs it only if Kryo has registered Comet's cached batch, as " + + "spark.kryo.registrator=org.apache.comet.CometKryoRegistrator does when set before " + + "creating the SparkContext, because Kryo would otherwise reject a cached block as soon " + + "as it is serialized, including the disk half of the default MEMORY_AND_DISK storage " + + "level.") .booleanConf - .createWithDefault(false) + .createWithDefault(true) val COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC: ConfigEntry[String] = conf("spark.comet.exec.inMemoryCache.compression.codec") diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index c1aa237970..ce354aed33 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -62,6 +62,15 @@ import org.apache.comet.shims.{CometTypeShim, ShimCometStreaming, ShimCometWindo object CometExecRule { + /** + * Whether `scan` reads a relation stored in Comet's cache format. Comet's serializer stores + * that format only for schemas it supports and delegates everything else to Spark's default + * cache format, which the native scan cannot read. + */ + private[rules] def readsCometCacheFormat(scan: InMemoryTableScanExec): Boolean = + scan.relation.cacheBuilder.serializer.isInstanceOf[ArrowCachedBatchSerializer] && + ArrowCachedBatchSerializer.supportsSchema(scan.relation.output) + private[rules] def removePlaceholders(plan: SparkPlan): SparkPlan = plan.transformUp { // revertUnsafePartialAggregates re-runs transform over already wrapped query stages, which // can produce CometSinkPlaceHolder(CometSinkPlaceHolder(stage)). Remove sinks bottom-up. @@ -381,10 +390,7 @@ case class CometExecRule(session: SparkSession) case scan: InMemoryTableScanExec => val serializer = scan.relation.cacheBuilder.serializer val usesCometCacheSerializer = serializer.isInstanceOf[ArrowCachedBatchSerializer] - // The serializer only stores Comet's Arrow format for schemas it supports and delegates - // everything else to Spark's default cache format, which the native scan cannot read. - val cometCacheFormat = usesCometCacheSerializer && - ArrowCachedBatchSerializer.supportsSchema(scan.relation.output) + val cometCacheFormat = CometExecRule.readsCometCacheFormat(scan) val nativeCacheEnabled = CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.get(conf) // Walks the cached plan, so it is lazy: only consulted once the native scan is otherwise // possible. See CometInMemoryTableScanExec.recordsObservedMetrics. @@ -756,6 +762,25 @@ case class CometExecRule(session: SparkSession) } } + /** + * A relation keeps the cache format it was stored in, since `spark.sql.cache.serializer` is + * static, so a plan that runs without Comet's native execution still reads relations cached in + * Comet's format. Spark's `InMemoryTableScanExec` reads that format more slowly than Spark's + * own (https://github.com/apache/datafusion-comet/issues/5485), and nothing else records a + * fallback reason in such a plan, so record one on each scan that does. + */ + private def explainSparkReadsOfCometCache(plan: SparkPlan, cause: String): Unit = + plan.foreach { + case scan: InMemoryTableScanExec if CometExecRule.readsCometCacheFormat(scan) => + val _ = withFallbackReason( + scan, + s"$cause, so Spark reads this relation from Comet's cache format, which is slower " + + "than reading Spark's own. Set " + + s"${CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key}=false when the application " + + "starts to cache in Spark's format instead.") + case _ => + } + override def apply(plan: SparkPlan): SparkPlan = { val newPlan = _apply(plan) if (showTransformations && !newPlan.fastEquals(plan)) { @@ -769,13 +794,17 @@ case class CometExecRule(session: SparkSession) private def _apply(plan: SparkPlan): SparkPlan = { // We shouldn't transform Spark query plan if Comet is not loaded. - if (!isCometLoaded(conf)) return plan + if (!isCometLoaded(conf)) { + explainSparkReadsOfCometCache(plan, "Comet is disabled") + return plan + } // Comet does not support structured streaming. Fall back to Spark for any plan that // belongs to a streaming query (detected via StreamSourceAwareSparkPlan.getStream). if (ShimCometStreaming.isStreamingPlan(plan)) return plan if (!CometConf.COMET_EXEC_ENABLED.get(conf)) { + explainSparkReadsOfCometCache(plan, s"${CometConf.COMET_EXEC_ENABLED.key} is false") // Comet exec is disabled, but for Spark shuffle, we still can use Comet columnar shuffle if (isCometShuffleEnabled(conf)) { applyCometShuffle(plan) diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala index c3246679d0..e79074f45d 100644 --- a/spark/src/main/scala/org/apache/spark/Plugins.scala +++ b/spark/src/main/scala/org/apache/spark/Plugins.scala @@ -30,6 +30,9 @@ import org.apache.spark.api.plugin.{DriverPlugin, ExecutorPlugin, PluginContext, import org.apache.spark.internal.Logging import org.apache.spark.internal.config.{EVENT_LOG_ENABLED, EXECUTOR_MEMORY_OVERHEAD, EXECUTOR_MEMORY_OVERHEAD_FACTOR} import org.apache.spark.scheduler.{SparkListener, SparkListenerApplicationEnd, SparkListenerExecutorMetricsUpdate, SparkListenerExecutorRemoved} +import org.apache.spark.serializer.KryoSerializer +import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer +import org.apache.spark.sql.comet.execution.shuffle.{CometCelebornShuffleManager, CometShuffleManager} import org.apache.spark.sql.internal.StaticSQLConf import org.apache.spark.util.{Clock, SystemClock} @@ -110,7 +113,7 @@ class CometDriverPlugin private[spark] (clock: Clock) extends DriverPlugin with val extraConfs = new ju.HashMap[String, String]() CometDriverPlugin.maybeSetCacheSerializer(sc.conf, extraConfs) - CometDriverPlugin.warnIfKryoRegistratorMissing(sc.conf) + CometDriverPlugin.warnIfKryoRegistrationsMissing(sc.conf) // register CometSparkSessionExtensions if it isn't already registered CometDriverPlugin.registerCometSessionExtension(sc.conf) @@ -174,7 +177,13 @@ object CometDriverPlugin extends Logging { // Use Comet's cache serializer only when the native in-memory cache scan can run, which needs // Comet and its native execution as well as the cache config. spark.sql.cache.serializer is // static, so an application that starts with Comet or native execution off would otherwise - // store every cache in Comet's format, with only Spark operators to read it. + // store every cache in Comet's format, with only Spark operators to read it. So would one that + // leaves Comet shuffle enabled without Comet's shuffle manager, since Comet then disables + // itself. + // Nor is it used where Kryo requires registration and has not registered Comet's cached batch: + // caching would then fail the first time Spark serialized a cached block. Where Kryo has + // registered it, by whatever means, Comet's format is used, since Spark registers its own + // cached batch only from 4.1. // If the application already set spark.sql.cache.serializer, leave that value // unchanged so Comet does not replace a user-selected cache format. private[apache] def maybeSetCacheSerializer( @@ -182,7 +191,9 @@ object CometDriverPlugin extends Logging { extraConfs: ju.HashMap[String, String]): Unit = { if (getBooleanConf(conf, CometConf.COMET_ENABLED) && getBooleanConf(conf, CometConf.COMET_EXEC_ENABLED) && - getBooleanConf(conf, CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED)) { + getBooleanConf(conf, CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED) && + (!getBooleanConf(conf, CometConf.COMET_SHUFFLE_ENABLED) || isCometShuffleManager(conf)) && + !unregisteredKryoClasses(conf).contains(ArrowCachedBatchSerializer.cachedBatchClass)) { val serializerKey = StaticSQLConf.SPARK_CACHE_SERIALIZER.key val serializerValue = "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer" @@ -207,27 +218,54 @@ object CometDriverPlugin extends Logging { // CometKryoRegistrator covers both, but spark.kryo.registrator is read when SparkEnv builds the // serializer, before any plugin runs, so it cannot be set from here. Say so while the // application is still starting up rather than leaving the user to attribute the failure later. - private[apache] def warnIfKryoRegistratorMissing(conf: SparkConf): Unit = { + private[apache] def warnIfKryoRegistrationsMissing(conf: SparkConf): Unit = { + val unregistered = unregisteredKryoClasses(conf) + if (unregistered.nonEmpty) { + logWarning("spark.kryo.registrationRequired=true but Kryo has not registered " + + s"${unregistered.map(_.getName).mkString(", ")}, which " + + s"${CometKryoRegistrator.CLASS_NAME} registers. Comet's native broadcast and in-memory " + + "cache fail with Kryo's \"Class is not registered\" when they serialize one of them, " + + "and Comet keeps Spark's cache format while its own cached batch is unregistered. " + + s"Add spark.kryo.registrator=${CometKryoRegistrator.CLASS_NAME} before creating the " + + "SparkContext; it cannot be set later.") + } + } + + // The classes CometKryoRegistrator registers that Kryo, configured as the application + // configured it, would reject: none unless it requires registration. They can be registered + // through CometKryoRegistrator, a registrator of the application's own or + // spark.kryo.classesToRegister, so ask a Kryo instance built from the conf rather than read the + // confs. If one cannot be built, take them as registered only if spark.kryo.registrator lists + // CometKryoRegistrator. + private[apache] def unregisteredKryoClasses(conf: SparkConf): Seq[Class[_]] = { val usingKryo = conf.get("spark.serializer", "") == "org.apache.spark.serializer.KryoSerializer" - val registrationRequired = conf.getBoolean("spark.kryo.registrationRequired", false) - val registered = conf - .get("spark.kryo.registrator", "") - .split(',') - .map(_.trim) - .contains(CometKryoRegistrator.CLASS_NAME) - - if (usingKryo && registrationRequired && !registered) { - logWarning( - "spark.kryo.registrationRequired=true but spark.kryo.registrator does not include " + - s"${CometKryoRegistrator.CLASS_NAME}. Comet's native broadcast and its in-memory " + - "cache format will fail with Kryo's \"Class is not registered\" as soon as their " + - "payloads are serialized. Add " + - s"spark.kryo.registrator=${CometKryoRegistrator.CLASS_NAME} before creating the " + - "SparkContext; it cannot be set later.") + if (!usingKryo || !conf.getBoolean("spark.kryo.registrationRequired", false)) { + Nil + } else { + // Qualified, because in this package org.apache.spark.Success, a TaskEndReason, hides an + // imported scala.util.Success on Scala 2.12. + Try(new KryoSerializer(conf).newKryo()) match { + case scala.util.Success(kryo) => + CometKryoRegistrator.classes.filter(kryo.getClassResolver.getRegistration(_) == null) + case scala.util.Failure(e) => + logDebug("Could not build Kryo to check Comet's registrations", e) + val listed = conf + .get("spark.kryo.registrator", "") + .split(',') + .map(_.trim) + .contains(CometKryoRegistrator.CLASS_NAME) + if (listed) Nil else CometKryoRegistrator.classes + } } } + // Comet's shuffle managers have no short name, so spark.shuffle.manager names one only by its + // class name. + private def isCometShuffleManager(conf: SparkConf): Boolean = + Set(classOf[CometShuffleManager].getName, classOf[CometCelebornShuffleManager].getName) + .contains(conf.get("spark.shuffle.manager", "")) + // Comet's native allocations are made by the Rust global allocator and live in the native heap. // In off-heap mode the share that operators reserve is charged against a memory pool, but // everything else -- expression kernels and Arrow array builders, decompression buffers, Parquet @@ -298,8 +336,13 @@ object CometDriverPlugin extends Logging { } } + // Reads a deprecated alternative too, such as spark.comet.exec.shuffle.enabled, as a session + // would. private def getBooleanConf(conf: SparkConf, entry: ConfigEntry[Boolean]): Boolean = - conf.getBoolean(entry.key, entry.defaultValue.get) + (entry.key +: entry.alternatives) + .find(conf.contains) + .map(conf.getBoolean(_, entry.defaultValue.get)) + .getOrElse(entry.defaultValue.get) def registerCometMetrics(sc: SparkContext): Unit = { if (sc.getConf.getBoolean( diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala index 103df61e17..853c731a94 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala @@ -739,6 +739,11 @@ object ArrowCachedBatchSerializer { def supportsSchema(schema: Seq[Attribute]): Boolean = schema.forall(a => supportsType(a.dataType)) + /** + * The class of Comet's cached batch, which Kryo has to have registered to store this format. + */ + private[apache] val cachedBatchClass: Class[_] = classOf[CometCachedBatch] + /** * The classes a `CometCachedBatch` adds on top of [[org.apache.comet.CometKryoRegistrator]]'s * shared Arrow-bytes classes. diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala index b403442531..4518f8ab2e 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala @@ -20,10 +20,16 @@ package org.apache.comet.exec import org.apache.spark.SparkConf +import org.apache.spark.serializer.KryoRegistrator import org.apache.spark.sql.{CometTestBase, Row} -import org.apache.spark.sql.execution.columnar.CometInMemoryRelationHelper +import org.apache.spark.sql.catalyst.expressions.GenericInternalRow +import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer +import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, DefaultCachedBatch} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.storage.StorageLevel +import org.apache.spark.unsafe.types.UTF8String + +import com.esotericsoftware.kryo.Kryo import org.apache.comet.{CometConf, CometKryoRegistrator} @@ -194,3 +200,106 @@ class CometInMemoryCacheKryoSuite extends CometTestBase { } } } + +/** + * Comet's driver plugin under Kryo with registration required, in an application that does not + * list [[CometKryoRegistrator]]. The plugin installs Comet's cache serializer only if Kryo has + * Comet's cached batch registered, whatever registered it, so the format it picks has to survive + * a `DISK_ONLY` cache, which serializes every block as it is put. + * + * Spark registers its own cached batch with Kryo only from 4.1, so on earlier versions an + * application whose caches work under registration registers it itself, as the suite that expects + * Spark's format does. + */ +abstract class CometInMemoryCacheKryoRegistrationSuite(expectedBatch: String) + extends CometTestBase { + + /** The Kryo registrations the application makes instead of listing CometKryoRegistrator. */ + protected def registrations: Seq[(String, String)] + + override protected def beforeAll(): Unit = { + CometInMemoryRelationHelper.clearSerializer() + super.beforeAll() + } + + override protected def afterAll(): Unit = { + try { + super.afterAll() + } finally { + CometInMemoryRelationHelper.clearSerializer() + } + } + + override protected def sparkConf: SparkConf = { + val conf = super.sparkConf + conf.set("spark.plugins", "org.apache.spark.CometPlugin") + conf.set(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key, "true") + conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer") + conf.set("spark.kryo.registrationRequired", "true") + conf.setAll(registrations) + } + + private def cachedBatchTypes(table: String): Array[String] = { + val cached = spark.sharedState.cacheManager.lookupCachedData(spark.table(table)).get + cached.cachedRepresentation.cacheBuilder.cachedColumnBuffers + .map(_.getClass.getName) + .distinct() + .collect() + } + + test("Comet plugin picks a cache format that Kryo can store") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + spark.catalog.clearCache() + try { + spark + .range(0, 100, 1, 2) + .selectExpr("id", "cast(id as string) AS s") + .createOrReplaceTempView("kryo_registration") + + spark.catalog.cacheTable("kryo_registration", StorageLevel.DISK_ONLY) + assert(spark.table("kryo_registration").count() == 100) + assert(cachedBatchTypes("kryo_registration").sameElements(Array(expectedBatch))) + + checkAnswer( + spark.sql("SELECT s FROM kryo_registration WHERE id > 97"), + Seq(Row("98"), Row("99"))) + } finally { + spark.catalog.clearCache() + } + } + } +} + +/** Registers Spark's cached batch and nothing of Comet's, so the plugin keeps Spark's format. */ +class CometInMemoryCacheKryoUnregisteredSuite + extends CometInMemoryCacheKryoRegistrationSuite(classOf[DefaultCachedBatch].getName) { + override protected def registrations: Seq[(String, String)] = + Seq("spark.kryo.registrator" -> classOf[SparkCachedBatchKryoRegistrator].getName) +} + +/** + * Registers Comet's classes through `spark.kryo.classesToRegister` rather than + * [[CometKryoRegistrator]], and not Spark's cached batch, so before Spark 4.1 only Comet's format + * can be stored, and the plugin has to install it. + */ +class CometInMemoryCacheKryoClassesToRegisterSuite + extends CometInMemoryCacheKryoRegistrationSuite( + ArrowCachedBatchSerializer.cachedBatchClass.getName) { + override protected def registrations: Seq[(String, String)] = Seq( + "spark.kryo.classesToRegister" -> CometKryoRegistrator.classes + .filterNot(_ == classOf[DefaultCachedBatch]) + .map(_.getName) + .mkString(",")) +} + +/** Registers what Spark's own cached batch needs, over a long and a string column, with Kryo. */ +class SparkCachedBatchKryoRegistrator extends KryoRegistrator { + override def registerClasses(kryo: Kryo): Unit = { + // The batch and its statistics row, whose bounds include UTF8String. + Seq( + classOf[DefaultCachedBatch], + classOf[GenericInternalRow], + classOf[Array[Any]], + classOf[UTF8String]).foreach(kryo.register) + } +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index ecce906752..9c25069645 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -39,6 +39,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeRefer import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec, CometSortExec, CometSortMergeJoinExec} import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} +import org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.SortExec import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec} @@ -51,7 +52,7 @@ import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.spark.storage.StorageLevel -import org.apache.comet.{CometArrowAllocator, CometConf, ExtendedExplainInfo} +import org.apache.comet.{CometArrowAllocator, CometConf, CometKryoRegistrator, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} import org.apache.comet.vector.{CometPlainVector, CometVector} @@ -474,6 +475,54 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Comet explains Spark's scan of a relation cached in Comet's format") { + // spark.sql.cache.serializer is static, so a relation cached in Comet's format stays in it + // after a session turns Comet or its native execution off, and from then on Spark's + // InMemoryTableScanExec reads it, which nothing else in the plan would record. A relation + // that Comet's serializer delegated to Spark's format gets no such reason. + withNativeCache { + spark + .sql("SELECT id, id % 7 AS k FROM range(100)") + .createOrReplaceTempView("comet_format_cache") + spark + .sql(s"SELECT id, ${unsupportedForArrowCache.head} FROM range(100)") + .createOrReplaceTempView("spark_format_cache") + spark.catalog.cacheTable("comet_format_cache") + spark.catalog.cacheTable("spark_format_cache") + assert( + cachedBatchTypes("comet_format_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + assert( + cachedBatchTypes("spark_format_cache").sameElements( + Array("org.apache.spark.sql.execution.columnar.DefaultCachedBatch"))) + + def reasons(query: String): Seq[String] = { + val df = spark.sql(query) + df.collect() + new ExtendedExplainInfo().getFallbackReasons(df.queryExecution.executedPlan) + } + + for { + (key, cause) <- Seq( + CometConf.COMET_ENABLED.key -> "Comet is disabled", + CometConf.COMET_EXEC_ENABLED.key -> s"${CometConf.COMET_EXEC_ENABLED.key} is false") + aqe <- Seq("false", "true") + } { + withSQLConf(key -> "false", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) { + val explained = reasons("SELECT k, count(*) FROM comet_format_cache GROUP BY k") + assert( + explained.exists( + _.startsWith(s"$cause, so Spark reads this relation from Comet's cache format")), + s"$key=false, AQE $aqe: $explained") + assert( + !reasons("SELECT count(id) FROM spark_format_cache").exists( + _.contains("Comet's cache format")), + s"$key=false, AQE $aqe") + } + } + } + } + test("Comet in-memory cache handles multi-partition cache") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", @@ -1021,6 +1070,7 @@ class CometInMemoryCacheSuite extends CometTestBase { val defaultConf = new SparkConf() .set(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key, "true") + .set("spark.shuffle.manager", shuffleManager) val defaultExtraConfs = new ju.HashMap[String, String]() // With no user serializer configured, the plugin should install Comet's @@ -1032,6 +1082,7 @@ class CometInMemoryCacheSuite extends CometTestBase { val userConf = new SparkConf() .set(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key, "true") + .set("spark.shuffle.manager", shuffleManager) .set(serializerKey, userSerializer) val userExtraConfs = new ju.HashMap[String, String]() @@ -1043,20 +1094,25 @@ class CometInMemoryCacheSuite extends CometTestBase { assert(!userExtraConfs.containsKey(serializerKey)) } - test("Comet plugin installs its cache serializer only if Comet can scan the cache natively") { + /** Whether the Comet plugin installs its cache serializer for an application's `settings`. */ + private def installsCacheSerializer(settings: (String, String)*): Boolean = { val serializerKey = StaticSQLConf.SPARK_CACHE_SERIALIZER.key + val conf = new SparkConf().setAll(settings) + val extraConfs = new ju.HashMap[String, String]() + CometDriverPlugin.maybeSetCacheSerializer(conf, extraConfs) + assert(conf.contains(serializerKey) == extraConfs.containsKey(serializerKey)) + extraConfs.containsKey(serializerKey) + } - def installed(settings: (String, String)*): Boolean = { - val conf = new SparkConf().setAll(settings) - val extraConfs = new ju.HashMap[String, String]() - CometDriverPlugin.maybeSetCacheSerializer(conf, extraConfs) - assert(conf.contains(serializerKey) == extraConfs.containsKey(serializerKey)) - extraConfs.containsKey(serializerKey) - } - + test("Comet plugin installs its cache serializer only if Comet can scan the cache natively") { val cometOn = CometConf.COMET_ENABLED.key -> "true" val execOn = CometConf.COMET_EXEC_ENABLED.key -> "true" val cacheOn = CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true" + // Without Comet's shuffle manager Comet disables itself, which the next test covers. + val cometShuffle = "spark.shuffle.manager" -> shuffleManager + + def installed(settings: (String, String)*): Boolean = + installsCacheSerializer(cometShuffle +: settings: _*) assert(installed(cometOn, execOn, cacheOn)) // An application that starts with Comet or its native execution off can never plan @@ -1074,6 +1130,85 @@ class CometInMemoryCacheSuite extends CometTestBase { CometConf.COMET_EXEC_ENABLED.defaultValue.get)) } + test("Comet plugin keeps Spark's cache format where Comet disables itself or Kryo rejects it") { + val cometShuffle = "spark.shuffle.manager" -> shuffleManager + + def installed(settings: (String, String)*): Boolean = + installsCacheSerializer( + Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") ++ settings: _*) + + assert(installed(cometShuffle)) + // Comet shuffle is enabled by default, and Comet disables itself while it is unless the + // application runs one of Comet's shuffle managers. + assert(!installed()) + assert(!installed("spark.shuffle.manager" -> "sort")) + assert(installed("spark.shuffle.manager" -> classOf[CometCelebornShuffleManager].getName)) + // Without Comet shuffle, under the key or its deprecated name, the shuffle manager does not + // matter. + assert(installed(CometConf.COMET_SHUFFLE_ENABLED.key -> "false")) + assert(installed("spark.comet.exec.shuffle.enabled" -> "false")) + + // Kryo with registration required rejects Comet's cached batch unless something registered + // it: CometKryoRegistrator, on its own or beside a registrator of the application's, or the + // application's own registrations. + val kryo = "spark.serializer" -> "org.apache.spark.serializer.KryoSerializer" + val registrationRequired = "spark.kryo.registrationRequired" -> "true" + val registrator = "spark.kryo.registrator" + val sparkOnly = classOf[SparkCachedBatchKryoRegistrator].getName + assert(!installed(cometShuffle, kryo, registrationRequired)) + assert(!installed(cometShuffle, kryo, registrationRequired, registrator -> sparkOnly)) + assert( + installed( + cometShuffle, + kryo, + registrationRequired, + registrator -> s"$sparkOnly, ${CometKryoRegistrator.CLASS_NAME}")) + assert( + installed( + cometShuffle, + kryo, + registrationRequired, + "spark.kryo.classesToRegister" -> ArrowCachedBatchSerializer.cachedBatchClass.getName)) + // A registrator that cannot be loaded leaves only spark.kryo.registrator to go by. + assert(!installed(cometShuffle, kryo, registrationRequired, registrator -> "com.example.R")) + assert( + installed( + cometShuffle, + kryo, + registrationRequired, + registrator -> s"com.example.R, ${CometKryoRegistrator.CLASS_NAME}")) + // Without registrationRequired, Kryo writes the class name of anything unregistered instead. + assert(installed(cometShuffle, kryo)) + } + + test("Comet plugin finds the Kryo registrations Comet needs however they were made") { + def unregistered(settings: (String, String)*): Seq[Class[_]] = + CometDriverPlugin.unregisteredKryoClasses(new SparkConf().setAll(settings)) + + val kryo = "spark.serializer" -> "org.apache.spark.serializer.KryoSerializer" + val registrationRequired = "spark.kryo.registrationRequired" -> "true" + assert(unregistered().isEmpty) + assert(unregistered(kryo).isEmpty) + assert( + unregistered(kryo, registrationRequired).contains( + ArrowCachedBatchSerializer.cachedBatchClass)) + assert( + unregistered( + kryo, + registrationRequired, + "spark.kryo.registrator" -> CometKryoRegistrator.CLASS_NAME).isEmpty) + assert( + unregistered( + kryo, + registrationRequired, + "spark.kryo.classesToRegister" -> CometKryoRegistrator.classes + .map(_.getName) + .mkString(",")).isEmpty) + } + test("Comet in-memory cache supports empty projection scan") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala index 9b0f652a19..8c8bbcfbe1 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala @@ -27,6 +27,8 @@ import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.AttributeReference import org.apache.spark.sql.comet.CometInMemoryTableScanExec import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec +import org.apache.spark.sql.execution.aggregate.HashAggregateExec import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, DefaultCachedBatchSerializer, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} @@ -236,6 +238,7 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { runCodecBenchmark(flatRelation) runSparkOperatorBenchmark(flatRelation) + runAdaptiveBenchmark(flatRelation) } } @@ -346,9 +349,9 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { * Reads that feed Spark operators rather than Comet ones, against Spark's own cache format. * * Comet is off in every case, so this measures Spark consuming the cached data: the shape where - * Comet's format has something to lose, and the reason the feature is off by default. Both - * formats are cached from the same relation, one copy at a time as in runCodecBenchmark, and - * each case checks which serializer cached the relation it reads. + * Comet's format has the most to lose. Both formats are cached from the same relation, one copy + * at a time as in runCodecBenchmark, and each case checks which serializer cached the relation + * it reads. */ private def runSparkOperatorBenchmark(relation: CachedRelation): Unit = { val view = s"${relation.table}_spark_operators" @@ -375,14 +378,7 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { cachedBy = serializer } - Seq( - ("row count only (0 of 6 columns)", s"SELECT count(*) FROM $view", 0), - ("narrow projection (1 of 6 columns)", s"SELECT count(k) FROM $view", 1), - ("3 of 6 columns", s"SELECT sum(id), sum(k), sum(v) FROM $view", 3), - ( - "full projection (6 of 6 columns)", - s"SELECT count(id), count(k), count(v), count(s1), count(s2), count(s3) FROM $view", - 6)).foreach { case (label, query, scanned) => + readShapes(view).foreach { case (label, query, scanned) => val benchmark = new Benchmark( s"in-memory cache read by Spark operators, $label", relation.rows.toLong, @@ -411,6 +407,129 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { } } + /** + * What the feature changes for a query that runs with Comet, against Spark's own cache format, + * with AQE on and Comet's other settings at their defaults. This is the comparison an + * application gets from turning the feature on or off. Both formats are cached from the same + * relation, one copy at a time as in runSparkOperatorBenchmark. + * + * Two shapes of plan read the cache. With Comet operators above the cache scan, Comet's format + * runs the whole query natively, while Spark's leaves the operators directly above its scan on + * Spark: Comet reads Spark's cache scan only through spark.comet.sparkToColumnar.enabled, which + * is off by default. With a Spark operator above the scan, Comet's format is read by the native + * scan and converted to rows for that operator, where Spark's is read by Spark's own scan. The + * Spark operator is the aggregate, with Comet's turned off, standing in for any operator Comet + * does not support. + */ + private def runAdaptiveBenchmark(relation: CachedRelation): Unit = { + val view = s"${relation.table}_adaptive" + val formats = Seq( + "Spark's cache format" -> classOf[DefaultCachedBatchSerializer].getName, + "Comet's cache format" -> classOf[ArrowCachedBatchSerializer].getName) + val operatorsAbove = Seq( + "Comet operators" -> Seq.empty[(String, String)], + "a Spark operator" -> Seq(CometConf.COMET_EXEC_AGGREGATE_ENABLED.key -> "false")) + + spark.catalog.clearCache() + withTempTable(view) { + spark + .sql(s"SELECT ${relation.columns.mkString(", ")} FROM ${relation.source}") + .createOrReplaceTempView(view) + + var cachedBy: String = null + def cacheBy(serializer: String): Unit = if (cachedBy != serializer) { + spark.catalog.uncacheTable(view) + cachedBy = null + withCacheSerializer(serializer) { + withSQLConf(adaptiveConf: _*) { + spark.catalog.cacheTable(view) + spark.table(view).count() + } + } + cachedBy = serializer + } + + for { + (operators, operatorConf) <- operatorsAbove + (label, query, scanned) <- readShapes(view) + } { + val benchmark = new Benchmark( + s"in-memory cache with AQE, $operators above the scan, $label", + relation.rows.toLong, + output = output) + formats.foreach { case (name, serializer) => + var verified = false + // Re-caching in this case's format is setup, so it is outside the timer, and it only + // happens on the case's first call, which is a warmup iteration. + benchmark.addTimerCase(name) { timer => + cacheBy(serializer) + withSQLConf(adaptiveConf ++ operatorConf: _*) { + if (!verified) { + verifyAdaptiveRead(query, scanned, serializer, operatorConf.nonEmpty) + verified = true + } + timer.startTiming() + spark.sql(query).noop() + timer.stopTiming() + } + } + } + benchmark.run() + } + + spark.catalog.uncacheTable(view) + } + } + + // The reads runSparkOperatorBenchmark and runAdaptiveBenchmark measure: no columns of the flat + // relation, one, three, and all six. + private def readShapes(view: String): Seq[(String, String, Int)] = Seq( + ("row count only (0 of 6 columns)", s"SELECT count(*) FROM $view", 0), + ("narrow projection (1 of 6 columns)", s"SELECT count(k) FROM $view", 1), + ("3 of 6 columns", s"SELECT sum(id), sum(k), sum(v) FROM $view", 3), + ( + "full projection (6 of 6 columns)", + s"SELECT count(id), count(k), count(v), count(s1), count(s2), count(s3) FROM $view", + 6)) + + // Pins what an adaptive case claims, in the plan AQE settles on, which it does only by running + // the query: one cache scan, native exactly when it reads Comet's format, reading the columns its + // label counts from a relation the named serializer cached. Spark aggregates run only where the + // case puts them: above Spark's scan, which nothing bridges into Comet, or wherever Comet's + // aggregate is turned off. + private def verifyAdaptiveRead( + query: String, + scanned: Int, + serializer: String, + sparkOperator: Boolean): Unit = { + val df = spark.sql(query) + df.collect() + val executed = df.queryExecution.executedPlan + val plan = executed.toString() + assert(executed.isInstanceOf[AdaptiveSparkPlanExec], s"Expected an adaptive plan:\n$plan") + + val nativeScans = collect(executed) { case s: CometInMemoryTableScanExec => s } + val sparkScans = collect(executed) { case s: InMemoryTableScanExec => s } + assert( + nativeScans.length + sparkScans.length == 1, + s"Expected exactly one cache scan:\n$plan") + val cometFormat = serializer == classOf[ArrowCachedBatchSerializer].getName + assert( + nativeScans.nonEmpty == cometFormat, + s"Expected a native scan exactly for Comet's format:\n$plan") + val (relation, columns) = nativeScans.headOption + .map(s => (s.originalPlan.relation, s.scanOutput.length)) + .getOrElse((sparkScans.head.relation, sparkScans.head.attributes.length)) + assert(columns == scanned, s"Expected the scan to read $scanned columns:\n$plan") + val actual = relation.cacheBuilder.serializer.getClass.getName + assert(actual == serializer, s"Expected a relation cached by $serializer, not $actual") + + val sparkAggregates = collect(executed) { case a: HashAggregateExec => a } + assert( + sparkAggregates.nonEmpty == (sparkOperator || !cometFormat), + s"Expected Spark aggregates only where this case puts them:\n$plan") + } + // spark.sql.cache.serializer is static, and InMemoryRelation memoizes the serializer it names // for the life of the JVM. It looks the name up in the active session's conf when a relation is // cached, though, so setting it there directly and clearing the memoized instance around one @@ -540,8 +659,8 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { // scan and convert" -- which is the overhead this feature exists to remove. // // Neither case is a baseline for Spark's own cache format: both read the same Comet-written - // CometCachedBatch. Spark's format is only measured by runSparkOperatorBenchmark, with Comet - // off, since that is the only comparison it answers. + // CometCachedBatch. Spark's format is measured by runSparkOperatorBenchmark, with Comet off, + // and by runAdaptiveBenchmark, with Comet on. withSQLConf(cacheConf(nativeCacheEnabled = true): _*) { spark .sql(s"SELECT ${relation.columns.mkString(", ")} FROM ${relation.source}") @@ -611,4 +730,14 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { CometConf.COMET_ENABLED.key -> "false", CometConf.COMET_EXEC_ENABLED.key -> "false", "spark.sql.inMemoryColumnarStorage.batchSize" -> "10000") + + // Comet and AQE on, and Comet's other settings at their defaults, unlike cacheConf. The batch + // size matches the other confs, and on-heap mode is what lets Comet run in this session. + private val adaptiveConf: Seq[(String, String)] = Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + "spark.comet.exec.onHeap.enabled" -> "true", + "spark.sql.inMemoryColumnarStorage.batchSize" -> "10000") }