diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md index 20091ac317..044fa0307c 100644 --- a/docs/source/user-guide/latest/in-memory-cache.md +++ b/docs/source/user-guide/latest/in-memory-cache.md @@ -68,6 +68,14 @@ A relation whose cached plan records observed metrics, from `Dataset.observe`, i Comet's format but is scanned by Spark's `InMemoryTableScanExec`, because Spark collects those metrics only through that scan. +Set `spark.comet.exec.inMemoryCache.deltaEncoding.enabled=true` to try delta encoding +for top-level `bigint` columns in compressed caches. A column uses deltas only when +its compressed data buffer is over 25% smaller than the plain representation. +This can reduce footprint for sequential values, but adds work and +a temporary data buffer during cache creation, and a prefix sum during reads. +It is off by default and is skipped with `compression.codec=none`. The setting +only affects newly cached batches; changing it leaves existing caches readable. + ## Storage format Each cached batch is stored as a single Arrow IPC record batch message and its body. @@ -113,6 +121,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.deltaEncoding.enabled` | `false` | Try delta encoding top-level long columns in compressed caches. Affects newly cached data only. | | `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. | diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index f2f19ac5ce..db21d9c447 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -361,6 +361,17 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_EXEC_IN_MEMORY_CACHE_DELTA_ENCODING_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.exec.inMemoryCache.deltaEncoding.enabled") + .category(CATEGORY_EXEC) + .doc("Store suitable long columns as deltas before Arrow buffer compression. " + + "Only used with compressed caches and only when the delta buffer is over 25% smaller. " + + "Can reduce cache footprint for sequential values, but adds cache-build " + + "work and a prefix sum on reads. Only affects newly cached data; existing batches " + + "remain readable when this setting changes.") + .booleanConf + .createWithDefault(false) + val COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC: ConfigEntry[String] = conf("spark.comet.exec.inMemoryCache.compression.codec") .category(CATEGORY_EXEC) 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 2771dc8d2c..cf35162488 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 @@ -51,7 +51,8 @@ import org.apache.comet.vector.NativeUtil * lets a scan decompress only the columns it projected: the message records every buffer's offset * and length, so `CachedBatchIpc.Projection.load` copies out just the selected columns' byte * ranges. The cache manager still owns storage and eviction; this class only changes the cached - * payload. + * payload. `deltaEncoded` records which top-level long data buffers contain deltas; readers + * restore the values before returning vectors, independently of the current write settings. * * `sizeInBytes` is not the payload's size. It is inherited from `SimpleMetricsCachedBatch`, which * sums the decoded per-column sizes in `stats`. `statsRow` explains why they are decoded sizes. @@ -59,7 +60,8 @@ import org.apache.comet.vector.NativeUtil private case class CometCachedBatch( override val numRows: Int, override val stats: InternalRow, - bytes: ChunkedByteBuffer) + bytes: ChunkedByteBuffer, + deltaEncoded: Array[Boolean] = Array.emptyBooleanArray) extends SimpleMetricsCachedBatch /** @@ -70,7 +72,11 @@ private case class CometCachedBatch( * executor `CometConf` would resolve against whatever `SQLConf` happens to be current on that * thread rather than against this session's. */ -private case class CacheWriteSettings(codecName: String, zstdLevel: Int, chunkSize: Int) +private case class CacheWriteSettings( + codecName: String, + zstdLevel: Int, + chunkSize: Int, + deltaEncoding: Boolean) /** * Cache serializer that stores Comet-compatible Arrow batches in Spark's in-memory cache. @@ -385,7 +391,8 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { CacheWriteSettings( CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC.get(conf), CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_ZSTD_LEVEL.get(conf), - CometConf.COMET_EXEC_IN_MEMORY_CACHE_CHUNK_SIZE.get(conf).toInt) + CometConf.COMET_EXEC_IN_MEMORY_CACHE_CHUNK_SIZE.get(conf).toInt, + CometConf.COMET_EXEC_IN_MEMORY_CACHE_DELTA_ENCODING_ENABLED.get(conf)) // Serialize each batch to Arrow, gathering the Spark-compatible cache stats first. The stats are // stored beside the Arrow bytes so Spark's cache filter can prune a CometCachedBatch without @@ -417,19 +424,31 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { val writeDirectly = Utils.isArrowBacked(batch) && CachedBatchIpc.matchesReaderLayout(batch, readerFields) - val (bytes, columnSizes) = if (writeDirectly) { - CachedBatchIpc.serialize(batch, codec, CometArrowAllocator, settings.chunkSize) + val (bytes, columnSizes, deltaEncoded) = if (writeDirectly) { + CachedBatchIpc.serialize( + batch, + codec, + CometArrowAllocator, + settings.chunkSize, + settings.deltaEncoding) } else { val arrowBatch = CometArrowConverters.columnarBatchToArrowBatch(batch, arrowSchema, CometArrowAllocator) - try CachedBatchIpc.serialize(arrowBatch, codec, CometArrowAllocator, settings.chunkSize) + try + CachedBatchIpc.serialize( + arrowBatch, + codec, + CometArrowAllocator, + settings.chunkSize, + settings.deltaEncoding) finally arrowBatch.close() } CometCachedBatch( numRows = numRows, stats = statsRow(lower, upper, nulls, numRows, columnSizes), - bytes = bytes) + bytes = bytes, + deltaEncoded = deltaEncoded) } } @@ -609,7 +628,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // Decoding happens during construction, so `batches` below can hand out the root directly. // `load` releases everything it allocated if it throws, so there is nothing to unwind here. - private val root = projection.load(cached.bytes, CometArrowAllocator) + private val root = projection.load(cached.bytes, CometArrowAllocator, cached.deltaEncoded) private var closed = false // A cached batch's columns all cover the same rows. Check rather than trust: a mismatch would diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala index 06c3e52759..3ae2a51b96 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala @@ -29,7 +29,7 @@ import scala.util.control.NonFatal import org.apache.arrow.compression.{CommonsCompressionFactory, ZstdCompressionCodec} import org.apache.arrow.flatbuf.{RecordBatch => FlatBufRecordBatch} import org.apache.arrow.memory.{ArrowBuf, BufferAllocator} -import org.apache.arrow.vector.{FieldVector, TypeLayout, ValueVector, VectorLoader, VectorSchemaRoot, VectorUnloader} +import org.apache.arrow.vector.{BigIntVector, FieldVector, TypeLayout, ValueVector, VectorLoader, VectorSchemaRoot, VectorUnloader} import org.apache.arrow.vector.compression.{CompressionCodec, CompressionUtil, NoCompressionCodec} import org.apache.arrow.vector.dictionary.DictionaryEncoder import org.apache.arrow.vector.ipc.{ReadChannel, WriteChannel} @@ -186,10 +186,13 @@ private[comet] object CachedBatchIpc { * batch arrives at whatever size the plan above produced. Chunks are appended rather than grown * and recopied, so the write also never holds the payload twice. * - * Returns the message and the decoded size of each top-level column. Each size is measured on - * the batch before compression, from the plain lengths of the column's own buffers, which is - * what `getBufferSize` reports for a vector. The caller records the sizes in the statistics - * row; `ArrowCachedBatchSerializer.statsRow` explains why they are decoded sizes. + * Returns the message, the decoded size of each top-level column, and flags for columns stored + * as deltas. Delta encoding touches only a top-level long column's data buffer; validity and + * statistics remain logical, and Projection restores values before exposing the root. Each size + * is measured on the batch before compression, from the plain lengths of the column's own + * buffers, which is what `getBufferSize` reports for a vector. The caller records the sizes in + * the statistics row; `ArrowCachedBatchSerializer.statsRow` explains why they are decoded + * sizes. * * Dictionary-encoded columns are decoded to their plain form first. A payload with no Schema * message cannot describe a dictionary encoding, and the schema the reader rebuilds from Spark @@ -204,7 +207,8 @@ private[comet] object CachedBatchIpc { batch: ColumnarBatch, codec: CompressionCodec, allocator: BufferAllocator, - chunkSize: Int): (ChunkedByteBuffer, Array[Long]) = { + chunkSize: Int, + deltaEncoding: Boolean = false): (ChunkedByteBuffer, Array[Long], Array[Boolean]) = { val (vectors, decoded) = decodeDictionaries(batch, allocator) try { val root = new VectorSchemaRoot(vectors.asJava) @@ -219,9 +223,20 @@ private[comet] object CachedBatchIpc { val fields = vectors.map(_.getField) val unloader = new VectorUnloader(root, true, NoCompressionCodec.INSTANCE, true) val plainBatch = unloader.getRecordBatch + val deltaEncoded = + if (deltaEncoding) new Array[Boolean](vectors.length) else Array.emptyBooleanArray + val deltaColumns = + if (deltaEncoding && codec.getCodecType != CompressionUtil.CodecType.NO_COMPRESSION) { + val starts = fields.scanLeft(0)(_ + fieldBufferCount(_)).toArray + vectors.indices.collect { + case i if vectors(i).isInstanceOf[BigIntVector] => (starts(i) + 1) -> i + }.toMap + } else Map.empty[Int, Int] val (sizes, recordBatch) = try { - (columnSizes(fields, plainBatch), compressed(plainBatch, codec, allocator)) + ( + columnSizes(fields, plainBatch), + compressed(plainBatch, codec, allocator, deltaColumns, deltaEncoded)) } finally { plainBatch.close() } @@ -237,7 +252,10 @@ private[comet] object CachedBatchIpc { } finally { out.close() } - (out.toChunkedByteBuffer, sizes) + ( + out.toChunkedByteBuffer, + sizes, + if (deltaEncoded.contains(true)) deltaEncoded else Array.emptyBooleanArray) } finally { recordBatch.close() } @@ -283,7 +301,13 @@ private[comet] object CachedBatchIpc { * selected are never read, let alone inflated. The windows are then decompressed in one pass; * see [[decompressed]] for why that is not left to `VectorLoader`. */ - def load(data: ChunkedByteBuffer, allocator: BufferAllocator): VectorSchemaRoot = { + def load( + data: ChunkedByteBuffer, + allocator: BufferAllocator, + deltaEncoded: Array[Boolean] = Array.emptyBooleanArray): VectorSchemaRoot = { + require( + deltaEncoded.isEmpty || deltaEncoded.length == arrowFields.length, + "Delta flags must match the cached schema") val readChannel = new ReadChannel(Channels.newChannel(data.toInputStream())) // Reads the message metadata only. The body stays in `data` and is copied selectively. val metadata = MessageSerializer.readMessage(readChannel) @@ -393,6 +417,25 @@ private[comet] object CachedBatchIpc { val root = VectorSchemaRoot.create(schema, allocator) try { new VectorLoader(root).load(plainBatch) + if (deltaEncoded.nonEmpty) { + selectedIndices.indices.foreach { i => + if (deltaEncoded(selectedIndices(i))) { + val vector = root.getVector(i) + require( + vector.isInstanceOf[BigIntVector], + "Delta-encoded cache column must contain longs") + val data = vector.getDataBuffer + var previous = 0L + var row = 0 + while (row < vector.getValueCount) { + val value = data.getLong(row * 8L) + previous + data.setLong(row * 8L, value) + previous = value + row += 1 + } + } + } + } root } catch { case NonFatal(e) => @@ -509,7 +552,9 @@ private[comet] object CachedBatchIpc { private def compressed( batch: ArrowRecordBatch, codec: CompressionCodec, - allocator: BufferAllocator): ArrowRecordBatch = { + allocator: BufferAllocator, + deltaColumns: Map[Int, Int], + deltaEncoded: Array[Boolean]): ArrowRecordBatch = { val buffers = new java.util.ArrayList[ArrowBuf](batch.getBuffers.size) try { batch.getBuffers.asScala.foreach { buffer => @@ -524,6 +569,45 @@ private[comet] object CachedBatchIpc { buffers.add(packed) } + deltaColumns.foreach { case (index, column) => + val buffer = batch.getBuffers.get(index) + val packed = buffers.get(index) + if (batch.getLength > 0) { + val deltas = allocator.buffer(buffer.writerIndex()) + try { + var previous = 0L + var smallDeltas = 0 + var row = 0 + while (row < batch.getLength) { + val value = buffer.getLong(row * 8L) + val delta = value - previous + deltas.setLong(row * 8L, delta) + if (delta == delta.toInt.toLong) smallDeltas += 1 + previous = value + row += 1 + } + deltas.writerIndex(buffer.writerIndex()) + // Skip a second compression for full-width irregular longs. The size comparison + // still rejects poorly compressing deltas from narrower distributions. + if (smallDeltas.toLong * 2 >= batch.getLength) { + deltas.getReferenceManager.retain() + val encoded = + try codec.compress(allocator, deltas) + catch { + case NonFatal(e) => + deltas.getReferenceManager.release() + throw e + } + if (encoded.writerIndex() < packed.writerIndex() * 3 / 4) { + buffers.set(index, encoded) + deltaEncoded(column) = true + packed.close() + } else encoded.close() + } + } finally deltas.close() + } + } + val result = new ArrowRecordBatch( batch.getLength, batch.getNodes, 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 05b80605db..ddcf5d35fa 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala @@ -23,7 +23,7 @@ import org.apache.spark.SparkConf import org.apache.spark.serializer.KryoRegistrator import org.apache.spark.sql.{CometTestBase, Row} import org.apache.spark.sql.catalyst.expressions.GenericInternalRow -import org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer +import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, DefaultCachedBatch} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.storage.StorageLevel @@ -109,6 +109,8 @@ class CometInMemoryCacheKryoSuite extends CometTestBase { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_DELTA_ENCODING_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC.key -> "zstd", CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { spark.catalog.clearCache() @@ -133,6 +135,15 @@ class CometInMemoryCacheKryoSuite extends CometTestBase { Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch")), "the payload Kryo serialized must be Comet's cached batch format") + val cached = spark.sharedState.cacheManager + .lookupCachedData(spark.table("kryo_cache")) + .get + assert( + cached.cachedRepresentation.cacheBuilder.cachedColumnBuffers + .map(b => CometCachedBatchHelper.columnsAreDeltaEncoded(b)(0)) + .collect() + .forall(identity)) + // Read the payload back rather than only the row count, so a Kryo round trip that // silently mangles the Arrow bytes fails too. The predicate also exercises the // statistics row, which is what carries UTF8String and Decimal through Kryo. 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..f29c1f2a25 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -29,7 +29,7 @@ import scala.jdk.CollectionConverters._ import org.apache.arrow.compression.ZstdCompressionCodec import org.apache.arrow.memory.{ArrowBuf, BufferAllocator} -import org.apache.arrow.vector.{BitVector, FixedSizeBinaryVector, IntVector, VarBinaryVector, VarCharVector} +import org.apache.arrow.vector.{BigIntVector, BitVector, FixedSizeBinaryVector, IntVector, VarBinaryVector, VarCharVector} import org.apache.arrow.vector.compression.{CompressionCodec, CompressionUtil, NoCompressionCodec} import org.apache.arrow.vector.types.pojo.ArrowType import org.apache.spark.CometDriverPlugin @@ -1226,6 +1226,8 @@ class CometInMemoryCacheSuite extends CometTestBase { SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", CometConf.COMET_SHUFFLE_MODE.key -> "jvm", SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_DELTA_ENCODING_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC.key -> "zstd", CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { spark.catalog.clearCache() @@ -1245,6 +1247,10 @@ class CometInMemoryCacheSuite extends CometTestBase { val cached = spark.sharedState.cacheManager.lookupCachedData(spark.table("disk_cache")).get + assert( + cached.cachedRepresentation.cacheBuilder.cachedColumnBuffers + .collect() + .forall(b => CometCachedBatchHelper.columnsAreDeltaEncoded(b)(0))) val rddId = cached.cachedRepresentation.cacheBuilder.cachedColumnBuffers.id val info = spark.sparkContext.getRDDStorageInfo .find(_.id == rddId) @@ -2701,6 +2707,115 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Comet in-memory cache skips delta encoding when disabled or uncompressed") { + for ((enabled, codec) <- Seq((false, "zstd"), (true, "none"))) { + withSQLConf( + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_DELTA_ENCODING_ENABLED.key -> enabled.toString, + CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC.key -> codec) { + val cached = spark.range(4096).toDF().cache() + try { + assert(cached.count() == 4096) + val relation = + spark.sharedState.cacheManager.lookupCachedData(cached).get.cachedRepresentation + val batches = relation.cacheBuilder.cachedColumnBuffers.collect() + assert( + batches.forall(b => !CometCachedBatchHelper.columnsAreDeltaEncoded(b).contains(true))) + checkAnswer(cached, (0L until 4096L).map(Row(_))) + } finally cached.unpersist() + } + } + } + + test("Comet in-memory cache releases buffers when delta compression fails") { + val longs = new BigIntVector("l", CometArrowAllocator) + try { + longs.allocateNew(4096) + (0 until 4096).foreach(i => longs.set(i, i.toLong)) + longs.setValueCount(4096) + val batch = new ColumnarBatch(Array[ColumnVector](new CometPlainVector(longs)), 4096) + // Validity and plain data succeed; delta compression fails with both the packed data + // and the scratch buffer live. Neither may escape the failure path. + val codec = new FailAfterCompressionCodec(succeedFor = 2) + val before = CometArrowAllocator.getAllocatedMemory + val thrown = intercept[Exception] { + CometCachedBatchHelper.serialize(batch, codec, CometArrowAllocator, deltaEncoding = true) + } + assert(causeChain(thrown).exists(t => + Option(t.getMessage).contains(FailAfterCompressionCodec.Message))) + assert(codec.compressed == 2) + assert(CometArrowAllocator.getAllocatedMemory == before) + } finally longs.close() + } + + test("Comet in-memory cache preserves delta-encoded longs across every reader") { + val random = new java.util.Random(5485) + val rows = (0 until 4096).map { i => + // Cross the signed-long boundary; delta reconstruction must preserve wrapping arithmetic. + val value = Long.MaxValue - 2048 + i + Row(value, if (i % 7 == 0) null else value, random.nextLong(), s"value_${i % 11}") + } + val schema = new StructType() + .add("seq", LongType, nullable = false) + .add("nullable", LongType, nullable = true) + .add("random", LongType, nullable = false) + .add("text", StringType, nullable = false) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_DELTA_ENCODING_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC.key -> "zstd", + "spark.comet.sparkToColumnar.enabled" -> "true") { + spark.catalog.clearCache() + val cached = spark.createDataFrame(spark.sparkContext.parallelize(rows, 2), schema).cache() + try { + assert(cached.count() == rows.length) + val relation = + spark.sharedState.cacheManager.lookupCachedData(cached).get.cachedRepresentation + val batches = relation.cacheBuilder.cachedColumnBuffers.collect() + assert(batches.forall(b => CometCachedBatchHelper.columnsAreDeltaEncoded(b)(0))) + assert(batches.forall(b => CometCachedBatchHelper.columnsAreDeltaEncoded(b)(1))) + assert(batches.forall(b => !CometCachedBatchHelper.columnsAreDeltaEncoded(b)(2))) + + val expected = + rows.map(r => Row(r.getLong(2), r.getLong(0), r.get(1), r.getLong(0), r.get(3))) + for ((native, vectorized) <- Seq((true, true), (false, true), (false, false))) { + withSQLConf( + CometConf.COMET_EXEC_ENABLED.key -> native.toString, + CometConf.COMET_EXEC_IN_MEMORY_CACHE_DELTA_ENCODING_ENABLED.key -> "false", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> vectorized.toString) { + val projected = + cached.selectExpr("random", "seq", "nullable", "seq AS repeated", "text") + checkAnswer(projected, expected) + val plan = projected.queryExecution.executedPlan + if (native) assert(plan.exists(_.isInstanceOf[CometInMemoryTableScanExec])) + else { + assert(plan.exists(_.isInstanceOf[InMemoryTableScanExec])) + assert(plan.exists(_.isInstanceOf[ColumnarToRowExec]) == vectorized) + } + } + } + + // Bounds and null counts stay logical even though the payload contains deltas. + checkAnswer( + cached.filter(cached.col("seq") === Long.MinValue), + rows.filter(_.getLong(0) == Long.MinValue)) + checkAnswer(cached.filter("nullable IS NULL"), rows.filter(_.isNullAt(1))) + + // Re-caching decoded vectors must neither encode them twice nor change the original cache. + withSQLConf( + CometConf.COMET_EXEC_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "false", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_DELTA_ENCODING_ENABLED.key -> "false") { + val recached = cached.union(cached).cache() + try checkAnswer(recached, rows ++ rows) + finally recached.unpersist() + checkAnswer(cached, rows) + } + } finally cached.unpersist() + } + } + test("Comet in-memory cache releases its vectors when a column fails to decode") { // Reading a batch allocates twice before anything can go wrong: the root that receives the // projected columns, and the off-heap body the selected buffers are copied into. A column diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala index 47e876b5f8..9920d5d40d 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometCachedBatchHelper.scala @@ -52,6 +52,9 @@ import org.apache.spark.util.io.{ChunkedByteBuffer, ChunkedByteBufferOutputStrea */ object CometCachedBatchHelper { + def columnsAreDeltaEncoded(batch: CachedBatch): Seq[Boolean] = + batch.asInstanceOf[CometCachedBatch].deltaEncoded.toSeq + /** The raw cached payload: one encapsulated Arrow IPC record batch message and its body. */ private def payload(batch: CachedBatch): ChunkedByteBuffer = batch.asInstanceOf[CometCachedBatch].bytes @@ -122,8 +125,9 @@ object CometCachedBatchHelper { batch: ColumnarBatch, codec: CompressionCodec, allocator: BufferAllocator, - chunkSize: Int = 1024 * 1024): ChunkedByteBuffer = - CachedBatchIpc.serialize(batch, codec, allocator, chunkSize)._1 + chunkSize: Int = 1024 * 1024, + deltaEncoding: Boolean = false): ChunkedByteBuffer = + CachedBatchIpc.serialize(batch, codec, allocator, chunkSize, deltaEncoding)._1 /** A payload [[serialize]] wrote, as the cached batch the writer would have stored it in. */ def cachedBatch(payload: ChunkedByteBuffer, numRows: Int): CachedBatch = @@ -139,7 +143,7 @@ object CometCachedBatchHelper { selected: Array[Int], allocator: BufferAllocator): VectorSchemaRoot = new CachedBatchIpc.Projection(arrowFields(cacheSchema).toIndexedSeq, selected) - .load(payload(batch), allocator) + .load(payload(batch), allocator, batch.asInstanceOf[CometCachedBatch].deltaEncoded) /** * A cached batch whose payload records `codec` as the byte that compressed its body.