diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff index cb0c8aaffa..681c497b1a 100644 --- a/dev/diffs/3.4.3.diff +++ b/dev/diffs/3.4.3.diff @@ -279,10 +279,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._ @@ -291,7 +295,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( @@ -301,6 +326,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..e25dc3d1993 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -367,6 +402,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 @@ -490,6 +539,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 @@ -513,10 +622,10 @@ index daef11ae4d6..9f3cc9181f2 100644 assert(exchanges.size == 2) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -index f33432ddb6f..b375e285dde 100644 +index f33432ddb6f..4e8522d3aab 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -@@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen +@@ -22,9 +22,11 @@ import org.scalatest.GivenWhenThen import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression} import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._ import org.apache.spark.sql.catalyst.plans.ExistenceJoin @@ -524,7 +633,11 @@ index f33432ddb6f..b375e285dde 100644 import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive._ -@@ -193,6 +194,7 @@ abstract class DynamicPartitionPruningSuiteBase ++import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec + import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec} + import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec +@@ -193,6 +195,7 @@ abstract class DynamicPartitionPruningSuiteBase } val subqueryBroadcast = dpExprs.collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b @@ -532,7 +645,19 @@ index f33432ddb6f..b375e285dde 100644 } val hasFilter = if (withSubquery) "Should" else "Shouldn't" -@@ -247,6 +249,8 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -229,7 +232,10 @@ abstract class DynamicPartitionPruningSuiteBase + } + + val isMainQueryAdaptive = plan.isInstanceOf[AdaptiveSparkPlanExec] +- subqueriesAll(plan).filterNot(subqueryBroadcast.contains).foreach { s => ++ // Comet's cache scan exposes Spark's own scan as its one subquery, for Spark's UI to draw ++ // the cached plan below it, but the query never runs it as a subquery. ++ subqueriesAll(plan).filterNot(subqueryBroadcast.contains) ++ .filterNot(_.isInstanceOf[InMemoryTableScanExec]).foreach { s => + val subquery = s match { + case r: ReusedSubqueryExec => r.child + case o => o +@@ -247,6 +253,8 @@ abstract class DynamicPartitionPruningSuiteBase val buf = collectDynamicPruningExpressions(df.queryExecution.executedPlan).collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b.index @@ -541,7 +666,7 @@ index f33432ddb6f..b375e285dde 100644 } assert(buf.distinct.size == n) } -@@ -262,6 +266,12 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -262,6 +270,12 @@ abstract class DynamicPartitionPruningSuiteBase case s: BatchScanExec => s.runtimeFilters.collect { case d: DynamicPruningExpression => d.child } @@ -554,7 +679,7 @@ index f33432ddb6f..b375e285dde 100644 case _ => Nil } } -@@ -1204,10 +1214,16 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1204,10 +1218,16 @@ abstract class DynamicPartitionPruningSuiteBase val plan = df.queryExecution.executedPlan val countSubqueryBroadcasts = @@ -573,7 +698,7 @@ index f33432ddb6f..b375e285dde 100644 assert(countSubqueryBroadcasts == 1) assert(countReusedSubqueryBroadcasts == 1) -@@ -1577,6 +1593,7 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1577,6 +1597,7 @@ abstract class DynamicPartitionPruningSuiteBase val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) { case s: SubqueryBroadcastExec => s @@ -581,7 +706,7 @@ index f33432ddb6f..b375e285dde 100644 } assert(subqueryBroadcastExecs.size === 1) subqueryBroadcastExecs.foreach { subqueryBroadcastExec => -@@ -1729,6 +1746,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat +@@ -1729,6 +1750,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat case s: BatchScanExec => // we use f1 col for v2 tables due to schema pruning s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1"))) @@ -1038,10 +1163,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"))) } @@ -1051,7 +1184,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 }) @@ -1068,7 +1201,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 } } @@ -1109,7 +1250,7 @@ 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..a164e273b76 100644 +index 48ad10992c5..cfd4b5ccb26 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 @@ -165,7 +165,15 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper { @@ -1128,7 +1269,15 @@ index 48ad10992c5..a164e273b76 100644 val extensions = create { extensions => // inject rule that will run during AQE query stage preparation and will add custom tags // to the plan -@@ -259,6 +267,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper { +@@ -208,6 +216,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) + } +@@ -259,6 +268,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper { } private def testInjectColumnar(enableAQE: Boolean): Unit = { @@ -1136,7 +1285,7 @@ index 48ad10992c5..a164e273b76 100644 def collectPlanSteps(plan: SparkPlan): Seq[Int] = plan match { case a: AdaptiveSparkPlanExec => assert(a.toString.startsWith("AdaptiveSparkPlan isFinalPlan=true")) -@@ -314,6 +323,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper { +@@ -314,6 +324,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper { } test("reset column vectors") { @@ -1144,7 +1293,7 @@ index 48ad10992c5..a164e273b76 100644 val session = SparkSession.builder() .master("local[1]") .config(COLUMN_BATCH_SIZE.key, 2) -@@ -482,6 +492,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper { +@@ -482,6 +493,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper { } test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE Optimizer") { @@ -1152,6 +1301,20 @@ index 48ad10992c5..a164e273b76 100644 def executedPlan(df: Dataset[java.lang.Long]): SparkPlan = { assert(df.queryExecution.executedPlan.isInstanceOf[AdaptiveSparkPlanExec]) df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan +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 @@ -2138,6 +2301,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 @@ -2608,10 +2916,10 @@ index 3a0bd35cb70..b28f06a757f 100644 val workDirPath = workDir.getAbsolutePath val input = spark.range(5).toDF("id") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -index 6333808b420..81b2704300c 100644 +index 6333808b420..5d99f1f5b22 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -@@ -21,7 +21,7 @@ import scala.reflect.ClassTag +@@ -21,12 +21,13 @@ import scala.reflect.ClassTag import org.apache.spark.AccumulatorSuite import org.apache.spark.internal.config.EXECUTOR_MEMORY @@ -2620,7 +2928,28 @@ index 6333808b420..81b2704300c 100644 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, Cast, Expression, Literal, ShiftLeft} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.Inner -@@ -486,7 +486,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + import org.apache.spark.sql.catalyst.plans.logical.BROADCAST + import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, PartitioningCollection} ++import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{DummySparkPlan, SparkPlan, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +@@ -156,12 +157,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) +@@ -486,7 +489,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils } } @@ -2630,6 +2959,21 @@ index 6333808b420..81b2704300c 100644 withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { val t1 = (0 until 100).map(i => (i % 5, i % 13)).toDF("i1", "j1") val t2 = (0 until 100).map(i => (i % 5, i % 14)).toDF("i2", "j2") +@@ -654,12 +658,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala index 26e61c6b58d..cb09d7e116a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala @@ -3250,10 +3594,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) @@ -3269,6 +3613,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, @@ -3354,10 +3703,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", @@ -3409,6 +3758,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 53bbd92c87..d5ad4abce9 100644 --- a/dev/diffs/3.5.9.diff +++ b/dev/diffs/3.5.9.diff @@ -260,10 +260,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._ @@ -272,7 +276,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( @@ -282,6 +307,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..52523429c87 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -348,6 +383,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 @@ -471,6 +520,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 @@ -494,10 +620,10 @@ index c4fb4fa943c..a04b23870a8 100644 assert(exchanges.size == 2) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -index f33432ddb6f..b375e285dde 100644 +index f33432ddb6f..4e8522d3aab 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -@@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen +@@ -22,9 +22,11 @@ import org.scalatest.GivenWhenThen import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression} import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._ import org.apache.spark.sql.catalyst.plans.ExistenceJoin @@ -505,7 +631,11 @@ index f33432ddb6f..b375e285dde 100644 import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive._ -@@ -193,6 +194,7 @@ abstract class DynamicPartitionPruningSuiteBase ++import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec + import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec} + import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec +@@ -193,6 +195,7 @@ abstract class DynamicPartitionPruningSuiteBase } val subqueryBroadcast = dpExprs.collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b @@ -513,7 +643,19 @@ index f33432ddb6f..b375e285dde 100644 } val hasFilter = if (withSubquery) "Should" else "Shouldn't" -@@ -247,6 +249,8 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -229,7 +232,10 @@ abstract class DynamicPartitionPruningSuiteBase + } + + val isMainQueryAdaptive = plan.isInstanceOf[AdaptiveSparkPlanExec] +- subqueriesAll(plan).filterNot(subqueryBroadcast.contains).foreach { s => ++ // Comet's cache scan exposes Spark's own scan as its one subquery, for Spark's UI to draw ++ // the cached plan below it, but the query never runs it as a subquery. ++ subqueriesAll(plan).filterNot(subqueryBroadcast.contains) ++ .filterNot(_.isInstanceOf[InMemoryTableScanExec]).foreach { s => + val subquery = s match { + case r: ReusedSubqueryExec => r.child + case o => o +@@ -247,6 +253,8 @@ abstract class DynamicPartitionPruningSuiteBase val buf = collectDynamicPruningExpressions(df.queryExecution.executedPlan).collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b.index @@ -522,7 +664,7 @@ index f33432ddb6f..b375e285dde 100644 } assert(buf.distinct.size == n) } -@@ -262,6 +266,12 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -262,6 +270,12 @@ abstract class DynamicPartitionPruningSuiteBase case s: BatchScanExec => s.runtimeFilters.collect { case d: DynamicPruningExpression => d.child } @@ -535,7 +677,7 @@ index f33432ddb6f..b375e285dde 100644 case _ => Nil } } -@@ -1204,10 +1214,16 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1204,10 +1218,16 @@ abstract class DynamicPartitionPruningSuiteBase val plan = df.queryExecution.executedPlan val countSubqueryBroadcasts = @@ -554,7 +696,7 @@ index f33432ddb6f..b375e285dde 100644 assert(countSubqueryBroadcasts == 1) assert(countReusedSubqueryBroadcasts == 1) -@@ -1577,6 +1593,7 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1577,6 +1597,7 @@ abstract class DynamicPartitionPruningSuiteBase val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) { case s: SubqueryBroadcastExec => s @@ -562,7 +704,7 @@ index f33432ddb6f..b375e285dde 100644 } assert(subqueryBroadcastExecs.size === 1) subqueryBroadcastExecs.foreach { subqueryBroadcastExec => -@@ -1729,6 +1746,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat +@@ -1729,6 +1750,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat case s: BatchScanExec => // we use f1 col for v2 tables due to schema pruning s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1"))) @@ -1045,10 +1187,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"))) } @@ -1058,7 +1208,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 }) @@ -1075,7 +1225,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 } } @@ -1117,7 +1275,7 @@ 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..5218e9125c6 100644 +index 8b4ac474f87..b097a14a972 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 @@ -167,7 +167,15 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt @@ -1136,7 +1294,15 @@ index 8b4ac474f87..5218e9125c6 100644 val extensions = create { extensions => // inject rule that will run during AQE query stage preparation and will add custom tags // to the plan -@@ -261,6 +269,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -210,6 +218,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) + } +@@ -261,6 +270,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } private def testInjectColumnar(enableAQE: Boolean): Unit = { @@ -1144,7 +1310,7 @@ index 8b4ac474f87..5218e9125c6 100644 def collectPlanSteps(plan: SparkPlan): Seq[Int] = plan match { case a: AdaptiveSparkPlanExec => assert(a.toString.startsWith("AdaptiveSparkPlan isFinalPlan=true")) -@@ -316,6 +325,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -316,6 +326,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("reset column vectors") { @@ -1152,7 +1318,7 @@ index 8b4ac474f87..5218e9125c6 100644 val session = SparkSession.builder() .master("local[1]") .config(COLUMN_BATCH_SIZE.key, 2) -@@ -484,6 +494,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -484,6 +495,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE Optimizer") { @@ -1160,6 +1326,20 @@ index 8b4ac474f87..5218e9125c6 100644 def executedPlan(df: Dataset[java.lang.Long]): SparkPlan = { assert(df.queryExecution.executedPlan.isInstanceOf[AdaptiveSparkPlanExec]) df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan +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 @@ -1735,7 +1915,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..a9567edf805 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._ @@ -2141,7 +2321,17 @@ index 2f8e401e743..7849c685b19 100644 }.size == (if (firstAccess) 2 else 0)) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -2938,7 +2981,8 @@ class AdaptiveQueryExecSuite +@@ -2923,7 +2966,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 +2982,8 @@ class AdaptiveQueryExecSuite } } @@ -2151,7 +2341,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 +3025,9 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -2162,6 +2352,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 @@ -2604,10 +2939,10 @@ index b8f3ea3c6f3..bbd44221288 100644 val workDirPath = workDir.getAbsolutePath val input = spark.range(5).toDF("id") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -index 5479be86e9f..07c81e4b830 100644 +index 5479be86e9f..103dab4657c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -@@ -21,7 +21,7 @@ import scala.reflect.ClassTag +@@ -21,12 +21,13 @@ import scala.reflect.ClassTag import org.apache.spark.AccumulatorSuite import org.apache.spark.internal.config.EXECUTOR_MEMORY @@ -2616,7 +2951,28 @@ index 5479be86e9f..07c81e4b830 100644 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, Cast, Expression, Literal, ShiftLeft} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.Inner -@@ -487,7 +487,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + import org.apache.spark.sql.catalyst.plans.logical.BROADCAST + import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, PartitioningCollection} ++import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{DummySparkPlan, SparkPlan, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +@@ -157,12 +158,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) +@@ -487,7 +490,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils } } @@ -2626,6 +2982,21 @@ index 5479be86e9f..07c81e4b830 100644 withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { val t1 = (0 until 100).map(i => (i % 5, i % 13)).toDF("i1", "j1") val t2 = (0 until 100).map(i => (i % 5, i % 14)).toDF("i2", "j2") +@@ -655,12 +659,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala index 5cdbdc27b32..307fba16578 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala @@ -3261,10 +3632,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) @@ -3280,6 +3651,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, @@ -3365,10 +3741,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", @@ -3420,6 +3796,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 e5c5793fd1..1ff2ef453d 100644 --- a/dev/diffs/4.0.4.diff +++ b/dev/diffs/4.0.4.diff @@ -375,10 +375,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._ @@ -387,7 +391,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( @@ -397,7 +422,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")) @@ -419,6 +464,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..e59a931c06a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -485,6 +538,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 @@ -608,6 +675,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 @@ -631,10 +789,10 @@ index 81713c777bc..b5f92ed9742 100644 assert(exchanges.size == 2) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -index 2c24cc7d570..12d897866da 100644 +index 2c24cc7d570..505157acd91 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -@@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen +@@ -22,9 +22,11 @@ import org.scalatest.GivenWhenThen import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression} import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._ import org.apache.spark.sql.catalyst.plans.ExistenceJoin @@ -642,7 +800,11 @@ index 2c24cc7d570..12d897866da 100644 import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive._ -@@ -193,6 +194,7 @@ abstract class DynamicPartitionPruningSuiteBase ++import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec + import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec} + import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec +@@ -193,6 +195,7 @@ abstract class DynamicPartitionPruningSuiteBase } val subqueryBroadcast = dpExprs.collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b @@ -650,7 +812,19 @@ index 2c24cc7d570..12d897866da 100644 } val hasFilter = if (withSubquery) "Should" else "Shouldn't" -@@ -247,6 +249,8 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -229,7 +232,10 @@ abstract class DynamicPartitionPruningSuiteBase + } + + val isMainQueryAdaptive = plan.isInstanceOf[AdaptiveSparkPlanExec] +- subqueriesAll(plan).filterNot(subqueryBroadcast.contains).foreach { s => ++ // Comet's cache scan exposes Spark's own scan as its one subquery, for Spark's UI to draw ++ // the cached plan below it, but the query never runs it as a subquery. ++ subqueriesAll(plan).filterNot(subqueryBroadcast.contains) ++ .filterNot(_.isInstanceOf[InMemoryTableScanExec]).foreach { s => + val subquery = s match { + case r: ReusedSubqueryExec => r.child + case o => o +@@ -247,6 +253,8 @@ abstract class DynamicPartitionPruningSuiteBase val buf = collectDynamicPruningExpressions(df.queryExecution.executedPlan).collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b.indices.map(idx => b.buildKeys(idx)) @@ -659,7 +833,7 @@ index 2c24cc7d570..12d897866da 100644 } assert(buf.distinct.size == n) } -@@ -262,6 +266,12 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -262,6 +270,12 @@ abstract class DynamicPartitionPruningSuiteBase case s: BatchScanExec => s.runtimeFilters.collect { case d: DynamicPruningExpression => d.child } @@ -672,7 +846,7 @@ index 2c24cc7d570..12d897866da 100644 case _ => Nil } } -@@ -1204,10 +1214,16 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1204,10 +1218,16 @@ abstract class DynamicPartitionPruningSuiteBase val plan = df.queryExecution.executedPlan val countSubqueryBroadcasts = @@ -691,7 +865,7 @@ index 2c24cc7d570..12d897866da 100644 assert(countSubqueryBroadcasts == 1) assert(countReusedSubqueryBroadcasts == 1) -@@ -1578,6 +1594,7 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1578,6 +1598,7 @@ abstract class DynamicPartitionPruningSuiteBase val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) { case s: SubqueryBroadcastExec => s @@ -699,7 +873,7 @@ index 2c24cc7d570..12d897866da 100644 } assert(subqueryBroadcastExecs.size === 1) subqueryBroadcastExecs.foreach { subqueryBroadcastExec => -@@ -1730,6 +1747,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat +@@ -1730,6 +1751,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat case s: BatchScanExec => // we use f1 col for v2 tables due to schema pruning s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1"))) @@ -1192,10 +1366,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"))) } @@ -1205,7 +1387,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 }) @@ -1222,6 +1404,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 @@ -1254,7 +1444,7 @@ 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..d068241d585 100644 +index c1c041509c3..6fc92c40ae7 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 @@ -179,7 +179,15 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt @@ -1273,7 +1463,15 @@ index c1c041509c3..d068241d585 100644 val extensions = create { extensions => // inject rule that will run during AQE query stage preparation and will add custom tags // to the plan -@@ -273,6 +281,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -222,6 +230,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) + } +@@ -273,6 +282,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } private def testInjectColumnar(enableAQE: Boolean): Unit = { @@ -1281,7 +1479,7 @@ index c1c041509c3..d068241d585 100644 def collectPlanSteps(plan: SparkPlan): Seq[Int] = plan match { case a: AdaptiveSparkPlanExec => assert(a.toString.startsWith("AdaptiveSparkPlan isFinalPlan=true")) -@@ -328,6 +337,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -328,6 +338,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("reset column vectors") { @@ -1289,7 +1487,7 @@ index c1c041509c3..d068241d585 100644 val session = SparkSession.builder() .master("local[1]") .config(COLUMN_BATCH_SIZE.key, 2) -@@ -496,6 +506,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -496,6 +507,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE Optimizer") { @@ -1312,6 +1510,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 @@ -1511,7 +1723,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} @@ -1519,7 +1731,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 @@ -1578,6 +1790,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 @@ -2382,7 +2603,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..e358c2e7b59 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 @@ -2867,7 +3088,17 @@ index fb8fab6a80f..403eb411920 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3048,7 +3101,8 @@ class AdaptiveQueryExecSuite +@@ -3033,7 +3086,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 +3102,8 @@ class AdaptiveQueryExecSuite } } @@ -2877,7 +3108,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 +3193,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -2887,6 +3118,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 @@ -3329,10 +3706,10 @@ index b8f3ea3c6f3..bbd44221288 100644 val workDirPath = workDir.getAbsolutePath val input = spark.range(5).toDF("id") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -index 69dd04e07d5..781018ecc66 100644 +index 69dd04e07d5..a78dd593985 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -@@ -21,7 +21,7 @@ import scala.reflect.ClassTag +@@ -21,13 +21,14 @@ import scala.reflect.ClassTag import org.apache.spark.AccumulatorSuite import org.apache.spark.internal.config.EXECUTOR_MEMORY @@ -3341,7 +3718,29 @@ index 69dd04e07d5..781018ecc66 100644 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, Cast, Expression, Literal, ShiftLeft} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.Inner -@@ -488,7 +488,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + import org.apache.spark.sql.catalyst.plans.logical.BROADCAST + import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, PartitioningCollection} + import org.apache.spark.sql.classic.{Dataset, SparkSession} ++import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{DummySparkPlan, SparkPlan, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +@@ -158,12 +159,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) +@@ -488,7 +491,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils } } @@ -3351,6 +3750,21 @@ index 69dd04e07d5..781018ecc66 100644 withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { val t1 = (0 until 100).map(i => (i % 5, i % 13)).toDF("i1", "j1") val t2 = (0 until 100).map(i => (i % 5, i % 14)).toDF("i2", "j2") +@@ -656,12 +660,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala index 0dd90925d3c..7d53ec845ef 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala @@ -4053,10 +4467,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) @@ -4073,6 +4487,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, @@ -4195,10 +4614,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", @@ -4243,6 +4662,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 57a6c69607..b100542451 100644 --- a/dev/diffs/4.1.3.diff +++ b/dev/diffs/4.1.3.diff @@ -400,11 +400,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 @@ -413,7 +443,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( @@ -423,7 +474,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")) @@ -445,6 +516,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..df8fb4c0a35 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -512,7 +607,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 @@ -528,7 +623,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 @@ -542,7 +647,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 @@ -556,7 +661,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 { @@ -569,7 +674,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 @@ -710,6 +815,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 @@ -733,10 +931,10 @@ index 6df8d66ee7f..35e270c7241 100644 assert(exchanges.size == 2) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -index e1a2fd33c7c..632f4b695df 100644 +index e1a2fd33c7c..ac0682e0e1a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -@@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen +@@ -22,9 +22,11 @@ import org.scalatest.GivenWhenThen import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression} import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._ import org.apache.spark.sql.catalyst.plans.ExistenceJoin @@ -744,7 +942,11 @@ index e1a2fd33c7c..632f4b695df 100644 import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive._ -@@ -193,6 +194,7 @@ abstract class DynamicPartitionPruningSuiteBase ++import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec + import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec} + import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec +@@ -193,6 +195,7 @@ abstract class DynamicPartitionPruningSuiteBase } val subqueryBroadcast = dpExprs.collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b @@ -752,7 +954,19 @@ index e1a2fd33c7c..632f4b695df 100644 } val hasFilter = if (withSubquery) "Should" else "Shouldn't" -@@ -247,6 +249,8 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -229,7 +232,10 @@ abstract class DynamicPartitionPruningSuiteBase + } + + val isMainQueryAdaptive = plan.isInstanceOf[AdaptiveSparkPlanExec] +- subqueriesAll(plan).filterNot(subqueryBroadcast.contains).foreach { s => ++ // Comet's cache scan exposes Spark's own scan as its one subquery, for Spark's UI to draw ++ // the cached plan below it, but the query never runs it as a subquery. ++ subqueriesAll(plan).filterNot(subqueryBroadcast.contains) ++ .filterNot(_.isInstanceOf[InMemoryTableScanExec]).foreach { s => + val subquery = s match { + case r: ReusedSubqueryExec => r.child + case o => o +@@ -247,6 +253,8 @@ abstract class DynamicPartitionPruningSuiteBase val buf = collectDynamicPruningExpressions(df.queryExecution.executedPlan).collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b.indices.map(idx => b.buildKeys(idx)) @@ -761,7 +975,7 @@ index e1a2fd33c7c..632f4b695df 100644 } assert(buf.distinct.size == n) } -@@ -262,6 +266,12 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -262,6 +270,12 @@ abstract class DynamicPartitionPruningSuiteBase case s: BatchScanExec => s.runtimeFilters.collect { case d: DynamicPruningExpression => d.child } @@ -774,7 +988,7 @@ index e1a2fd33c7c..632f4b695df 100644 case _ => Nil } } -@@ -1204,10 +1214,16 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1204,10 +1218,16 @@ abstract class DynamicPartitionPruningSuiteBase val plan = df.queryExecution.executedPlan val countSubqueryBroadcasts = @@ -793,7 +1007,7 @@ index e1a2fd33c7c..632f4b695df 100644 assert(countSubqueryBroadcasts == 1) assert(countReusedSubqueryBroadcasts == 1) -@@ -1579,6 +1595,7 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1579,6 +1599,7 @@ abstract class DynamicPartitionPruningSuiteBase val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) { case s: SubqueryBroadcastExec => s @@ -801,7 +1015,7 @@ index e1a2fd33c7c..632f4b695df 100644 } assert(subqueryBroadcastExecs.size === 1) subqueryBroadcastExecs.foreach { subqueryBroadcastExec => -@@ -1731,6 +1748,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat +@@ -1731,6 +1752,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat case s: BatchScanExec => // we use f1 col for v2 tables due to schema pruning s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1"))) @@ -1313,10 +1527,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"))) } @@ -1326,7 +1548,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 }) @@ -1343,6 +1565,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 @@ -1388,7 +1618,7 @@ 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..a330d142eb1 100644 +index 66826a9ca76..db56c426dc0 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 @@ -196,7 +196,15 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt @@ -1407,7 +1637,15 @@ index 66826a9ca76..a330d142eb1 100644 val extensions = create { extensions => // inject rule that will run during AQE query stage preparation and will add custom tags // to the plan -@@ -290,6 +298,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -239,6 +247,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) + } +@@ -290,6 +299,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } private def testInjectColumnar(enableAQE: Boolean): Unit = { @@ -1415,7 +1653,7 @@ index 66826a9ca76..a330d142eb1 100644 def collectPlanSteps(plan: SparkPlan): Seq[Int] = plan match { case a: AdaptiveSparkPlanExec => assert(a.toString.startsWith("AdaptiveSparkPlan isFinalPlan=true")) -@@ -345,6 +354,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -345,6 +355,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("reset column vectors") { @@ -1423,7 +1661,7 @@ index 66826a9ca76..a330d142eb1 100644 val session = SparkSession.builder() .master("local[1]") .config(COLUMN_BATCH_SIZE.key, 2) -@@ -513,6 +523,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -513,6 +524,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE Optimizer") { @@ -1446,6 +1684,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 @@ -1645,7 +1897,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} @@ -1653,7 +1905,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 @@ -1712,6 +1964,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 @@ -2516,7 +2777,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..6c2d3216604 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 @@ -3001,7 +3262,17 @@ index 188a28ff1c0..8fdccf31749 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3229,7 +3282,8 @@ class AdaptiveQueryExecSuite +@@ -3214,7 +3267,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 +3283,8 @@ class AdaptiveQueryExecSuite } } @@ -3011,7 +3282,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 +3374,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -3022,10 +3293,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 => @@ -3050,6 +3340,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 @@ -3537,10 +3973,10 @@ index b8f3ea3c6f3..bbd44221288 100644 val workDirPath = workDir.getAbsolutePath val input = spark.range(5).toDF("id") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -index 9bd858608cb..2682ba53513 100644 +index 9bd858608cb..896d33a8f0e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -@@ -21,7 +21,7 @@ import scala.reflect.ClassTag +@@ -21,13 +21,14 @@ import scala.reflect.ClassTag import org.apache.spark.AccumulatorSuite import org.apache.spark.internal.config.EXECUTOR_MEMORY @@ -3549,7 +3985,29 @@ index 9bd858608cb..2682ba53513 100644 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, Cast, Expression, Literal, ShiftLeft} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.Inner -@@ -488,7 +488,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + import org.apache.spark.sql.catalyst.plans.logical.BROADCAST + import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, PartitioningCollection} + import org.apache.spark.sql.classic.{Dataset, SparkSession} ++import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{DummySparkPlan, SparkPlan, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +@@ -158,12 +159,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) +@@ -488,7 +491,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils } } @@ -3559,6 +4017,21 @@ index 9bd858608cb..2682ba53513 100644 withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { val t1 = (0 until 100).map(i => (i % 5, i % 13)).toDF("i1", "j1") val t2 = (0 until 100).map(i => (i % 5, i % 14)).toDF("i2", "j2") +@@ -656,12 +660,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala index f2e9121d566..2c9f517034f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala @@ -4366,10 +4839,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) @@ -4386,6 +4859,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, @@ -4508,10 +4986,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", @@ -4556,6 +5034,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 3fed2699cf..06f5e72544 100644 --- a/dev/diffs/4.2.0.diff +++ b/dev/diffs/4.2.0.diff @@ -418,11 +418,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 @@ -431,7 +461,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( @@ -441,7 +492,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")) @@ -463,6 +534,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..eb3de47f387 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -541,7 +636,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 @@ -557,7 +652,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 @@ -571,7 +676,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 @@ -585,7 +690,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 { @@ -598,7 +703,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 @@ -739,6 +844,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 @@ -762,10 +960,10 @@ index 879569045b6..f3ff89067d2 100644 assert(exchanges.size == 2) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -index c68f64d52f2..32489f2cb15 100644 +index c68f64d52f2..59f9f9bb092 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -@@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen +@@ -22,9 +22,11 @@ import org.scalatest.GivenWhenThen import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression} import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._ import org.apache.spark.sql.catalyst.plans.ExistenceJoin @@ -773,7 +971,11 @@ index c68f64d52f2..32489f2cb15 100644 import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive._ -@@ -192,6 +193,7 @@ abstract class DynamicPartitionPruningSuiteBase ++import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec + import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec} + import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec +@@ -192,6 +194,7 @@ abstract class DynamicPartitionPruningSuiteBase } val subqueryBroadcast = dpExprs.collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b @@ -781,7 +983,19 @@ index c68f64d52f2..32489f2cb15 100644 } val hasFilter = if (withSubquery) "Should" else "Shouldn't" -@@ -246,6 +248,8 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -228,7 +231,10 @@ abstract class DynamicPartitionPruningSuiteBase + } + + val isMainQueryAdaptive = plan.isInstanceOf[AdaptiveSparkPlanExec] +- subqueriesAll(plan).filterNot(subqueryBroadcast.contains).foreach { s => ++ // Comet's cache scan exposes Spark's own scan as its one subquery, for Spark's UI to draw ++ // the cached plan below it, but the query never runs it as a subquery. ++ subqueriesAll(plan).filterNot(subqueryBroadcast.contains) ++ .filterNot(_.isInstanceOf[InMemoryTableScanExec]).foreach { s => + val subquery = s match { + case r: ReusedSubqueryExec => r.child + case o => o +@@ -246,6 +252,8 @@ abstract class DynamicPartitionPruningSuiteBase val buf = collectDynamicPruningExpressions(df.queryExecution.executedPlan).collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b.indices.map(idx => b.buildKeys(idx)) @@ -790,7 +1004,7 @@ index c68f64d52f2..32489f2cb15 100644 } assert(buf.distinct.size == n) } -@@ -261,6 +265,12 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -261,6 +269,12 @@ abstract class DynamicPartitionPruningSuiteBase case s: BatchScanExec => s.runtimeFilters.collect { case d: DynamicPruningExpression => d.child } @@ -803,7 +1017,7 @@ index c68f64d52f2..32489f2cb15 100644 case _ => Nil } } -@@ -1203,10 +1213,16 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1203,10 +1217,16 @@ abstract class DynamicPartitionPruningSuiteBase val plan = df.queryExecution.executedPlan val countSubqueryBroadcasts = @@ -822,7 +1036,7 @@ index c68f64d52f2..32489f2cb15 100644 assert(countSubqueryBroadcasts == 1) assert(countReusedSubqueryBroadcasts == 1) -@@ -1578,6 +1594,7 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1578,6 +1598,7 @@ abstract class DynamicPartitionPruningSuiteBase val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) { case s: SubqueryBroadcastExec => s @@ -830,7 +1044,7 @@ index c68f64d52f2..32489f2cb15 100644 } assert(subqueryBroadcastExecs.size === 1) subqueryBroadcastExecs.foreach { subqueryBroadcastExec => -@@ -1730,6 +1747,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat +@@ -1730,6 +1751,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat case s: BatchScanExec => // we use f1 col for v2 tables due to schema pruning s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1"))) @@ -1373,10 +1587,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"))) } @@ -1386,7 +1608,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 }) @@ -1403,6 +1625,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 @@ -1469,7 +1699,7 @@ 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..6c177bff0e3 100644 +index bfcf583a705..749e6d7f1a4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -217,7 +217,15 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp @@ -1488,7 +1718,15 @@ index bfcf583a705..6c177bff0e3 100644 val extensions = create { extensions => // inject rule that will run during AQE query stage preparation and will add custom tags // to the plan -@@ -311,6 +319,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp +@@ -260,6 +268,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) + } +@@ -311,6 +320,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp } private def testInjectColumnar(enableAQE: Boolean): Unit = { @@ -1496,7 +1734,7 @@ index bfcf583a705..6c177bff0e3 100644 def collectPlanSteps(plan: SparkPlan): Seq[Int] = plan match { case a: AdaptiveSparkPlanExec => assert(a.toString.startsWith("AdaptiveSparkPlan isFinalPlan=true")) -@@ -366,6 +375,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp +@@ -366,6 +376,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp } test("reset column vectors") { @@ -1504,7 +1742,7 @@ index bfcf583a705..6c177bff0e3 100644 val session = SparkSession.builder() .master("local[1]") .config(COLUMN_BATCH_SIZE.key, 2) -@@ -553,6 +563,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp +@@ -553,6 +564,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp } test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE Optimizer") { @@ -1527,6 +1765,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 @@ -1726,7 +1978,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} @@ -1734,7 +1986,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 @@ -1793,6 +2045,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 @@ -2607,7 +2868,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..f1fc6dee563 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 @@ -3083,7 +3344,17 @@ index d6d19d21e65..751ad50a569 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3278,7 +3330,8 @@ class AdaptiveQueryExecSuite +@@ -3263,7 +3315,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 +3331,8 @@ class AdaptiveQueryExecSuite } } @@ -3093,7 +3364,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 +3422,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -3104,10 +3375,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 => @@ -3132,6 +3422,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 @@ -3618,10 +4054,10 @@ index b8f3ea3c6f3..bbd44221288 100644 val workDirPath = workDir.getAbsolutePath val input = spark.range(5).toDF("id") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -index 52746720eba..15df5e78440 100644 +index 52746720eba..98266bf06f5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -@@ -21,7 +21,7 @@ import scala.reflect.ClassTag +@@ -21,13 +21,14 @@ import scala.reflect.ClassTag import org.apache.spark.AccumulatorSuite import org.apache.spark.internal.config.EXECUTOR_MEMORY @@ -3630,7 +4066,29 @@ index 52746720eba..15df5e78440 100644 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, Cast, Expression, Literal, ShiftLeft} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.Inner -@@ -487,7 +487,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest + import org.apache.spark.sql.catalyst.plans.logical.BROADCAST + import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, PartitioningCollection} + import org.apache.spark.sql.classic.{Dataset, SparkSession} ++import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{DummySparkPlan, SparkPlan, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +@@ -157,12 +158,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) +@@ -487,7 +490,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest } } @@ -3640,6 +4098,21 @@ index 52746720eba..15df5e78440 100644 withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { val t1 = (0 until 100).map(i => (i % 5, i % 13)).toDF("i1", "j1") val t2 = (0 until 100).map(i => (i % 5, i % 14)).toDF("i2", "j2") +@@ -654,12 +658,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala index 2ae2ea6339f..30c8959a057 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala @@ -4513,10 +4986,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) @@ -4533,6 +5006,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, @@ -4611,10 +5089,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", @@ -4659,6 +5137,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/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 467e48ecb5..67036eccbd 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 20091ac317..4266c09f22 100644 --- a/docs/source/user-guide/latest/in-memory-cache.md +++ b/docs/source/user-guide/latest/in-memory-cache.md @@ -21,16 +21,17 @@ 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.convert.inMemoryCache.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 @@ -112,7 +113,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. | @@ -161,10 +162,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.convert.inMemoryCache.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.convert.inMemoryCache.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 @@ -220,6 +254,9 @@ operator reads a relation cached in Comet's format depends on the scan below it: and operators that do not take part in code generation, such as exchanges and limits, or a query that returns the cached rows as they are. +Whenever Spark's scan reads Comet's format, which also happens for a relation whose cached plan +records `Dataset.observe` metrics, Comet records a fallback reason on the scan. + Measured by the same benchmark over the same 5M-row relation, with native execution off so that Spark operators consume the cached data, Comet disabled for the row reader and enabled for the fused reader (Apple M4, JDK 17, Spark 4.1; the average of two runs): @@ -234,8 +271,9 @@ reader (Apple M4, JDK 17, Spark 4.1; the average of two runs): The fused reader is faster than Spark's own format for the narrowest reads and within 20% of it for the others. The row reader takes up to 1.6 times as long, and it is the only reader for relations wider than `spark.sql.codegen.maxFields`: reading every column of relations of 100, 200 and 1500 -nullable `bigint` columns took 2.2 to 2.5 times as long as from Spark's format. These gaps are why -the feature is still off by default; +nullable `bigint` columns took 2.2 to 2.5 times as long as from Spark's format. A Spark operator +above Comet's native cache scan does not pay this; see [Performance](#performance). These gaps are +the main reason the feature is still described as experimental; [#5485](https://github.com/apache/datafusion-comet/issues/5485) tracks them. Comet's serializer exists because Spark's own Arrow cache format diff --git a/docs/source/user-guide/latest/migration-guide.md b/docs/source/user-guide/latest/migration-guide.md index 0a9731a2fa..96261f2ec8 100644 --- a/docs/source/user-guide/latest/migration-guide.md +++ b/docs/source/user-guide/latest/migration-guide.md @@ -59,6 +59,33 @@ 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 can read 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. + ### Deprecated and Removed Settings Using `spark.comet.sparkToColumnar.enabled` and `spark.comet.sparkToColumnar.supportedOperatorList` diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index aea4fc8a88..d03d89f154 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -65,7 +65,7 @@ omitted from the tables below and may be reconsidered based on demand: | `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. | | `RangeExec` | ⚠️ | Disabled by default. Set `spark.comet.exec.range.enabled=true` to generate the rows of `spark.range` and SQL `range()` in native code, so the operators above them run natively. It can be slower than Spark when those operators are only cheap expressions, such as a filter, which Spark compiles together with the range into one loop. Ranges whose arithmetic overflows the `Long` range fall back to Spark. | -| `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 f2f19ac5ce..8fd23c1c35 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -359,7 +359,7 @@ object CometConf extends ShimCometConf { "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/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index 3a698af33b..11472de837 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf._ import org.apache.comet.iceberg.IcebergWriteStrategy -import org.apache.comet.rules.{CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometRule, CometSpark34AqeDppFallbackRule} +import org.apache.comet.rules.{CometCoalesceShufflePartitions, CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometRule, CometSpark34AqeDppFallbackRule} import org.apache.comet.shims.ShimCometSparkSessionExtensions /** @@ -70,9 +70,13 @@ import org.apache.comet.shims.ShimCometSparkSessionExtensions * 1. queryStageOptimizerRules: * a. PlanAdaptiveDynamicPruningFilters (Spark) -- skips wrapped SABs * b. ReuseAdaptiveSubquery (Spark) - * c. CometPlanAdaptiveDynamicPruningFilters -- converts wrapped SABs to + * c. OptimizeSkewInRebalancePartitions, CoalesceShufflePartitions, + * OptimizeShuffleWithLocalRead (Spark) + * d. CometPlanAdaptiveDynamicPruningFilters -- converts wrapped SABs to * CometSubqueryBroadcastExec with BroadcastQueryStageExec for broadcast reuse - * d. CometReuseSubquery -- deduplicates converted subqueries + * e. CometReuseSubquery -- deduplicates converted subqueries + * f. CometCoalesceShufflePartitions -- coalesces the shuffles that Comet + * unions and broadcast joins keep Spark's CoalesceShufflePartitions from reaching * 2. postStageCreationRules -> ApplyColumnarRulesAndInsertTransitions: * a. preColumnarTransitions: CometRule (no-op, already converted) * b. insertTransitions @@ -81,9 +85,9 @@ import org.apache.comet.shims.ShimCometSparkSessionExtensions * }}} * * On Spark 3.4, injectQueryStageOptimizerRule is unavailable. CometExecRule does not wrap SABs, - * and CometPlanAdaptiveDynamicPruningFilters/CometReuseSubquery are not registered. AQE DPP scans - * fall back to Spark so that Spark's PlanAdaptiveDynamicPruningFilters handles them natively - * (with DPP). + * and CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery and + * CometCoalesceShufflePartitions are not registered. AQE DPP scans fall back to Spark so that + * Spark's PlanAdaptiveDynamicPruningFilters handles them natively (with DPP). */ class CometSparkSessionExtensions extends (SparkSessionExtensions => Unit) @@ -107,6 +111,7 @@ class CometSparkSessionExtensions } injectQueryStageOptimizerRuleShim(extensions, CometPlanAdaptiveDynamicPruningFilters) injectQueryStageOptimizerRuleShim(extensions, CometReuseSubquery) + injectQueryStageOptimizerRuleShim(extensions, CometCoalesceShufflePartitions) extensions.injectPlannerStrategy { session => IcebergWriteStrategy(session) } } diff --git a/spark/src/main/scala/org/apache/comet/rules/CometCoalesceShufflePartitions.scala b/spark/src/main/scala/org/apache/comet/rules/CometCoalesceShufflePartitions.scala new file mode 100644 index 0000000000..731ff1eade --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/rules/CometCoalesceShufflePartitions.scala @@ -0,0 +1,156 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.rules + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.plans.physical.UnknownPartitioning +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.comet.CometExec +import org.apache.spark.sql.execution.{LeafExecNode, SparkPlan, UnionExec} +import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, AQEShuffleReadRule, CoalesceShufflePartitions, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.exchange.ShuffleOrigin +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec} + +/** + * Coalesces the shuffle partitions of a query stage that Spark's CoalesceShufflePartitions leaves + * alone because Comet operators stand where it looks for Spark ones. + * + * Spark coalesces each child of a `UnionExec` as a group of its own, and from Spark 3.5 each + * child of a `CartesianProductExec`, `BroadcastHashJoinExec` or `BroadcastNestedLoopJoinExec` + * too. It matches those classes, and the Comet operators that replace them are other classes, so + * it falls through to the case that coalesces only when every leaf below the operator is an + * exchange stage. A union with a scan or a table-cache stage in one branch then keeps every + * partition of the shuffles in the others: `spark.sql.shuffle.partitions` tasks for a query that + * needs a few. + * + * Comet replaces these operators while AQE prepares a stage, before its optimizer rules run, so + * this runs after Spark's rule instead. It rebuilds the Spark operator that each such Comet + * operator replaced, over the Comet children, has Spark's own rule coalesce the whole stage, and + * swaps the Comet operators back in. Spark's code makes every decision, for its version, so the + * partitions come out as Spark would have coalesced them, including the smaller target size it + * uses below a Cartesian product or a nested loop join from Spark 4.0. A shuffle that Spark's + * rule has coalesced already, or that another AQE rule reads its own way, stays as it is. + * + * It rebuilds only operators whose output partitioning is unknown. Spark can coalesce the + * children of a union differently, and from Spark 4.1 a union whose children share a partitioning + * reports it, so an aggregate above can rely on it instead of a shuffle. AQE discards a + * coalescing that breaks a distribution an operator requires, but Comet's operators state none, + * so the aggregate would read one key from more than one partition. A broadcast join reports the + * partitioning of its streamed side, so the same applies to it. + * + * When every leaf below such an operator is an exchange stage, Spark's rule already coalesces its + * shuffles, together rather than child by child, and this leaves them as they are. It steps in + * there only when Spark's rule cannot coalesce them together, as when they differ in partition + * count or one of them is a single-partition shuffle. + * + * Extending `AQEShuffleReadRule` gets this the same treatment from AQE as Spark's rule: it is + * skipped for the final stage when that stage's shuffle optimizations are off, and its result is + * discarded if it breaks a distribution required above it. + */ +case object CometCoalesceShufflePartitions extends AQEShuffleReadRule { + + // The Comet operator that a stand-in Spark operator was rebuilt from. + private val COMET_OPERATOR = TreeNodeTag[SparkPlan]("cometCoalesceShufflePartitions") + + // Required by the trait. Which shuffles are coalesced is decided by Spark's rule, which applies + // its own list. + override protected def supportedShuffleOrigins: Seq[ShuffleOrigin] = + CoalesceShufflePartitions(SparkSession.active).supportedShuffleOrigins + + override def apply(plan: SparkPlan): SparkPlan = { + if (!conf.coalesceShufflePartitionsEnabled || + !plan.exists(_.isInstanceOf[ShuffleQueryStageExec])) { + return plan + } + val asSpark = plan.transformDown { case comet @ Replaced(original) => + standIn(comet, original) + } + if (asSpark eq plan) { + return plan + } + // A read over a shuffle stage means an AQE rule has already decided how to read it: Spark's + // rule just coalesced it, below a Cartesian product say, or it is a skew-split or local read. + // Spark's rule expects to coalesce only reads that split a skewed partition, so hide each read + // behind a leaf that is not an exchange stage. Spark's rule then leaves alone every shuffle it + // would coalesce together with a read one, and coalesces the others. The hidden groups no + // longer count when it shares the minimum number of partitions among groups, so the others + // can keep more partitions than on Spark, never fewer. + val withDecidedReads = asSpark.transformUp { case read: AQEShuffleReadExec => + DecidedRead(read) + } + val coalesced = CoalesceShufflePartitions(SparkSession.active).apply(withDecidedReads) + if (coalesced eq withDecidedReads) plan else restore(coalesced) + } + + // A Comet operator whose Spark original's children Spark's rule coalesces one by one, with that + // original. The class match mirrors Spark's, and Spark's rule decides, for its version, which + // of these it actually treats that way. Comet has no counterpart of `CartesianProductExec`. + private object Replaced { + def unapply(plan: SparkPlan): Option[SparkPlan] = plan match { + case comet: CometExec if comet.outputPartitioning.isInstanceOf[UnknownPartitioning] => + comet.originalPlan match { + case original @ (_: UnionExec | _: BroadcastHashJoinExec | + _: BroadcastNestedLoopJoinExec) + if original.children.length == comet.children.length => + Some(original) + case _ => None + } + case _ => None + } + } + + private def standIn(comet: SparkPlan, original: SparkPlan): SparkPlan = { + val standIn = original.withNewChildren(comet.children) match { + // `withNewChildren` hands back the original itself when its children are these already, + // as for a union that AQE planned again over materialized stages. The tag must not land on + // the operator that the Comet one keeps, so copy it. + case same if same eq original => + original.makeCopy(original.productIterator.map(_.asInstanceOf[AnyRef]).toArray) + case copy => copy + } + standIn.setTagValue(COMET_OPERATOR, comet) + standIn + } + + // Stands in for a read that an AQE rule has already decided on while Spark's rule runs. + private case class DecidedRead(read: AQEShuffleReadExec) extends LeafExecNode { + override def output: Seq[Attribute] = read.output + + override protected def doExecute(): RDD[InternalRow] = + throw new UnsupportedOperationException(s"$nodeName never runs") + } + + // Put each Comet operator back over the children of its stand-in, and each decided read back in + // place of its leaf. Rebuilt by hand rather than with transformUp, which copies a replaced + // node's tags onto a replacement that has none, and so could leave the stand-in's tag on the + // Comet operator. + private def restore(plan: SparkPlan): SparkPlan = plan match { + case DecidedRead(read) => read + case _ => + val children = plan.children.map(restore) + plan.getTagValue(COMET_OPERATOR) match { + case Some(comet) => comet.withNewChildren(children) + case None => plan.withNewChildren(children) + } + } +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 82407881b8..e43346fc81 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -38,10 +38,10 @@ import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog import org.apache.spark.sql.execution._ -import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, BroadcastQueryStageExec, LogicalQueryStage} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, LogicalQueryStage} import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryTableScanExec} import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat -import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, ReusedExchangeExec, ShuffleExchangeExec} +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike} import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec, SortMergeJoinExec} import org.apache.spark.sql.execution.reuse.ReuseExchangeAndSubquery import org.apache.spark.sql.execution.window.WindowExec @@ -53,6 +53,7 @@ import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.{CometConf, CometExecIterator, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark41Plus, isSpark42Plus} +import org.apache.comet.rules.CometCoalesceShufflePartitions import org.apache.comet.serde.Config.ConfigMap import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} @@ -3594,6 +3595,233 @@ class CometExecSuite extends CometTestBase { } } + // https://github.com/apache/datafusion-comet/issues/6454 + test("AQE coalesces the shuffle partitions of a union whose other branch is a scan") { + // Spark coalesces each child of a union as its own group, but its rule did not recognize + // Comet's union, so the shuffled branch of a union with a scan kept every shuffle partition. + // Comet's rule defers to Spark's, so the query should come out partitioned as it is on Spark. + assume(isSpark35Plus, "Comet's query-stage optimizer rules need Spark 3.5+") + withTempPath { dir => + spark.range(0, 100, 1, 1).toDF("c").write.parquet(dir.getCanonicalPath) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1", + SQLConf.SHUFFLE_PARTITIONS.key -> "200") { + def query() = spark + .range(0, 10, 1, 2) + .toDF("c") + .repartition($"c") + .union(spark.read.parquet(dir.getCanonicalPath)) + var sparkPartitions = 0 + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + val df = query() + df.collect() + sparkPartitions = df.rdd.getNumPartitions + } + assert(sparkPartitions < 200, "Spark should have coalesced the shuffled branch") + + val df = query() + checkSparkAnswer(df) + // checkSparkAnswer runs copies of the query, so run this one to finalize its own plan. + df.collect() + val plan = df.queryExecution.executedPlan + assert(plan.asInstanceOf[AdaptiveSparkPlanExec].isFinalPlan) + assert(collect(plan) { case u: CometUnionExec => u }.size == 1) + assert(collect(plan) { case r: AQEShuffleReadExec if r.isCoalescedRead => r }.size == 1) + assert(df.rdd.getNumPartitions == sparkPartitions) + } + } + } + + // https://github.com/apache/spark/blob/v4.1.3/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala#L466-L484 + test("AQE coalesces each branch of a union whose shuffles have different partition counts") { + // Every leaf below the union is a shuffle stage, so Spark's rule coalesces its shuffles + // together, and gives up because the aggregate's single-partition shuffle cannot be + // coalesced. Below a Spark union it coalesces the join's shuffles as a group of their own. + assume(isSpark35Plus, "Comet's query-stage optimizer rules need Spark 3.5+") + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1", + SQLConf.SHUFFLE_PARTITIONS.key -> "5") { + val df = spark.range(3).join(spark.range(3), "id").union(spark.range(3).groupBy().sum()) + checkAnswer(df, (0 to 3).map(i => Row(i.toLong))) + val plan = df.queryExecution.executedPlan + assert(plan.asInstanceOf[AdaptiveSparkPlanExec].isFinalPlan) + assert(collect(plan) { case u: CometUnionExec => u }.size == 1) + assert( + collect(plan) { case r: AQEShuffleReadExec if r.isCoalescedRead => r }.size == 2, + plan) + } + } + + test("AQE leaves the shuffles of a union alone when an aggregate relies on its partitioning") { + // From Spark 4.1 a union whose children share a hash partitioning reports it, and the + // aggregate above reads the union without a shuffle. Coalescing the branch that can be + // coalesced, and not the other, would split each key between the branches' partitions, and + // Comet's aggregate, which requires no distribution, would return each key twice. + assume(isSpark35Plus, "Comet's query-stage optimizer rules need Spark 3.5+") + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1", + SQLConf.SHUFFLE_PARTITIONS.key -> "20") { + val keys = spark.range(0, 100, 1, 2).selectExpr("id % 10 AS k") + // repartition($"k") can be coalesced, and repartition(20, $"k") cannot. + val df = keys.repartition($"k").union(keys.repartition(20, $"k")).groupBy("k").count() + checkAnswer(df, (0 until 10).map(k => Row(k.toLong, 20L))) + val plan = df.queryExecution.executedPlan + assert(plan.asInstanceOf[AdaptiveSparkPlanExec].isFinalPlan) + assert(collect(plan) { case u: CometUnionExec => u }.size == 1) + if (isSpark41Plus) { + assert(collect(plan) { case e: ShuffleExchangeLike => e }.size == 2, plan) + assert(collect(plan) { case r: AQEShuffleReadExec => r }.isEmpty, plan) + } + } + } + + test("AQE coalesces the shuffle partitions of a union planned again over its stages") { + // When AQE plans a query again after its stages materialize, a union directly over those + // stages has the same children as the Spark union that Comet replaced. + assume(isSpark35Plus, "Comet's query-stage optimizer rules need Spark 3.5+") + withTempPath { dir => + spark.range(0, 100, 1, 1).toDF("c").write.parquet(dir.getCanonicalPath) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1", + SQLConf.SHUFFLE_PARTITIONS.key -> "200") { + val df = spark + .range(0, 10, 1, 2) + .toDF("c") + .repartition($"c") + .union(spark.read.parquet(dir.getCanonicalPath)) + df.collect() + val union = collect(df.queryExecution.executedPlan) { case u: CometUnionExec => u }.head + val children = union.children.map(_.transformUp { case r: AQEShuffleReadExec => r.child }) + val replanned = CometUnionExec(UnionExec(children), union.output, children) + val coalesced = CometCoalesceShufflePartitions(replanned) + assert(collect(coalesced) { + case r: AQEShuffleReadExec if r.isCoalescedRead => r + }.size == 1) + assert(coalesced.isInstanceOf[CometUnionExec]) + } + } + } + + test("AQE coalesces the shuffle partitions of a union below a join as Spark does") { + // Spark coalesces each child of a broadcast join, nested loop join or Cartesian product as a + // group of its own, and from Spark 4.0 it targets the minimum partition size rather than the + // advisory size below the latter two, which join every row with every row of the other side. + // The shuffled branch of the union should keep as many partitions as it does on Spark. + assume(isSpark35Plus, "Comet's query-stage optimizer rules need Spark 3.5+") + withTempPath { dir => + spark.range(0, 10, 1, 1).toDF("c").write.parquet(dir.getCanonicalPath) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_SIZE.key -> "1", + SQLConf.SHUFFLE_PARTITIONS.key -> "10") { + def union() = spark + .range(0, 1000, 1, 2) + .toDF("c") + .repartition($"c") + .union(spark.read.parquet(dir.getCanonicalPath)) + def other() = spark.range(0, 3).toDF("d") + // The number of partitions the union reads from its shuffled branch. + def shuffledBranchPartitions(df: DataFrame): Int = { + df.collect() + val unions = collect(df.queryExecution.executedPlan) { + case u: UnionExec => u + case u: CometUnionExec => u + } + assert(unions.size == 1, df.queryExecution.executedPlan) + unions.head.children.head.outputPartitioning.numPartitions + } + // Each join, the operator Comet plans for it, and whether it joins every row with every + // row of the other side. + val joins = Seq( + ( + "Cartesian product", + () => union().crossJoin(other()), + classOf[CartesianProductExec], + true), + ( + "broadcast nested loop join", + () => union().join(broadcast(other()), $"c" < $"d"), + classOf[CometBroadcastNestedLoopJoinExec], + true), + ( + "broadcast hash join", + () => union().join(broadcast(other()), $"c" === $"d"), + classOf[CometBroadcastHashJoinExec], + false)) + joins.foreach { case (join, query, joinClass, explodingJoin) => + var sparkPartitions = 0 + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkPartitions = shuffledBranchPartitions(query()) + } + if (isSpark40Plus && explodingJoin) { + assert(sparkPartitions == 10, s"Spark should not coalesce below a $join") + } else { + assert(sparkPartitions == 1, s"Spark should coalesce below a $join") + } + + val df = query() + checkSparkAnswer(df) + assert(shuffledBranchPartitions(df) == sparkPartitions, join) + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case u: CometUnionExec => u }.size == 1, plan) + assert(collect(plan) { case j if joinClass.isInstance(j) => j }.size == 1, plan) + } + } + } + } + + test("AQE coalesces the shuffle partitions of a union in a stage Spark coalesced in part") { + // Spark's rule coalesces the shuffle on the other side of a Cartesian product as a group of + // its own before Comet's rule runs. The read it leaves there should not keep the union's + // shuffled branch from being coalesced as it is on Spark. + assume(isSpark35Plus, "Comet's query-stage optimizer rules need Spark 3.5+") + withTempPath { dir => + spark.range(0, 10, 1, 1).toDF("c").write.parquet(dir.getCanonicalPath) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + // Keeps AQE from broadcasting the small side, which would move its shuffle and read into + // a stage of their own. + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1", + SQLConf.SHUFFLE_PARTITIONS.key -> "10") { + def query() = spark + .range(0, 1000, 1, 2) + .toDF("c") + .repartition($"c") + .union(spark.read.parquet(dir.getCanonicalPath)) + .crossJoin(spark.range(0, 3).toDF("d").repartition($"d")) + // The number of partitions that each coalesced read of the final plan reads. + def coalescedReads(df: DataFrame): Seq[Int] = { + df.collect() + collect(df.queryExecution.executedPlan) { + case r: AQEShuffleReadExec if r.isCoalescedRead => r.partitionSpecs.length + } + } + var sparkReads = Seq.empty[Int] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sparkReads = coalescedReads(query()) + } + assert(sparkReads == Seq(1, 1), "Spark should coalesce both sides of the product") + + val df = query() + checkSparkAnswer(df) + assert(coalescedReads(df) == sparkReads, df.queryExecution.executedPlan) + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case u: CometUnionExec => u }.size == 1, plan) + assert(collect(plan) { case j: CartesianProductExec => j }.size == 1, plan) + } + } + } + test("native execution after coalesce") { withTable("t1") { (0 until 5) 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 4c90612c3a..b4d3f6c66c 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -199,6 +199,26 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + // https://github.com/apache/spark/blob/v4.1.2/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala#L3178-L3191 + test("AQE SPARK-42101: coalesce the shuffle partitions of a union with a table cache stage") { + assume(isSpark35Plus, "Table-cache query stages require Spark 3.5+") + withAQECache { + withSQLConf(SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1") { + val cached = Seq(1).toDF("c").cache() + val df = Seq(2).toDF("c").repartition($"c").union(cached) + checkAnswer(df, Seq(Row(1), Row(2))) + val plan = df.queryExecution.executedPlan + assert(plan.asInstanceOf[AdaptiveSparkPlanExec].isFinalPlan) + assert(collect(plan) { case u: org.apache.spark.sql.comet.CometUnionExec => u }.size == 1) + assert(collect(plan) { case r @ AQEShuffleReadExec(_: ShuffleQueryStageExec, _) => + r + }.size == 1) + assert(collect(plan) { case s: QueryStageExec if isTableCacheStage(s) => s }.size == 1) + assert(collect(plan) { case s: CometInMemoryTableScanExec => s }.size == 1) + } + } + } + // https://github.com/apache/spark/blob/v4.1.2/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala#L2780-L2832 test("AQE SPARK-37742: use valid Comet cache statistics for join selection") { withAQECache { 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 046fe3b822..036b6fd8de 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 @@ -28,6 +28,8 @@ 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.ColumnarToRowExec +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} @@ -237,6 +239,7 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { runCodecBenchmark(flatRelation) runSparkOperatorBenchmark(flatRelation) + runAdaptiveBenchmark(flatRelation) } runWideSparkOperatorBenchmark() } @@ -348,13 +351,13 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { * Reads that feed Spark operators rather than Comet ones, against Spark's own cache format. * * Native execution 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. Comet's format is read two ways. With Comet off, every Spark operator reads rows - * from the cache scan's row reader. With Comet on, a generated Spark operator instead reads the - * cached vectors through a ColumnarToRowExec fused into its generated code, which is how these - * reads run when Comet is enabled without native execution. With native execution on, Spark - * operators above the cache read CometInMemoryTableScan's batches through CometColumnarToRow - * instead, which this does not measure. + * shape where Comet's format has the most to lose. Comet's format is read two ways. With Comet + * off, every Spark operator reads rows from the cache scan's row reader. With Comet on, a + * generated Spark operator instead reads the cached vectors through a ColumnarToRowExec fused + * into its generated code, which is how these reads run when Comet is enabled without native + * execution. With native execution on, Spark operators above the cache read + * CometInMemoryTableScan's batches through CometColumnarToRow instead, which + * runAdaptiveBenchmark measures. * * 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 and which reader it uses. @@ -377,14 +380,7 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { .createOrReplaceTempView(view) val cache = new OneCachedCopy(view) - 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, @@ -446,17 +442,18 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { fused: Boolean = false) /** - * Holds one cached copy of `view` at a time. Two copies could not coexist anyway: the cache - * manager keys on the plan rather than the name, so a second one would find the first. + * Holds one cached copy of `view` at a time, cached under `conf`. Two copies could not coexist + * anyway: the cache manager keys on the plan rather than the name, so a second one would find + * the first. */ - private class OneCachedCopy(view: String) { + private class OneCachedCopy(view: String, conf: Seq[(String, String)] = sparkOperatorConf) { private var cachedBy: String = _ def cacheBy(serializer: String): Unit = if (cachedBy != serializer) { spark.catalog.uncacheTable(view) cachedBy = null withCacheSerializer(serializer) { - withSQLConf(sparkOperatorConf: _*) { + withSQLConf(conf: _*) { spark.catalog.cacheTable(view) spark.table(view).count() } @@ -488,6 +485,116 @@ 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.convert.inMemoryCache.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" -> sparkSerializer, "Comet's cache format" -> cometSerializer) + 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) + + val cache = new OneCachedCopy(view, adaptiveConf) + 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 => + cache.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 == cometSerializer + 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 @@ -629,8 +736,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}") @@ -713,6 +820,16 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { "spark.comet.exec.onHeap.enabled" -> "true", "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") + private val sparkSerializer = classOf[DefaultCachedBatchSerializer].getName private val cometSerializer = classOf[ArrowCachedBatchSerializer].getName }