Skip to content

perf: add opt-in delta encoding for cached longs #5869

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
peterxcli wants to merge 2 commits into
base: main
Choose a base branch
from
Open

perf: add opt-in delta encoding for cached longs #5869

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
72c3d8c
perf: add opt-in delta encoding for cached longs
peterxcli Oct 5, 2026
8701879
docs: list the cache delta encoding setting
peterxcli Oct 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions docs/source/user-guide/latest/in-memory-cache.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

@andygrove andygrove Oct 10, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The code compares the compressed deltas with the compressed values (encoded.writerIndex() < packed.writerIndex() * 3 / 4), not with the plain representation, so this sentence describes a different threshold. The CometConf doc string has the same gap, since it only says the delta buffer is over 25% smaller. Could both say the comparison is between the two compressed forms?

It would also help to give numbers here, the way the codec table under Storage format does, perhaps from a delta arm in runCodecBenchmark. That benchmark's relation is a good case for this option. I measured 54.8 MiB without deltas and 44.2 MiB with them. id, k and v went from 6.3, 0.9 and 6.5 MiB to 1.3, 0.4 and 1.5 MiB, and reads were no slower. The build cost on columns that are not sequential is worth stating too.

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.
Expand Down Expand Up @@ -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. |

Expand Down
11 changes: 11 additions & 0 deletions spark/src/main/scala/org/apache/comet/CometConf.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
37 changes: 28 additions & 9 deletions ...rc/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,15 +51,17 @@ 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.
*/
private case class CometCachedBatch(
override val numRows: Int,
override val stats: InternalRow,
bytes: ChunkedByteBuffer)
bytes: ChunkedByteBuffer,
deltaEncoded: Array[Boolean] = Array.emptyBooleanArray)
extends SimpleMetricsCachedBatch

/**
Expand All @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
}
}

Expand Down Expand Up @@ -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
Expand Down
104 changes: 94 additions & 10 deletions spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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()
}
Expand All @@ -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()
}
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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) =>
Expand Down Expand Up @@ -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 =>
Expand All @@ -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) {

@andygrove andygrove Oct 10, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This check sends every column whose values fit in an int through a second zstd pass, whether or not the deltas end up smaller. I cached a 5M-row relation of three such bigint columns: uniform below 1e6, uniform below 2^31, and a 1 ms step with up to 1 ms of jitter. With the option on, none of the 512 batches kept deltas for any column and the footprint stayed at 46.9 MiB. The build went from 210 ms to 299 ms, median of 7, and every run with the option on was slower than every run without it. Per 10,000-row column, uniform values below 1000 took 328 us to serialize instead of 169 us. Their deltas compress to 1.38 times the size of the values, so they can never pass the bar.

Could the loop also compare how wide the deltas are with how wide the values are, and only try the second pass when the deltas are clearly narrower? For example, summing 64 - numberOfLeadingZeros of each value's zigzag form, and trying only when the deltas come out at least four bits narrower on average. That skips every uniform case I tried and still tries all the sequential, constant-step and cyclic ones.

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,
Expand Down
13 changes: 12 additions & 1 deletion spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheKryoSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand All @@ -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.
Expand Down
Loading
Loading