diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index d85eb14d2c4..f6b85808f2d 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -515,6 +515,8 @@ jobs: org.apache.comet.exec.CometInMemoryCacheSuite org.apache.comet.exec.CometInMemoryCachePruningSuite org.apache.comet.exec.CometInMemoryCacheKryoSuite + org.apache.comet.exec.CometInMemoryCacheKryoUnregisteredSuite + org.apache.comet.exec.CometInMemoryCacheKryoClassesToRegisterSuite org.apache.comet.exec.CometGenerateExecSuite org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 02d3077e4fc..7d6b1ac9fc9 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -219,6 +219,8 @@ jobs: org.apache.comet.exec.CometInMemoryCacheSuite org.apache.comet.exec.CometInMemoryCachePruningSuite org.apache.comet.exec.CometInMemoryCacheKryoSuite + org.apache.comet.exec.CometInMemoryCacheKryoUnregisteredSuite + org.apache.comet.exec.CometInMemoryCacheKryoClassesToRegisterSuite org.apache.comet.exec.CometGenerateExecSuite org.apache.comet.exec.CometWindowExecSuite org.apache.comet.exec.CometJoinSuite diff --git a/docs/source/contributor-guide/plugin_overview.md b/docs/source/contributor-guide/plugin_overview.md index 89586bf7b6d..072f59e01a2 100644 --- a/docs/source/contributor-guide/plugin_overview.md +++ b/docs/source/contributor-guide/plugin_overview.md @@ -48,6 +48,8 @@ and skips the remaining steps. Otherwise it: - Appends `CometSparkSessionExtensions` to `spark.sql.extensions`, unless it is already listed. - Sets `spark.sql.cache.serializer` to Comet's `ArrowCachedBatchSerializer` when `spark.comet.exec.inMemoryCache.enabled=true`, unless the application has chosen a different serializer. + It leaves Spark's serializer in place where Comet could not scan its format natively or Kryo would + reject it; see [In-Memory Cache](../user-guide/latest/in-memory-cache.md). - Registers `CometSource` with Spark's metrics system and adds `CometMetricsListener` to `spark.sql.queryExecutionListeners` when `spark.comet.metrics.enabled=true`. - Logs a warning for settings that are likely to cause problems, such as an unset `spark.executor.memoryOverhead`. diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md index 8de4759957f..fb3ba9cc632 100644 --- a/docs/source/user-guide/latest/in-memory-cache.md +++ b/docs/source/user-guide/latest/in-memory-cache.md @@ -34,11 +34,13 @@ $SPARK_HOME/bin/spark-shell \ ``` It has to be set before the `SparkContext` starts. Comet's driver plugin chooses -`spark.sql.cache.serializer` once, while the context is initializing, so a session that started -with the default goes on using Spark's cache format however the config is set afterwards. The -plugin installs Comet's serializer only if `spark.comet.enabled` and `spark.comet.exec.enabled` -are enabled at that point too, because an application that starts without native execution could -not scan Comet's format natively. +`spark.sql.cache.serializer` once, while the context is initializing, so an application keeps the +cache format it started with however the config is set afterwards. The plugin installs Comet's +serializer only if `spark.comet.enabled` and `spark.comet.exec.enabled` are enabled at that point +too, because an application that starts without native execution could not scan Comet's format +natively. It also keeps Spark's format when Comet shuffle is enabled but `spark.shuffle.manager` is +not one of Comet's shuffle managers, since Comet then disables itself, and when Kryo would reject +Comet's format; see [Kryo](#kryo). ## What changes when it is enabled @@ -179,11 +181,20 @@ spark.kryo.registrator=org.apache.comet.CometKryoRegistrator Comet cannot set `spark.kryo.registrator` for you the way it sets `spark.sql.cache.serializer`: `KryoSerializer` reads it when `SparkEnv` builds the serializer, which happens before any plugin -runs. Without it, caching fails with a "Class is not registered" error that does not name this -feature. Comet's driver plugin warns at startup when it sees Kryo, `registrationRequired`, and no -registrator. Native broadcast needs the same registrator even when the cache is disabled; see +runs. Without it, Kryo would reject Comet's cached batch with a "Class is not registered" error +that does not name this feature. So when Kryo requires registration and has not registered +Comet's cached batch, Comet's driver plugin does not install Comet's serializer, and caches stay in +Spark's format. Registrations made another way, through a registrator of the application's own or +`spark.kryo.classesToRegister`, count as well. The plugin warns at startup when Kryo requires +registration and has not registered every class `CometKryoRegistrator` registers. An application +that sets `spark.sql.cache.serializer` to Comet's serializer itself gets the error instead. Native +broadcast needs the same registrator even when the cache is disabled; see [Kryo serialization](installation.md#kryo-serialization). +Spark registers its own cached batch with Kryo only from Spark 4.1, so on earlier versions caching +in either format under `registrationRequired` needs a registrator. `CometKryoRegistrator` registers +Spark's cached batch too. + ## Limitations Reads that feed **Spark** operators rather than Comet ones are slower than Spark's own cache diff --git a/docs/source/user-guide/latest/installation.md b/docs/source/user-guide/latest/installation.md index b9d28b74b03..c6ca172ab06 100644 --- a/docs/source/user-guide/latest/installation.md +++ b/docs/source/user-guide/latest/installation.md @@ -269,8 +269,9 @@ If the application uses Kryo (`spark.serializer=org.apache.spark.serializer.Kryo Without it, any query that uses Comet's native broadcast exchange, which is enabled by default, fails with Kryo's "Class is not registered" error, for example on the first broadcast hash join. -The [in-memory cache](in-memory-cache.md#kryo) needs the same registrator. Set it before the -`SparkContext` is created: `KryoSerializer` reads it before Comet's plugin runs, so Comet cannot -add it for you. `spark.kryo.registrator` accepts a comma-separated list, so an application with -its own registrator can list both. Comet logs a warning at startup when Kryo requires registration -and this registrator is missing. +Comet's [in-memory cache](in-memory-cache.md#kryo) format needs the same registrations, and while +Kryo has not registered Comet's cached batch, Comet's plugin keeps caches in Spark's format. Set it +before the `SparkContext` is created: `KryoSerializer` reads it before Comet's plugin runs, so +Comet cannot add it for you. `spark.kryo.registrator` accepts a comma-separated list, so an +application with its own registrator can list both. Comet logs a warning at startup when Kryo +requires registration and has not registered the classes this registrator covers. diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6cfe5bc92bd..985b18b9cb1 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -275,7 +275,8 @@ object CometConf extends ShimCometConf { .doc("Whether to enable Comet native execution for in-memory cached tables. Its value at " + "startup also decides whether CometDriverPlugin installs Comet's cache serializer, " + "which stores cached data in Arrow format. The plugin installs it only if " + - "spark.comet.enabled and spark.comet.exec.enabled are also enabled at startup. " + + "spark.comet.enabled and spark.comet.exec.enabled are also enabled at startup, and " + + "only with one of Comet's shuffle managers while Comet shuffle is enabled. " + "Because spark.sql.cache.serializer is a " + "static config, the cached format is fixed for the application, and disabling this " + "at runtime only sends cached scans back to Spark's execution path. Relations whose " + @@ -284,10 +285,12 @@ object CometConf extends ShimCometConf { "zstd compression, and a scan copies out only the buffers of the columns it projected, " + "so the unselected ones are never decompressed. Reads that feed Spark operators rather " + "than Comet ones still pay a row conversion the default format avoids, and can be " + - "slower than Spark's cache. With spark.kryo.registrationRequired=true, also set " + - "spark.kryo.registrator=org.apache.comet.CometKryoRegistrator before creating the " + - "SparkContext, otherwise caching fails as soon as a block is serialized, including " + - "the disk half of the default MEMORY_AND_DISK storage level.") + "slower than Spark's cache. With spark.kryo.registrationRequired=true, the plugin " + + "installs it only if Kryo has registered Comet's cached batch, as " + + "spark.kryo.registrator=org.apache.comet.CometKryoRegistrator does when set before " + + "creating the SparkContext, because Kryo would otherwise reject a cached block as soon " + + "as it is serialized, including the disk half of the default MEMORY_AND_DISK storage " + + "level.") .booleanConf .createWithDefault(false) diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala index c3246679d02..e79074f45d2 100644 --- a/spark/src/main/scala/org/apache/spark/Plugins.scala +++ b/spark/src/main/scala/org/apache/spark/Plugins.scala @@ -30,6 +30,9 @@ import org.apache.spark.api.plugin.{DriverPlugin, ExecutorPlugin, PluginContext, import org.apache.spark.internal.Logging import org.apache.spark.internal.config.{EVENT_LOG_ENABLED, EXECUTOR_MEMORY_OVERHEAD, EXECUTOR_MEMORY_OVERHEAD_FACTOR} import org.apache.spark.scheduler.{SparkListener, SparkListenerApplicationEnd, SparkListenerExecutorMetricsUpdate, SparkListenerExecutorRemoved} +import org.apache.spark.serializer.KryoSerializer +import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer +import org.apache.spark.sql.comet.execution.shuffle.{CometCelebornShuffleManager, CometShuffleManager} import org.apache.spark.sql.internal.StaticSQLConf import org.apache.spark.util.{Clock, SystemClock} @@ -110,7 +113,7 @@ class CometDriverPlugin private[spark] (clock: Clock) extends DriverPlugin with val extraConfs = new ju.HashMap[String, String]() CometDriverPlugin.maybeSetCacheSerializer(sc.conf, extraConfs) - CometDriverPlugin.warnIfKryoRegistratorMissing(sc.conf) + CometDriverPlugin.warnIfKryoRegistrationsMissing(sc.conf) // register CometSparkSessionExtensions if it isn't already registered CometDriverPlugin.registerCometSessionExtension(sc.conf) @@ -174,7 +177,13 @@ object CometDriverPlugin extends Logging { // Use Comet's cache serializer only when the native in-memory cache scan can run, which needs // Comet and its native execution as well as the cache config. spark.sql.cache.serializer is // static, so an application that starts with Comet or native execution off would otherwise - // store every cache in Comet's format, with only Spark operators to read it. + // store every cache in Comet's format, with only Spark operators to read it. So would one that + // leaves Comet shuffle enabled without Comet's shuffle manager, since Comet then disables + // itself. + // Nor is it used where Kryo requires registration and has not registered Comet's cached batch: + // caching would then fail the first time Spark serialized a cached block. Where Kryo has + // registered it, by whatever means, Comet's format is used, since Spark registers its own + // cached batch only from 4.1. // If the application already set spark.sql.cache.serializer, leave that value // unchanged so Comet does not replace a user-selected cache format. private[apache] def maybeSetCacheSerializer( @@ -182,7 +191,9 @@ object CometDriverPlugin extends Logging { extraConfs: ju.HashMap[String, String]): Unit = { if (getBooleanConf(conf, CometConf.COMET_ENABLED) && getBooleanConf(conf, CometConf.COMET_EXEC_ENABLED) && - getBooleanConf(conf, CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED)) { + getBooleanConf(conf, CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED) && + (!getBooleanConf(conf, CometConf.COMET_SHUFFLE_ENABLED) || isCometShuffleManager(conf)) && + !unregisteredKryoClasses(conf).contains(ArrowCachedBatchSerializer.cachedBatchClass)) { val serializerKey = StaticSQLConf.SPARK_CACHE_SERIALIZER.key val serializerValue = "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer" @@ -207,27 +218,54 @@ object CometDriverPlugin extends Logging { // CometKryoRegistrator covers both, but spark.kryo.registrator is read when SparkEnv builds the // serializer, before any plugin runs, so it cannot be set from here. Say so while the // application is still starting up rather than leaving the user to attribute the failure later. - private[apache] def warnIfKryoRegistratorMissing(conf: SparkConf): Unit = { + private[apache] def warnIfKryoRegistrationsMissing(conf: SparkConf): Unit = { + val unregistered = unregisteredKryoClasses(conf) + if (unregistered.nonEmpty) { + logWarning("spark.kryo.registrationRequired=true but Kryo has not registered " + + s"${unregistered.map(_.getName).mkString(", ")}, which " + + s"${CometKryoRegistrator.CLASS_NAME} registers. Comet's native broadcast and in-memory " + + "cache fail with Kryo's \"Class is not registered\" when they serialize one of them, " + + "and Comet keeps Spark's cache format while its own cached batch is unregistered. " + + s"Add spark.kryo.registrator=${CometKryoRegistrator.CLASS_NAME} before creating the " + + "SparkContext; it cannot be set later.") + } + } + + // The classes CometKryoRegistrator registers that Kryo, configured as the application + // configured it, would reject: none unless it requires registration. They can be registered + // through CometKryoRegistrator, a registrator of the application's own or + // spark.kryo.classesToRegister, so ask a Kryo instance built from the conf rather than read the + // confs. If one cannot be built, take them as registered only if spark.kryo.registrator lists + // CometKryoRegistrator. + private[apache] def unregisteredKryoClasses(conf: SparkConf): Seq[Class[_]] = { val usingKryo = conf.get("spark.serializer", "") == "org.apache.spark.serializer.KryoSerializer" - val registrationRequired = conf.getBoolean("spark.kryo.registrationRequired", false) - val registered = conf - .get("spark.kryo.registrator", "") - .split(',') - .map(_.trim) - .contains(CometKryoRegistrator.CLASS_NAME) - - if (usingKryo && registrationRequired && !registered) { - logWarning( - "spark.kryo.registrationRequired=true but spark.kryo.registrator does not include " + - s"${CometKryoRegistrator.CLASS_NAME}. Comet's native broadcast and its in-memory " + - "cache format will fail with Kryo's \"Class is not registered\" as soon as their " + - "payloads are serialized. Add " + - s"spark.kryo.registrator=${CometKryoRegistrator.CLASS_NAME} before creating the " + - "SparkContext; it cannot be set later.") + if (!usingKryo || !conf.getBoolean("spark.kryo.registrationRequired", false)) { + Nil + } else { + // Qualified, because in this package org.apache.spark.Success, a TaskEndReason, hides an + // imported scala.util.Success on Scala 2.12. + Try(new KryoSerializer(conf).newKryo()) match { + case scala.util.Success(kryo) => + CometKryoRegistrator.classes.filter(kryo.getClassResolver.getRegistration(_) == null) + case scala.util.Failure(e) => + logDebug("Could not build Kryo to check Comet's registrations", e) + val listed = conf + .get("spark.kryo.registrator", "") + .split(',') + .map(_.trim) + .contains(CometKryoRegistrator.CLASS_NAME) + if (listed) Nil else CometKryoRegistrator.classes + } } } + // Comet's shuffle managers have no short name, so spark.shuffle.manager names one only by its + // class name. + private def isCometShuffleManager(conf: SparkConf): Boolean = + Set(classOf[CometShuffleManager].getName, classOf[CometCelebornShuffleManager].getName) + .contains(conf.get("spark.shuffle.manager", "")) + // Comet's native allocations are made by the Rust global allocator and live in the native heap. // In off-heap mode the share that operators reserve is charged against a memory pool, but // everything else -- expression kernels and Arrow array builders, decompression buffers, Parquet @@ -298,8 +336,13 @@ object CometDriverPlugin extends Logging { } } + // Reads a deprecated alternative too, such as spark.comet.exec.shuffle.enabled, as a session + // would. private def getBooleanConf(conf: SparkConf, entry: ConfigEntry[Boolean]): Boolean = - conf.getBoolean(entry.key, entry.defaultValue.get) + (entry.key +: entry.alternatives) + .find(conf.contains) + .map(conf.getBoolean(_, entry.defaultValue.get)) + .getOrElse(entry.defaultValue.get) def registerCometMetrics(sc: SparkContext): Unit = { if (sc.getConf.getBoolean( diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala index 103df61e17f..853c731a94a 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala @@ -739,6 +739,11 @@ object ArrowCachedBatchSerializer { def supportsSchema(schema: Seq[Attribute]): Boolean = schema.forall(a => supportsType(a.dataType)) + /** + * The class of Comet's cached batch, which Kryo has to have registered to store this format. + */ + private[apache] val cachedBatchClass: Class[_] = classOf[CometCachedBatch] + /** * The classes a `CometCachedBatch` adds on top of [[org.apache.comet.CometKryoRegistrator]]'s * shared Arrow-bytes classes. diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala index b403442531c..4518f8ab2ee 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala @@ -20,10 +20,16 @@ package org.apache.comet.exec import org.apache.spark.SparkConf +import org.apache.spark.serializer.KryoRegistrator import org.apache.spark.sql.{CometTestBase, Row} -import org.apache.spark.sql.execution.columnar.CometInMemoryRelationHelper +import org.apache.spark.sql.catalyst.expressions.GenericInternalRow +import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer +import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, DefaultCachedBatch} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.storage.StorageLevel +import org.apache.spark.unsafe.types.UTF8String + +import com.esotericsoftware.kryo.Kryo import org.apache.comet.{CometConf, CometKryoRegistrator} @@ -194,3 +200,106 @@ class CometInMemoryCacheKryoSuite extends CometTestBase { } } } + +/** + * Comet's driver plugin under Kryo with registration required, in an application that does not + * list [[CometKryoRegistrator]]. The plugin installs Comet's cache serializer only if Kryo has + * Comet's cached batch registered, whatever registered it, so the format it picks has to survive + * a `DISK_ONLY` cache, which serializes every block as it is put. + * + * Spark registers its own cached batch with Kryo only from 4.1, so on earlier versions an + * application whose caches work under registration registers it itself, as the suite that expects + * Spark's format does. + */ +abstract class CometInMemoryCacheKryoRegistrationSuite(expectedBatch: String) + extends CometTestBase { + + /** The Kryo registrations the application makes instead of listing CometKryoRegistrator. */ + protected def registrations: Seq[(String, String)] + + override protected def beforeAll(): Unit = { + CometInMemoryRelationHelper.clearSerializer() + super.beforeAll() + } + + override protected def afterAll(): Unit = { + try { + super.afterAll() + } finally { + CometInMemoryRelationHelper.clearSerializer() + } + } + + override protected def sparkConf: SparkConf = { + val conf = super.sparkConf + conf.set("spark.plugins", "org.apache.spark.CometPlugin") + conf.set(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key, "true") + conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer") + conf.set("spark.kryo.registrationRequired", "true") + conf.setAll(registrations) + } + + private def cachedBatchTypes(table: String): Array[String] = { + val cached = spark.sharedState.cacheManager.lookupCachedData(spark.table(table)).get + cached.cachedRepresentation.cacheBuilder.cachedColumnBuffers + .map(_.getClass.getName) + .distinct() + .collect() + } + + test("Comet plugin picks a cache format that Kryo can store") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + spark.catalog.clearCache() + try { + spark + .range(0, 100, 1, 2) + .selectExpr("id", "cast(id as string) AS s") + .createOrReplaceTempView("kryo_registration") + + spark.catalog.cacheTable("kryo_registration", StorageLevel.DISK_ONLY) + assert(spark.table("kryo_registration").count() == 100) + assert(cachedBatchTypes("kryo_registration").sameElements(Array(expectedBatch))) + + checkAnswer( + spark.sql("SELECT s FROM kryo_registration WHERE id > 97"), + Seq(Row("98"), Row("99"))) + } finally { + spark.catalog.clearCache() + } + } + } +} + +/** Registers Spark's cached batch and nothing of Comet's, so the plugin keeps Spark's format. */ +class CometInMemoryCacheKryoUnregisteredSuite + extends CometInMemoryCacheKryoRegistrationSuite(classOf[DefaultCachedBatch].getName) { + override protected def registrations: Seq[(String, String)] = + Seq("spark.kryo.registrator" -> classOf[SparkCachedBatchKryoRegistrator].getName) +} + +/** + * Registers Comet's classes through `spark.kryo.classesToRegister` rather than + * [[CometKryoRegistrator]], and not Spark's cached batch, so before Spark 4.1 only Comet's format + * can be stored, and the plugin has to install it. + */ +class CometInMemoryCacheKryoClassesToRegisterSuite + extends CometInMemoryCacheKryoRegistrationSuite( + ArrowCachedBatchSerializer.cachedBatchClass.getName) { + override protected def registrations: Seq[(String, String)] = Seq( + "spark.kryo.classesToRegister" -> CometKryoRegistrator.classes + .filterNot(_ == classOf[DefaultCachedBatch]) + .map(_.getName) + .mkString(",")) +} + +/** Registers what Spark's own cached batch needs, over a long and a string column, with Kryo. */ +class SparkCachedBatchKryoRegistrator extends KryoRegistrator { + override def registerClasses(kryo: Kryo): Unit = { + // The batch and its statistics row, whose bounds include UTF8String. + Seq( + classOf[DefaultCachedBatch], + classOf[GenericInternalRow], + classOf[Array[Any]], + classOf[UTF8String]).foreach(kryo.register) + } +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index ecce9067522..e913a21a69b 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -39,6 +39,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeRefer import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec, CometSortExec, CometSortMergeJoinExec} import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} +import org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.SortExec import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec} @@ -51,7 +52,7 @@ import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.spark.storage.StorageLevel -import org.apache.comet.{CometArrowAllocator, CometConf, ExtendedExplainInfo} +import org.apache.comet.{CometArrowAllocator, CometConf, CometKryoRegistrator, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} import org.apache.comet.vector.{CometPlainVector, CometVector} @@ -1021,6 +1022,7 @@ class CometInMemoryCacheSuite extends CometTestBase { val defaultConf = new SparkConf() .set(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key, "true") + .set("spark.shuffle.manager", shuffleManager) val defaultExtraConfs = new ju.HashMap[String, String]() // With no user serializer configured, the plugin should install Comet's @@ -1032,6 +1034,7 @@ class CometInMemoryCacheSuite extends CometTestBase { val userConf = new SparkConf() .set(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key, "true") + .set("spark.shuffle.manager", shuffleManager) .set(serializerKey, userSerializer) val userExtraConfs = new ju.HashMap[String, String]() @@ -1043,20 +1046,25 @@ class CometInMemoryCacheSuite extends CometTestBase { assert(!userExtraConfs.containsKey(serializerKey)) } - test("Comet plugin installs its cache serializer only if Comet can scan the cache natively") { + /** Whether the Comet plugin installs its cache serializer for an application's `settings`. */ + private def installsCacheSerializer(settings: (String, String)*): Boolean = { val serializerKey = StaticSQLConf.SPARK_CACHE_SERIALIZER.key + val conf = new SparkConf().setAll(settings) + val extraConfs = new ju.HashMap[String, String]() + CometDriverPlugin.maybeSetCacheSerializer(conf, extraConfs) + assert(conf.contains(serializerKey) == extraConfs.containsKey(serializerKey)) + extraConfs.containsKey(serializerKey) + } - def installed(settings: (String, String)*): Boolean = { - val conf = new SparkConf().setAll(settings) - val extraConfs = new ju.HashMap[String, String]() - CometDriverPlugin.maybeSetCacheSerializer(conf, extraConfs) - assert(conf.contains(serializerKey) == extraConfs.containsKey(serializerKey)) - extraConfs.containsKey(serializerKey) - } - + test("Comet plugin installs its cache serializer only if Comet can scan the cache natively") { val cometOn = CometConf.COMET_ENABLED.key -> "true" val execOn = CometConf.COMET_EXEC_ENABLED.key -> "true" val cacheOn = CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true" + // Without Comet's shuffle manager Comet disables itself, which the next test covers. + val cometShuffle = "spark.shuffle.manager" -> shuffleManager + + def installed(settings: (String, String)*): Boolean = + installsCacheSerializer(cometShuffle +: settings: _*) assert(installed(cometOn, execOn, cacheOn)) // An application that starts with Comet or its native execution off can never plan @@ -1074,6 +1082,85 @@ class CometInMemoryCacheSuite extends CometTestBase { CometConf.COMET_EXEC_ENABLED.defaultValue.get)) } + test("Comet plugin keeps Spark's cache format where Comet disables itself or Kryo rejects it") { + val cometShuffle = "spark.shuffle.manager" -> shuffleManager + + def installed(settings: (String, String)*): Boolean = + installsCacheSerializer( + Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") ++ settings: _*) + + assert(installed(cometShuffle)) + // Comet shuffle is enabled by default, and Comet disables itself while it is unless the + // application runs one of Comet's shuffle managers. + assert(!installed()) + assert(!installed("spark.shuffle.manager" -> "sort")) + assert(installed("spark.shuffle.manager" -> classOf[CometCelebornShuffleManager].getName)) + // Without Comet shuffle, under the key or its deprecated name, the shuffle manager does not + // matter. + assert(installed(CometConf.COMET_SHUFFLE_ENABLED.key -> "false")) + assert(installed("spark.comet.exec.shuffle.enabled" -> "false")) + + // Kryo with registration required rejects Comet's cached batch unless something registered + // it: CometKryoRegistrator, on its own or beside a registrator of the application's, or the + // application's own registrations. + val kryo = "spark.serializer" -> "org.apache.spark.serializer.KryoSerializer" + val registrationRequired = "spark.kryo.registrationRequired" -> "true" + val registrator = "spark.kryo.registrator" + val sparkOnly = classOf[SparkCachedBatchKryoRegistrator].getName + assert(!installed(cometShuffle, kryo, registrationRequired)) + assert(!installed(cometShuffle, kryo, registrationRequired, registrator -> sparkOnly)) + assert( + installed( + cometShuffle, + kryo, + registrationRequired, + registrator -> s"$sparkOnly, ${CometKryoRegistrator.CLASS_NAME}")) + assert( + installed( + cometShuffle, + kryo, + registrationRequired, + "spark.kryo.classesToRegister" -> ArrowCachedBatchSerializer.cachedBatchClass.getName)) + // A registrator that cannot be loaded leaves only spark.kryo.registrator to go by. + assert(!installed(cometShuffle, kryo, registrationRequired, registrator -> "com.example.R")) + assert( + installed( + cometShuffle, + kryo, + registrationRequired, + registrator -> s"com.example.R, ${CometKryoRegistrator.CLASS_NAME}")) + // Without registrationRequired, Kryo writes the class name of anything unregistered instead. + assert(installed(cometShuffle, kryo)) + } + + test("Comet plugin finds the Kryo registrations Comet needs however they were made") { + def unregistered(settings: (String, String)*): Seq[Class[_]] = + CometDriverPlugin.unregisteredKryoClasses(new SparkConf().setAll(settings)) + + val kryo = "spark.serializer" -> "org.apache.spark.serializer.KryoSerializer" + val registrationRequired = "spark.kryo.registrationRequired" -> "true" + assert(unregistered().isEmpty) + assert(unregistered(kryo).isEmpty) + assert( + unregistered(kryo, registrationRequired).contains( + ArrowCachedBatchSerializer.cachedBatchClass)) + assert( + unregistered( + kryo, + registrationRequired, + "spark.kryo.registrator" -> CometKryoRegistrator.CLASS_NAME).isEmpty) + assert( + unregistered( + kryo, + registrationRequired, + "spark.kryo.classesToRegister" -> CometKryoRegistrator.classes + .map(_.getName) + .mkString(",")).isEmpty) + } + test("Comet in-memory cache supports empty projection scan") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",