diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md index d6867defeab..f0c2cd6a31e 100644 --- a/docs/source/user-guide/latest/in-memory-cache.md +++ b/docs/source/user-guide/latest/in-memory-cache.md @@ -91,36 +91,47 @@ column out of six does roughly a sixth of the decompression work, and a `SELECT selects no columns at all, answers from the row count stored beside the payload without touching it. -Compression defaults to `zstd`, for footprint rather than for speed. Over the same 5M-row, -six-column relation the tables under [Performance](#performance) use — and measured by the same -benchmark — it holds the data in a sixth of the memory and pays for that on both sides: about 40% -longer to materialize, and, on a read wide enough to inflate everything, close to five times -longer. A narrow projection pays far less, because it only inflates the columns it asked for. +Three codecs are offered: `zstd`, the default, `lz4` and `none`. Measured by the same benchmark +over the same 5M-row, six-column relation the tables under [Performance](#performance) use (Apple +M3 Ultra, JDK 17, Spark 4.1, release build): | Codec | Materialize | Footprint | Read 1 of 6 | Read 6 of 6 | | ------ | ----------: | --------: | ----------: | ----------: | -| `zstd` | 1507 ms | 55 MiB | 45 ms | 295 ms | -| `none` | 1081 ms | 315 MiB | 35 ms | 64 ms | +| `zstd` | 1232 ms | 51 MiB | 36 ms | 248 ms | +| `lz4` | 1115 ms | 145 MiB | 33 ms | 110 ms | +| `none` | 857 ms | 315 MiB | 31 ms | 54 ms | + +`zstd` is the default, for footprint rather than for speed. It holds the data in a sixth of the +memory `none` takes and pays for that on both sides: about 45% longer to materialize, and, on a read +wide enough to inflate everything, more than four times longer. A narrow projection pays far less, +because it only inflates the columns it asked for. + +`lz4` decompresses several times faster than `zstd`, so a read of every column takes less than half +as long, at nearly three times the footprint. It is written with lz4-java, the JNI-accelerated +library behind `spark.io.compression.codec`, in Arrow's standard `LZ4_FRAME` format. Arrow's own +LZ4 codec is not used: it is commons-compress's pure-Java implementation, and is orders of +magnitude slower to write than `zstd` while also producing larger output. + +What `lz4` saves depends more on the data than what `zstd` saves does. It holds this relation in +less memory than Spark's own cache format does (217 MiB), but not a relation of sequential `bigint` +columns, which Spark's format delta-encodes: each of the wide relations under +[Limitations](#limitations) takes about 74 MiB with `lz4`, against 26 to 27 MiB in Spark's format +and with `zstd`. `none` is the better setting for a relation that fits in memory uncompressed and is read at close -to full width. The default is the other way round because a cache that does not fit costs more than -one that is slower to read, and Spark's own cache format compresses by default too. - -Arrow's other IPC codec, LZ4, is deliberately not offered and the config rejects it. It is -commons-compress's pure-Java implementation, unrelated to the JNI-accelerated lz4-java behind -`spark.io.compression.codec`, and is orders of magnitude slower to write than `zstd` while also -producing larger output. +to full width. The default is `zstd` because a cache that does not fit costs more than one that is +slower to read, and Spark's own cache format compresses by default too. Dictionary-encoded columns are decoded before they are stored. A payload with no schema message has nowhere to record either that a column is dictionary encoded or the dictionary itself. ## Configuration -| Config | Default | Description | -| ------------------------------------------------------- | ------- | ---------------------------------------------------------------------------------------------------------------------------------------------- | -| `spark.comet.exec.inMemoryCache.enabled` | `true` | Whether to store and scan Spark's in-memory cache in Comet's format. Read at startup. Defaults to `false` on Spark 3.4. | -| `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. | +| Config | Default | Description | +| ------------------------------------------------------- | ------- | ---------------------------------------------------------------------------------------------------------------------------------------------------- | +| `spark.comet.exec.inMemoryCache.enabled` | `true` | Whether to store and scan Spark's in-memory cache in Comet's format. Read at startup. Defaults to `false` on Spark 3.4. | +| `spark.comet.exec.inMemoryCache.compression.codec` | `zstd` | Arrow IPC compression codec for cached data: `zstd`, `lz4` 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. | ## Performance @@ -205,6 +216,12 @@ all of them longs, is about 10% slower under either kind of operator. That cost 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. +Measured on an Apple M3 Ultra, Comet's format with `lz4` is faster than Spark's in every shape +under both kinds of operator. The three-column read is 1.4x faster than Spark's format with either +kind of operator above the scan, where `zstd` is about 15% slower on the same machine, and the full +projection is 2.5x faster with Comet operators above the scan and 1.7x faster with a Spark operator +above it. + Building the cache is measured the same way, from the same source, with its rows produced either by Comet operators or by Spark operators. Comet's format is written straight from the Arrow batches of Comet operators, while Spark's is always built from rows, so above Comet operators it first converts @@ -292,6 +309,13 @@ the others. The row reader takes up to 1.6 times as long, and it is the only rea 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. A Spark operator above Comet's native cache scan does not pay this; see [Performance](#performance). + +Most of what the row reader pays is decompression, so `lz4` narrows these gaps. Measured on an +Apple M3 Ultra, the row reader takes up to 1.25 times as long as Spark's format with `lz4`, against +up to 1.7 times with `zstd`, and the fused reader is faster than Spark's format for every read. +Reading every column of the wide relations takes 1.3 to 2.0 times as long as from Spark's format +with `lz4`, against 2.1 to 2.6 times with `zstd`. `lz4` takes more memory than Spark's format for +those relations, though; see [Storage format](#storage-format). [#5485](https://github.com/apache/datafusion-comet/issues/5485) tracks these gaps. Comet's serializer exists because Spark's own Arrow cache format diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 4aab7c72a11..2beaf8f6a3c 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -388,7 +388,7 @@ object CometConf extends ShimCometConf { "without the fused reader. Relations whose schema Comet's Arrow writer does not " + "support are always cached in Spark's default " + "format. Each cached batch is stored as one Arrow IPC record batch with per-buffer " + - "zstd compression, and a scan copies out only the buffers of the columns it " + + "compression, and a scan copies out only the buffers of the columns it " + "projected, so the unselected ones are never decompressed. Eligible Spark " + "whole-stage codegen consumers read cached vectors directly when vectorized cache " + "reading is enabled; other Spark row consumers use a reusable row buffer. Decoding " + @@ -409,20 +409,20 @@ object CometConf extends ShimCometConf { .category(CATEGORY_EXEC) .doc( "The Arrow IPC compression codec used when Comet's cache serializer writes cached " + - "data. Unlike spark.io.compression.codec, this compresses each Arrow buffer " + - "separately rather than the batch as a whole, which is what lets a projected scan " + - "decompress only the columns it selected. zstd is the default for footprint rather " + - "than speed: it stores cached data in a fraction of the memory, but makes " + - "materializing slower, and a read slower in proportion to how many columns it " + - "selects. Set to none when a relation fits in memory uncompressed and is read at " + - "close to full width. " + + "data: zstd, lz4 or none. Unlike spark.io.compression.codec, this compresses each " + + "Arrow buffer separately rather than the batch as a whole, which is what lets a " + + "projected scan decompress only the columns it selected. zstd is the default for " + + "footprint rather than speed: it stores cached data in a fraction of the memory, but " + + "makes materializing slower, and a read slower in proportion to how many columns it " + + "selects. lz4 decompresses several times faster than zstd, which makes wide reads " + + "much faster, at around three times zstd's footprint. It uses lz4-java, the " + + "JNI-accelerated library behind spark.io.compression.codec, rather than Arrow's own " + + "pure-Java lz4. Set to none when a relation fits in memory uncompressed and is read " + + "at close to full width. " + "Only affects newly cached data; the codec a batch was written with is recorded in " + - "the batch itself and is what the read path uses. Arrow's lz4 is deliberately not " + - "offered: it is a pure-Java implementation, unrelated to the JNI-accelerated lz4 " + - "behind spark.io.compression.codec, and is orders of magnitude slower to write than " + - "zstd while also producing larger output.") + "the batch itself and is what the read path uses.") .stringConf - .checkValues(Set("none", "zstd")) + .checkValues(Set("none", "zstd", "lz4")) .createWithDefault("zstd") val COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_ZSTD_LEVEL: ConfigEntry[Int] = 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 06c3e527592..16dc6813899 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 @@ -73,24 +73,28 @@ private[comet] object CachedBatchIpc { // Constructed directly rather than through CompressionCodec.Factory, which ignores the level // and always builds a codec at zstd's default. case "zstd" => new ZstdCompressionCodec(zstdLevel) - // Arrow's other codec, LZ4_FRAME, is not offered. It is commons-compress's pure-Java LZ4 -- - // no relation to the JNI-accelerated lz4-java behind spark.io.compression.codec -- and - // measures three orders of magnitude slower to write than zstd while also producing larger - // output, so nothing prefers it. Reads still accept it, since the factory the read path uses - // handles whatever codec a batch records. + // Arrow's LZ4_FRAME format, written with lz4-java rather than Arrow's own pure-Java codec; + // see Lz4FrameCompressionCodec. + case "lz4" => Lz4FrameCompressionCodec case other => throw new SparkException( s"Unsupported Arrow compression codec for Comet's cache: $other. " + - "Supported values: none, zstd") + "Supported values: none, zstd, lz4") } // Decompressors are stateless and shared. Resolving one per cached batch would allocate a codec - // per batch on every scan, and the enum lookup walks the CodecType values each time. + // per batch on every scan, and the enum lookup walks the CodecType values each time. LZ4_FRAME + // is read with the codec that writes it rather than with Arrow's pure-Java one, which inflates + // through heap copies and streams. private val readCodecs: Map[CompressionUtil.CodecType, CompressionCodec] = CompressionUtil.CodecType .values() .filter(_ != CompressionUtil.CodecType.NO_COMPRESSION) - .map(t => t -> CommonsCompressionFactory.INSTANCE.createCodec(t)) + .map { + case CompressionUtil.CodecType.LZ4_FRAME => + CompressionUtil.CodecType.LZ4_FRAME -> Lz4FrameCompressionCodec + case t => t -> CommonsCompressionFactory.INSTANCE.createCodec(t) + } .toMap /** diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/Lz4FrameCompressionCodec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/Lz4FrameCompressionCodec.scala new file mode 100644 index 00000000000..b0cc91501b4 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/Lz4FrameCompressionCodec.scala @@ -0,0 +1,209 @@ +/* + * 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.spark.sql.comet.execution.arrow + +import scala.util.control.NonFatal + +import org.apache.arrow.memory.{ArrowBuf, BufferAllocator} +import org.apache.arrow.vector.compression.{AbstractCompressionCodec, CompressionUtil} +import org.apache.spark.SparkException +import org.apache.spark.internal.Logging + +import net.jpountz.lz4.{LZ4Compressor, LZ4Exception, LZ4Factory, LZ4SafeDecompressor} + +/** + * Arrow's LZ4_FRAME compression codec, on lz4-java. + * + * Arrow ships an LZ4_FRAME codec of its own, but it is commons-compress's pure-Java LZ4, which + * writes orders of magnitude more slowly than zstd and produces larger output. This writes the + * same format with lz4-java, the JNI-accelerated library behind `spark.io.compression.codec`: + * each buffer is one standard LZ4 frame behind Arrow's 8-byte uncompressed-length prefix, which + * `AbstractCompressionCodec` reads and writes, so any Arrow reader can decode it. lz4-java comes + * from Spark's own classpath, as zstd-jni does for Arrow's zstd codec. + * + * Every frame has the same header: independent blocks of up to 4 MiB, and no checksums, content + * size or dictionary. The reader accepts no other, since it only ever reads frames this codec + * wrote. It checks every length the frame records before reading through it, because Comet turns + * off Arrow's own bounds checks (`arrow.enable_unsafe_memory_access`), so a corrupt length would + * otherwise read or write past the buffers. + * + * Lengths are written in the platform's byte order, as the rest of the payload is, which is the + * little-endian order the format requires on every platform Comet runs on. + */ +private[comet] object Lz4FrameCompressionCodec extends AbstractCompressionCodec with Logging { + + private val MaxBlockSize = 4 * 1024 * 1024 + + // The magic number, then the frame descriptor: FLG (version 01, independent blocks), BD (a + // 4 MiB maximum block size) and HC, the second byte of the xxHash32 of FLG and BD with a zero + // seed. HC is a constant because FLG and BD are. + private val Magic = 0x184d2204 + private val Flg: Byte = 0x60 + private val Bd: Byte = 0x70 + private val Hc: Byte = 0x73 + private val HeaderSize = 7 + + // Each block is led by its length, whose high bit marks a block stored as it was because + // compressing it would not have made it smaller. A zero length ends the frame. + private val LengthSize = 4 + private val StoredBlock = 0x80000000 + + private val Prefix = CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH + + // fastestInstance() takes lz4-java's JNI implementation only when lz4-java was loaded by the + // system class loader, and silently settles for its Java one otherwise, which is the case under + // `mvn exec:java`. Ask for the JNI one directly, and use Java only where it cannot load. + private lazy val factory: LZ4Factory = + try LZ4Factory.nativeInstance() + catch { + case e @ (NonFatal(_) | _: LinkageError) => + logWarning( + "Could not load lz4-java's native library, so Comet's in-memory cache uses its " + + "slower Java implementation of lz4", + e) + LZ4Factory.fastestJavaInstance() + } + private lazy val compressor: LZ4Compressor = factory.fastCompressor() + private lazy val decompressor: LZ4SafeDecompressor = factory.safeDecompressor() + + override def getCodecType: CompressionUtil.CodecType = CompressionUtil.CodecType.LZ4_FRAME + + override protected def doCompress( + allocator: BufferAllocator, + uncompressed: ArrowBuf): ArrowBuf = { + val length = uncompressed.writerIndex() + val fullBlocks = length / MaxBlockSize + val lastBlock = (length % MaxBlockSize).toInt + val bound = Prefix + HeaderSize + + fullBlocks * (LengthSize + compressor.maxCompressedLength(MaxBlockSize)) + + (if (lastBlock > 0) LengthSize + compressor.maxCompressedLength(lastBlock) else 0) + + LengthSize + val frame = allocator.buffer(bound) + try { + frame.setInt(Prefix, Magic) + frame.setByte(Prefix + 4, Flg) + frame.setByte(Prefix + 5, Bd) + frame.setByte(Prefix + 6, Hc) + var at = Prefix + HeaderSize + var from = 0L + while (from < length) { + val n = math.min(MaxBlockSize.toLong, length - from).toInt + val room = compressor.maxCompressedLength(n) + val packed = compressor.compress( + uncompressed.nioBuffer(from, n), + 0, + n, + frame.nioBuffer(at + LengthSize, room), + 0, + room) + if (packed < n) { + frame.setInt(at, packed) + at += LengthSize + packed + } else { + frame.setInt(at, n | StoredBlock) + frame.setBytes(at + LengthSize, uncompressed, from, n.toLong) + at += LengthSize + n + } + from += n + } + frame.setInt(at, 0) + frame.writerIndex(at + LengthSize) + frame + } catch { + case NonFatal(e) => + frame.close() + throw e + } + } + + override protected def doDecompress( + allocator: BufferAllocator, + compressed: ArrowBuf): ArrowBuf = { + val length = readUncompressedLength(compressed) + val end = compressed.writerIndex() + if (end < Prefix + HeaderSize || + compressed.getInt(Prefix) != Magic || + compressed.getByte(Prefix + 4) != Flg || + compressed.getByte(Prefix + 5) != Bd || + compressed.getByte(Prefix + 6) != Hc) { + throw corrupt("the buffer does not begin with the frame header Comet writes") + } + + val plain = allocator.buffer(length) + try { + var at = Prefix + HeaderSize + var to = 0L + var ended = false + while (!ended) { + if (end - at < LengthSize) { + throw corrupt("the frame ends before its end mark") + } + val word = compressed.getInt(at) + at += LengthSize + if (word == 0) { + ended = true + } else { + val n = word & ~StoredBlock + if (n > end - at) { + throw corrupt(s"a block of $n bytes runs past the end of the frame") + } + if ((word & StoredBlock) != 0) { + if (n > length - to) { + throw corrupt(s"the frame holds more than its recorded $length bytes") + } + plain.setBytes(to, compressed, at, n.toLong) + to += n + } else { + // A block never decodes to more than the maximum block size, and must not decode + // past the length the frame records. + val room = math.min(length - to, MaxBlockSize.toLong).toInt + val decoded = + try { + decompressor.decompress( + compressed.nioBuffer(at, n), + 0, + n, + plain.nioBuffer(to, room), + 0, + room) + } catch { + case e: LZ4Exception => + throw corrupt(s"a block does not decode within $room bytes", e) + } + to += decoded + } + at += n + } + } + if (to != length) { + throw corrupt(s"the frame decodes to $to bytes rather than its recorded $length") + } + plain.writerIndex(length) + plain + } catch { + case NonFatal(e) => + plain.close() + throw e + } + } + + private def corrupt(reason: String, cause: Throwable = null): SparkException = + new SparkException(s"Comet cached batch holds a corrupt LZ4 frame: $reason", cause) +} 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 aba97d80219..a2943675018 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -20,6 +20,7 @@ package org.apache.comet.exec import java.{util => ju} +import java.io.ByteArrayInputStream import java.nio.charset.StandardCharsets import scala.concurrent.{Await, Future} @@ -32,8 +33,8 @@ import org.apache.arrow.memory.{ArrowBuf, BufferAllocator} import org.apache.arrow.vector.{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.{SparkConf, SparkException} import org.apache.spark.CometDriverPlugin -import org.apache.spark.SparkConf import org.apache.spark.sql.{CometTestBase, DataFrame, Observation, QueryTest, Row} import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeReference, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, LessThan, Literal} import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} @@ -57,6 +58,8 @@ import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plu import org.apache.comet.rules.CometCacheColumnarRule import org.apache.comet.vector.{CometPlainVector, CometVector} +import net.jpountz.lz4.LZ4FrameInputStream + class CometInMemoryCacheSuite extends CometTestBase { import testImplicits._ @@ -2129,7 +2132,7 @@ class CometInMemoryCacheSuite extends CometTestBase { // Every codec the config accepts, not just the default. `none` takes a different path on read // -- the payload records no codec, so nothing is decompressed -- and shipped broken for a // while because the only tests that ran were on the default codec. - Seq("none", "zstd").foreach { codec => + Seq("none", "zstd", "lz4").foreach { codec => withConversions( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", @@ -2176,6 +2179,103 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("Comet in-memory cache writes lz4 buffers as standard LZ4 frames") { + // Arrow's LZ4_FRAME codec is the standard LZ4 frame format, so check what the codec writes + // against lz4-java's own frame reader rather than against the codec that wrote it. The inputs + // span the codec's 4 MiB block: a buffer of one block, and one of three whose first block is + // incompressible, which a frame stores verbatim. + val random = new scala.util.Random(7) + def compressible(n: Int): Array[Byte] = Array.tabulate[Byte](n)(i => (i / 16).toByte) + val inputs = Seq( + "one block" -> compressible(100 * 1000), + "three blocks" -> + (Array.fill[Byte](4 << 20)(random.nextInt().toByte) ++ compressible(5 << 20))) + val codec = CometCachedBatchHelper.compressionCodec("lz4") + assert(codec.getCodecType == CompressionUtil.CodecType.LZ4_FRAME) + + val allocator = CometArrowAllocator.newChildAllocator("lz4-frames", 0, Long.MaxValue) + try { + inputs.foreach { case (label, data) => + val plain = allocator.buffer(data.length.toLong) + plain.setBytes(0, data) + plain.writerIndex(data.length.toLong) + // Takes ownership of `plain`. + val packed = codec.compress(allocator, plain) + try { + // Arrow prefixes a compressed buffer with its uncompressed length, and a buffer it + // stored verbatim with -1. + assert(packed.getLong(0) == data.length, s"$label was stored verbatim") + val frame = new Array[Byte]((packed.writerIndex() - 8).toInt) + packed.getBytes(8, frame) + val in = new LZ4FrameInputStream(new ByteArrayInputStream(frame)) + val decoded = + try in.readAllBytes() + finally in.close() + assert(decoded.sameElements(data), s"lz4-java read $label back differently") + } finally { + packed.close() + } + } + assert(allocator.getAllocatedMemory == 0) + } finally { + allocator.close() + } + } + + test("Comet in-memory cache rejects a corrupt lz4 frame without leaking") { + // A corrupt frame has to fail as itself, naming the codec, and release the buffer it was + // decoding into. Offsets are from the LZ4 frame format behind Arrow's 8-byte + // uncompressed-length prefix: the frame's 4-byte magic number at 8, then a 3-byte frame + // descriptor, then the first block's 4-byte length at 15. A frame ends with a 4-byte end + // mark. Truncated frames are decoded from copies of exactly their length, so reading past + // one is an error of its own rather than a read of the bytes that follow. + val data = Array.tabulate[Byte](100 * 1000)(i => (i / 16).toByte) + val codec = CometCachedBatchHelper.compressionCodec("lz4") + + val allocator = CometArrowAllocator.newChildAllocator("lz4-corrupt", 0, Long.MaxValue) + try { + def truncated(frame: ArrowBuf, length: Long): ArrowBuf = { + val copy = allocator.buffer(length) + copy.setBytes(0, frame, 0, length) + copy.writerIndex(length) + copy + } + // Each corruption returns the buffer to decode: the frame itself, or a truncated copy. + val corruptions: Seq[(String, ArrowBuf => ArrowBuf)] = Seq( + "a buffer that is not an LZ4 frame" -> { b => b.setInt(8, 0); b }, + "a frame cut off inside its header" -> (b => truncated(b, 12)), + "a frame cut off before its end mark" -> (b => truncated(b, b.writerIndex() - 4)), + "a block that runs past the end of the frame" -> { b => b.setInt(15, 0x7fffff00); b }, + "a frame that decodes to less than its recorded length" -> { b => + b.setLong(0, data.length + 1L); b + }, + "a frame that decodes to more than its recorded length" -> { b => + b.setLong(0, data.length - 1L); b + }) + + corruptions.foreach { case (label, corrupt) => + val plain = allocator.buffer(data.length.toLong) + plain.setBytes(0, data) + plain.writerIndex(data.length.toLong) + val packed = codec.compress(allocator, plain) + assert(packed.getLong(0) == data.length, "the test needs a compressed buffer") + val input = corrupt(packed) + try { + val thrown = intercept[SparkException](codec.decompress(allocator, input).close()) + assert(thrown.getMessage.contains("LZ4"), s"$label failed with: $thrown") + } finally { + // A failed decompress leaves its input to the caller, as Arrow's own codecs do. + Seq(packed, input).distinct.foreach { b => + if (b.getReferenceManager.getRefCount > 0) b.close() + } + } + assert(allocator.getAllocatedMemory == 0, s"$label leaked") + } + } finally { + allocator.close() + } + } + test("Comet in-memory cache rejects a payload that records an unknown compression codec") { // `CodecType.fromCompressionType` answers NO_COMPRESSION for any byte outside its enum, so a // reader that took its word for it would read a compressed body as plain bytes and hand back @@ -2771,8 +2871,12 @@ class CometInMemoryCacheSuite extends CometTestBase { // reachable from nothing the failure path can see. The second is the one that catches a leak // in `VectorLoader`; the first is the one that catches a cleanup path releasing the shared // body twice. Both run over a multi-chunk payload too, where the copy into the body walks - // several chunks before the failure. - chunkSizes.foreach { case (_, chunkSize) => + // several chunks before the failure, and under each codec that compresses, since each fails + // in a decoder of its own. + for { + codec <- Seq("zstd", "lz4") + (_, chunkSize) <- chunkSizes + } withSQLConf(CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC.key -> codec) { withProjectionCache(chunkSize) { (relation, batches) => val cacheSchema = Utils.fromAttributes(relation.output) val pristine = CometCachedBatchHelper.snapshotPayloads(batches) @@ -2801,7 +2905,7 @@ class CometInMemoryCacheSuite extends CometTestBase { } assert( CometArrowAllocator.getAllocatedMemory == before, - s"everything allocated before a failure in $where must be released") + s"everything allocated before a failure in $where under $codec must be released") } } } 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 230126a7025..a7b8412eb72 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 @@ -19,8 +19,11 @@ package org.apache.spark.sql.benchmark +import java.io.PrintStream import java.nio.charset.StandardCharsets +import scala.collection.mutable + import org.apache.spark.SparkConf import org.apache.spark.benchmark.Benchmark import org.apache.spark.sql.SparkSession @@ -73,11 +76,8 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { Seq("id", "sc", "deep", "wide", "tail", "d"), nestedNumRows) - // Every value spark.comet.exec.inMemoryCache.compression.codec accepts, default first. Arrow's - // other IPC codec, LZ4_FRAME, is not one of them: it is commons-compress's pure-Java LZ4 rather - // than the JNI-accelerated lz4-java behind spark.io.compression.codec, and the write path - // rejects it. - private val codecs = Seq("zstd", "none") + // Every value spark.comet.exec.inMemoryCache.compression.codec accepts, default first. + private val codecs = Seq("zstd", "lz4", "none") // A sink for the benchmarked call's result: written but never read, so that neither the // compiler nor the JIT can treat `gatherColumnStats` as dead code. Not `private`, because @@ -444,19 +444,29 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { * CometInMemoryTableScan's batches through CometColumnarToRow instead, which * runAdaptiveBenchmark measures. * + * Comet's format is cached under zstd, the default, and under lz4, since most of what a Spark + * operator pays to read it is decompression. + * * 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. */ private def runSparkOperatorBenchmark(relation: CachedRelation): Unit = { val view = s"${relation.table}_spark_operators" - val reads = Seq( - SparkOperatorRead("Spark's cache format", sparkSerializer, sparkOperatorConf), - SparkOperatorRead("Comet's cache format, row reader", cometSerializer, sparkOperatorConf), - SparkOperatorRead( - "Comet's cache format, fused reader", - cometSerializer, - fusedReaderConf, - fused = true)) + val reads = SparkOperatorRead("Spark's cache format", sparkSerializer, sparkOperatorConf) +: + Seq("zstd", "lz4").flatMap { codec => + Seq( + SparkOperatorRead( + s"Comet's cache format, $codec, row reader", + cometSerializer, + sparkOperatorConf, + codec = codec), + SparkOperatorRead( + s"Comet's cache format, $codec, fused reader", + cometSerializer, + fusedReaderConf, + fused = true, + codec = codec)) + } spark.catalog.clearCache() withTempTable(view) { @@ -465,14 +475,16 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { .createOrReplaceTempView(view) val cache = new OneCachedCopy(view) - readShapes(view).foreach { case (label, query, scanned) => + val benchmarks = readShapes(view).map { case (label, query, scanned) => val benchmark = new Benchmark( s"in-memory cache read by Spark operators, $label", relation.rows.toLong, output = output) reads.foreach(addSparkOperatorCase(benchmark, cache, query, scanned, _)) benchmark.run() + benchmark } + cache.printFootprints(benchmarks.last.out) spark.catalog.uncacheTable(view) } @@ -488,12 +500,18 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { * spark.sql.codegen.maxFields (100 by default) that is how any Spark operator reads a cached * relation, because the scan stops offering columnar output and nothing can fuse with it. Every * width holds about the same number of values, so the cases differ in width rather than in data - * volume. + * volume. The columns are sequential longs, which Spark's format delta-encodes, so the + * footprint of each format is printed too. */ private def runWideSparkOperatorBenchmark(): Unit = { - val reads = Seq( - SparkOperatorRead("Spark's cache format", sparkSerializer, sparkOperatorConf), - SparkOperatorRead("Comet's cache format, row reader", cometSerializer, sparkOperatorConf)) + val reads = SparkOperatorRead("Spark's cache format", sparkSerializer, sparkOperatorConf) +: + Seq("zstd", "lz4").map { codec => + SparkOperatorRead( + s"Comet's cache format, $codec, row reader", + cometSerializer, + sparkOperatorConf, + codec = codec) + } Seq(100, 200, 1500).foreach { width => val rows = 20L * 1000 * 1000 / width @@ -513,37 +531,57 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { output = output) reads.foreach(addSparkOperatorCase(benchmark, cache, s"SELECT * FROM $view", width, _)) benchmark.run() + cache.printFootprints(benchmark.out) spark.catalog.uncacheTable(view) } } } - /** How Spark operators read a cached relation: its format and the settings of the read. */ + /** + * How Spark operators read a cached relation: its format, the codec Comet's format is written + * with, and the settings of the read. + */ private case class SparkOperatorRead( name: String, serializer: String, conf: Seq[(String, String)], - fused: Boolean = false) + fused: Boolean = false, + codec: String = "zstd") /** * 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. + * + * Records what each format occupies as it is cached, since a format that reads faster is only + * worth having at a footprint the application can afford. */ private class OneCachedCopy(view: String, conf: Seq[(String, String)] = sparkOperatorConf) { - private var cachedBy: String = _ + private var cachedBy: (String, String) = _ + private val footprints = mutable.LinkedHashMap.empty[String, Long] - def cacheBy(serializer: String): Unit = if (cachedBy != serializer) { - spark.catalog.uncacheTable(view) - cachedBy = null - withCacheSerializer(serializer) { - withSQLConf(conf: _*) { - spark.catalog.cacheTable(view) - spark.table(view).count() + // The codec only matters to Comet's serializer, and Spark's ignores its config. + def cacheBy(serializer: String, codec: String = "zstd"): Unit = + if (cachedBy != ((serializer, codec))) { + spark.catalog.uncacheTable(view) + cachedBy = null + withCacheSerializer(serializer) { + withSQLConf(conf ++ codecConf(codec): _*) { + spark.catalog.cacheTable(view) + spark.table(view).count() + } } + cachedBy = (serializer, codec) + val format = + if (serializer == sparkSerializer) "Spark's cache format" + else s"Comet's cache format, $codec" + footprints(format) = cachedBytes(view) } - cachedBy = serializer + + // As in runCodecBenchmark, footprint has no column in a Benchmark table. + def printFootprints(out: PrintStream): Unit = footprints.foreach { case (format, bytes) => + out.println(f"Cached footprint ($format): ${bytes / (1024.0 * 1024.0)}%.1f MiB") } } @@ -557,7 +595,7 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { // 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(read.name) { timer => - cache.cacheBy(read.serializer) + cache.cacheBy(read.serializer, read.codec) withSQLConf(read.conf: _*) { if (!verified) { verifySparkOperatorRead(query, scanned, read) @@ -573,8 +611,9 @@ 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. + * application gets from turning the feature on or off, and, through lz4, from changing the + * codec. 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 @@ -586,8 +625,8 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { */ 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 formats = ("Spark's cache format", sparkSerializer, "zstd") +: + Seq("zstd", "lz4").map(codec => (s"Comet's cache format, $codec", cometSerializer, codec)) val operatorsAbove = Seq( "Comet operators" -> Seq.empty[(String, String)], "a Spark operator" -> Seq(CometConf.COMET_EXEC_AGGREGATE_ENABLED.key -> "false")) @@ -607,12 +646,12 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { s"in-memory cache with AQE, $operators above the scan, $label", relation.rows.toLong, output = output) - formats.foreach { case (name, serializer) => + formats.foreach { case (name, serializer, codec) => 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) + cache.cacheBy(serializer, codec) withSQLConf(adaptiveConf ++ operatorConf: _*) { if (!verified) { verifyAdaptiveRead(query, scanned, serializer, operatorConf.nonEmpty) 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 47e876b5f86..7789e13ed7f 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 @@ -125,6 +125,14 @@ object CometCachedBatchHelper { chunkSize: Int = 1024 * 1024): ChunkedByteBuffer = CachedBatchIpc.serialize(batch, codec, allocator, chunkSize)._1 + /** + * The codec the cache writer uses for a `spark.comet.exec.inMemoryCache.compression.codec` + * value. A thin shim, for tests about a codec's own encoding rather than about a cached + * relation. + */ + def compressionCodec(name: String): CompressionCodec = + CachedBatchIpc.compressionCodec(name, zstdLevel = 1) + /** A payload [[serialize]] wrote, as the cached batch the writer would have stored it in. */ def cachedBatch(payload: ChunkedByteBuffer, numRows: Int): CachedBatch = CometCachedBatch(numRows, InternalRow.empty, payload)