|
| 1 | +/* |
| 2 | + * Licensed to the Apache Software Foundation (ASF) under one |
| 3 | + * or more contributor license agreements. See the NOTICE file |
| 4 | + * distributed with this work for additional information |
| 5 | + * regarding copyright ownership. The ASF licenses this file |
| 6 | + * to you under the Apache License, Version 2.0 (the |
| 7 | + * "License"); you may not use this file except in compliance |
| 8 | + * with the License. You may obtain a copy of the License at |
| 9 | + * |
| 10 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | + * |
| 12 | + * Unless required by applicable law or agreed to in writing, |
| 13 | + * software distributed under the License is distributed on an |
| 14 | + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 15 | + * KIND, either express or implied. See the License for the |
| 16 | + * specific language governing permissions and limitations |
| 17 | + * under the License. |
| 18 | + */ |
| 19 | + |
| 20 | +package org.apache.comet.exec |
| 21 | + |
| 22 | +import org.apache.spark.SparkConf |
| 23 | +import org.apache.spark.sql.CometTestBase |
| 24 | +import org.apache.spark.sql.execution.columnar.CometInMemoryRelationHelper |
| 25 | +import org.apache.spark.sql.internal.SQLConf |
| 26 | +import org.apache.spark.storage.StorageLevel |
| 27 | + |
| 28 | +import org.apache.comet.{CometConf, CometKryoRegistrator} |
| 29 | + |
| 30 | +/** |
| 31 | + * Covers Comet's cached batch format under `spark.kryo.registrationRequired=true`. |
| 32 | + * |
| 33 | + * Kryo then rejects any class it has not been told about, and Spark serializes a `CachedBatch` |
| 34 | + * whenever a cached block leaves the heap: the disk half of the default `MEMORY_AND_DISK`, the |
| 35 | + * `_SER` levels, replication, and cross-executor fetches. So this is not a `DISK_ONLY`-only |
| 36 | + * concern -- a plain `df.cache()` that spills is enough to reach it. Spark registers its own |
| 37 | + * `ArrowCachedBatch` in `KryoSerializer.loadableSparkClasses`; Comet cannot add to that list, so |
| 38 | + * [[CometKryoRegistrator]] has to be set explicitly, and this suite is what proves it is |
| 39 | + * sufficient. |
| 40 | + * |
| 41 | + * This needs its own suite because `spark.serializer` and `spark.kryo.registrator` are read when |
| 42 | + * `SparkEnv` builds the serializer, so they cannot be changed per test. |
| 43 | + */ |
| 44 | +class CometInMemoryCacheKryoSuite extends CometTestBase { |
| 45 | + |
| 46 | + import testImplicits._ |
| 47 | + |
| 48 | + override protected def beforeAll(): Unit = { |
| 49 | + CometInMemoryRelationHelper.clearSerializer() |
| 50 | + super.beforeAll() |
| 51 | + } |
| 52 | + |
| 53 | + override protected def afterAll(): Unit = { |
| 54 | + try { |
| 55 | + super.afterAll() |
| 56 | + } finally { |
| 57 | + CometInMemoryRelationHelper.clearSerializer() |
| 58 | + } |
| 59 | + } |
| 60 | + |
| 61 | + override protected def sparkConf: SparkConf = { |
| 62 | + val conf = super.sparkConf |
| 63 | + conf.set("spark.plugins", "org.apache.spark.CometPlugin") |
| 64 | + conf.set( |
| 65 | + "spark.sql.cache.serializer", |
| 66 | + "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") |
| 67 | + conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer") |
| 68 | + conf.set("spark.kryo.registrationRequired", "true") |
| 69 | + conf.set("spark.kryo.registrator", CometKryoRegistrator.CLASS_NAME) |
| 70 | + conf |
| 71 | + } |
| 72 | + |
| 73 | + private def cachedBatchTypes(table: String): Array[String] = { |
| 74 | + val cached = spark.sharedState.cacheManager.lookupCachedData(spark.table(table)).get |
| 75 | + cached.cachedRepresentation.cacheBuilder.cachedColumnBuffers |
| 76 | + .map(_.getClass.getName) |
| 77 | + .distinct() |
| 78 | + .collect() |
| 79 | + } |
| 80 | + |
| 81 | + // Every type whose bounds gatherColumnStats records, so the statistics row carries one of each |
| 82 | + // internal representation Kryo has to write: boxed primitives, UTF8String, and Decimal at both |
| 83 | + // sides of the long/BigDecimal split. |
| 84 | + private val statsColumns = Seq( |
| 85 | + "id AS c_long", |
| 86 | + "cast(id % 2 = 0 as boolean) AS c_bool", |
| 87 | + "cast(id % 100 as byte) AS c_byte", |
| 88 | + "cast(id % 100 as short) AS c_short", |
| 89 | + "cast(id as int) AS c_int", |
| 90 | + "cast(id as float) AS c_float", |
| 91 | + "cast(id as double) AS c_double", |
| 92 | + "cast(id as decimal(9,2)) AS c_dec_short", |
| 93 | + "cast(id as decimal(30,4)) AS c_dec_long", |
| 94 | + "cast(id as string) AS c_string", |
| 95 | + "cast(date '2020-01-01' + cast(id as int) as date) AS c_date", |
| 96 | + "timestamp '2020-01-01 00:00:00' + make_interval(0, 0, 0, 0, 0, 0, id) AS c_ts") |
| 97 | + |
| 98 | + // DISK_ONLY and MEMORY_AND_DISK_SER both serialize the block on put, so each one reaches Kryo |
| 99 | + // deterministically in local mode. The default MEMORY_AND_DISK reaches it only once a partition |
| 100 | + // spills, which is the case that makes this more than a DISK_ONLY concern but is not something a |
| 101 | + // test can force cheaply. |
| 102 | + Seq(StorageLevel.DISK_ONLY, StorageLevel.MEMORY_AND_DISK_SER) |
| 103 | + .foreach { level => |
| 104 | + test(s"Comet in-memory cache round-trips through Kryo at $level") { |
| 105 | + withSQLConf( |
| 106 | + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", |
| 107 | + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", |
| 108 | + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", |
| 109 | + "spark.comet.sparkToColumnar.enabled" -> "true") { |
| 110 | + |
| 111 | + spark.catalog.clearCache() |
| 112 | + try { |
| 113 | + spark |
| 114 | + .range(0, 200, 1, 4) |
| 115 | + .selectExpr(statsColumns: _*) |
| 116 | + .createOrReplaceTempView("kryo_cache") |
| 117 | + |
| 118 | + spark.catalog.cacheTable("kryo_cache", level) |
| 119 | + assert(spark.table("kryo_cache").count() == 200) |
| 120 | + |
| 121 | + assert( |
| 122 | + cachedBatchTypes("kryo_cache").sameElements( |
| 123 | + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch")), |
| 124 | + "the payload Kryo serialized must be Comet's cached batch format") |
| 125 | + |
| 126 | + // Read the payload back rather than only the row count, so a Kryo round trip that |
| 127 | + // silently mangles the Arrow bytes fails too. The predicate also exercises the |
| 128 | + // statistics row, which is what carries UTF8String and Decimal through Kryo. |
| 129 | + checkSparkAnswer( |
| 130 | + spark.sql("SELECT c_long, c_string, c_dec_long, c_ts FROM kryo_cache " + |
| 131 | + "WHERE c_dec_short >= 100 AND c_string > '1'")) |
| 132 | + } finally { |
| 133 | + spark.catalog.clearCache() |
| 134 | + } |
| 135 | + } |
| 136 | + } |
| 137 | + } |
| 138 | + |
| 139 | + test("Comet broadcast exchange survives Kryo with registration required") { |
| 140 | + // Not about the cache: CometBroadcastExchangeExec broadcasts an Array[ChunkedByteBuffer], and |
| 141 | + // Spark registers ChunkedByteBuffer but not an array of them, so this fails on main today |
| 142 | + // under registrationRequired=true. The registrator this suite installs covers it because the |
| 143 | + // cache write path hands back the same type. Kept here rather than split out because that |
| 144 | + // registration is the thing under test. |
| 145 | + withSQLConf( |
| 146 | + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", |
| 147 | + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", |
| 148 | + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", |
| 149 | + CometConf.COMET_EXEC_BROADCAST_EXCHANGE_ENABLED.key -> "true", |
| 150 | + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") { |
| 151 | + withParquetTable((0 until 100).map(i => (i, i.toString)), "kryo_bcast_a") { |
| 152 | + withParquetTable((0 until 10).map(i => (i, i.toString)), "kryo_bcast_b") { |
| 153 | + val df = spark.sql( |
| 154 | + "SELECT /*+ BROADCAST(b) */ a._1, b._2 " + |
| 155 | + "FROM kryo_bcast_a a JOIN kryo_bcast_b b ON a._1 = b._1") |
| 156 | + assert( |
| 157 | + df.queryExecution.executedPlan.toString().contains("CometBroadcastExchange"), |
| 158 | + "the broadcast has to run through Comet for this to test anything") |
| 159 | + checkSparkAnswer(df) |
| 160 | + } |
| 161 | + } |
| 162 | + } |
| 163 | + } |
| 164 | + |
| 165 | + test("Comet in-memory cache falls back to Spark's format under Kryo for unsupported types") { |
| 166 | + // A relation Comet cannot store is delegated to DefaultCachedBatch, which Spark registers |
| 167 | + // itself. Pins that the fallback path is not collateral damage of the registration work. |
| 168 | + withSQLConf( |
| 169 | + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", |
| 170 | + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { |
| 171 | + |
| 172 | + spark.catalog.clearCache() |
| 173 | + try { |
| 174 | + spark |
| 175 | + .range(0, 100, 1, 2) |
| 176 | + .selectExpr("id", "make_interval(0, 0, 0, 0, 0, 0, id) AS iv") |
| 177 | + .createOrReplaceTempView("kryo_cache_fallback") |
| 178 | + |
| 179 | + spark.catalog.cacheTable("kryo_cache_fallback", StorageLevel.DISK_ONLY) |
| 180 | + assert(spark.table("kryo_cache_fallback").count() == 100) |
| 181 | + assert( |
| 182 | + cachedBatchTypes("kryo_cache_fallback").sameElements( |
| 183 | + Array("org.apache.spark.sql.execution.columnar.DefaultCachedBatch"))) |
| 184 | + |
| 185 | + checkSparkAnswer(spark.sql("SELECT id FROM kryo_cache_fallback WHERE id > 90")) |
| 186 | + } finally { |
| 187 | + spark.catalog.clearCache() |
| 188 | + } |
| 189 | + } |
| 190 | + } |
| 191 | +} |
0 commit comments