You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/01 21:51:57 UTC
(datafusion-comet) 01/01: fix: report a cached relation's decoded size to the planner, not its compressed size (#6412)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6412-a76fd5a5862097297552d1bb28fc4d257be990ed
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit c083c705b3dc3bb09e0cb5d8bc0f347f6e6c419f
Author: Andy Grove <ag...@apache.org>
AuthorDate: Thu Oct 1 21:51:32 2026 +0000
fix: report a cached relation's decoded size to the planner, not its compressed size (#6412)
* fix: report a cached relation's decoded size to the planner, not its compressed size
Comet's cache format reported each batch's sizeInBytes as its stored,
compressed payload. Spark's planner reads the sum of those as the size of a
materialized cached relation, for the broadcast threshold and the shuffled
hash join build side among others, and both of Spark's own cache formats
report decoded sizes there. With the default zstd codec a cached relation
could look several times smaller than in Spark's format and be broadcast
where Spark would shuffle it.
Record each column's decoded Arrow size, measured before compression as
Spark's ArrowCachedBatchSerializer does, and let CometCachedBatch inherit
sizeInBytes from SimpleMetricsCachedBatch as the sum of those.
* test: check the dictionary path's recorded sizes, and state the size rationale once
---
docs/source/user-guide/latest/in-memory-cache.md | 3 +
.../arrow/ArrowCachedBatchSerializer.scala | 20 +++---
.../sql/comet/execution/arrow/CachedBatchIpc.scala | 30 ++++----
.../comet/exec/CometInMemoryCacheSuite.scala | 82 ++++++++++++++++------
.../benchmark/CometInMemoryCacheBenchmark.scala | 10 +--
.../execution/arrow/CometCachedBatchHelper.scala | 16 ++++-
6 files changed, 114 insertions(+), 47 deletions(-)
diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md
index 5704365ac9..8d2ccd2521 100644
--- a/docs/source/user-guide/latest/in-memory-cache.md
+++ b/docs/source/user-guide/latest/in-memory-cache.md
@@ -48,6 +48,9 @@ With Comet's serializer installed as `spark.sql.cache.serializer`:
- Cached tables are scanned by `CometInMemoryTableScan`, which feeds Comet operators directly.
- Per-batch column statistics are recorded in the layout Spark's `SimpleMetricsCachedBatchSerializer`
expects, so Spark can prune whole cached batches on a predicate before any of them is decoded.
+- The size Spark's planner sees for a cached relation is its decoded Arrow size, not the compressed
+ size it occupies in memory, as with Spark's own cache formats. Compression therefore does not
+ change how queries over a cached relation are planned, such as whether a join broadcasts it.
Relations whose schema Comet's Arrow writer cannot store — interval types, most notably — are
delegated in full to Spark's default cache format, per relation. Which format a relation uses does
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 4769155684..103df61e17 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
@@ -52,10 +52,12 @@ import org.apache.comet.vector.NativeUtil
* 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.
+ *
+ * `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 sizeInBytes: Long,
override val stats: InternalRow,
bytes: ChunkedByteBuffer)
extends SimpleMetricsCachedBatch
@@ -350,11 +352,14 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
values(base + 1) = upper(c)
values(base + 2) = nulls(c)
values(base + 3) = numRows
- // The stored size of the column's own Arrow buffers, taken from the message's buffer
- // layout, so it is exact rather than an estimate. Cache pruning uses
- // bounds/null-count/row-count rather than this field, but Spark reserves it and reports it,
- // so record the real value. The per-batch message framing is not attributed to any column,
- // so these sum to slightly less than sizeInBytes.
+ // The column's decoded size: the plain length of its own Arrow buffers before compression.
+ // SimpleMetricsCachedBatch sums these into sizeInBytes, which Spark's planner reads as the
+ // size of a materialized cached relation, for the broadcast threshold and the shuffled hash
+ // join build side among others. Spark's DefaultCachedBatchSerializer also records each
+ // column's uncompressed size here (ColumnStats.sizeInBytes). The compressed payload can be
+ // several times smaller, and a relation reported at that size would be planned differently
+ // from the same relation in Spark's cache format, for example broadcast where Spark's cache
+ // would have it shuffled.
values(base + 4) = columnSizes(c)
c += 1
}
@@ -401,7 +406,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
batches.map { batch =>
// Bounds and null counts are read from the input batch before it is serialized, and the row
- // is only assembled once the per-column sizes the message reports are known.
+ // is only assembled once serialize has measured the per-column sizes.
val (lower, upper, nulls) = gatherColumnStats(batch, attrs, orderings)
val numRows = batch.numRows()
@@ -423,7 +428,6 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
CometCachedBatch(
numRows = numRows,
- sizeInBytes = bytes.size,
stats = statsRow(lower, upper, nulls, numRows, columnSizes),
bytes = bytes)
}
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 d92971c2bf..06c3e52759 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
@@ -186,9 +186,10 @@ 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 on-body compressed size of each top-level column, which the
- * caller records in the statistics row. The sizes come from the message's own buffer layout, so
- * they are the real stored sizes rather than an estimate.
+ * 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.
*
* 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
@@ -215,16 +216,18 @@ private[comet] object CachedBatchIpc {
// Unloaded plain and compressed afterwards rather than by handing the codec to the unloader;
// see compressed for why.
+ val fields = vectors.map(_.getField)
val unloader = new VectorUnloader(root, true, NoCompressionCodec.INSTANCE, true)
val plainBatch = unloader.getRecordBatch
- val recordBatch =
- try compressed(plainBatch, codec, allocator)
- finally plainBatch.close()
+ val (sizes, recordBatch) =
+ try {
+ (columnSizes(fields, plainBatch), compressed(plainBatch, codec, allocator))
+ } finally {
+ plainBatch.close()
+ }
try {
- val fields = vectors.map(_.getField)
// Leaves the batch in the state serializeBatches leaves one. The record batch holds its
- // own buffers by now, so this does not touch it, and getField still answers afterwards:
- // clearing releases buffers, not the schema.
+ // own buffers by now, so this does not touch it.
root.clear()
val out = new ChunkedByteBufferOutputStream(chunkSize, ByteBuffer.allocate)
@@ -234,7 +237,7 @@ private[comet] object CachedBatchIpc {
} finally {
out.close()
}
- (out.toChunkedByteBuffer, columnSizes(fields, recordBatch))
+ (out.toChunkedByteBuffer, sizes)
} finally {
recordBatch.close()
}
@@ -608,11 +611,10 @@ private[comet] object CachedBatchIpc {
}
/**
- * The on-body compressed size of each top-level column.
+ * The size of each top-level column's buffers in `recordBatch`.
*
- * Each column owns the run of buffers its subtree occupies, so its stored size is the sum of
- * those buffers' recorded lengths. With one payload per batch these are the only per-column
- * sizes available -- there is no separate stream to measure -- and they are exact.
+ * Each column owns the run of buffers its subtree occupies, so its size is the sum of those
+ * buffers' recorded lengths.
*/
private def columnSizes(fields: Seq[Field], recordBatch: ArrowRecordBatch): Array[Long] = {
val buffers = recordBatch.getBuffersLayout
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 d82f6c021b..c963de3589 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
@@ -191,10 +191,11 @@ class CometInMemoryCacheSuite extends CometTestBase {
// https://github.com/apache/spark/blob/v4.1.2/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala#L2780-L2832
test("AQE SPARK-37742: use valid Comet cache statistics for join selection") {
withAQECache {
- // Comet reports compressed Arrow bytes, so use a threshold below the compressed
- // large cache as well as below its logical estimate. The single-row side still fits.
+ // Spark's own threshold. The large cache has to stay above it once materialized: its 60k
+ // 20-byte keys are about 1.4 MB decoded but compress to a small fraction of that, so a cache
+ // that reported its compressed size would be broadcast by the third join.
withSQLConf(
- SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1024",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1048584",
SQLConf.ADAPTIVE_OPTIMIZER_EXCLUDED_RULES.key ->
"org.apache.spark.sql.execution.adaptive.AQEPropagateEmptyRelation") {
withTempView("cache_large", "cache_other", "cache_small") {
@@ -235,7 +236,7 @@ class CometInMemoryCacheSuite extends CometTestBase {
val stats = relation.computeStats()
assert(stats.rowCount.contains(BigInt(60000)))
assert(stats.sizeInBytes == batches.map(_.sizeInBytes).sum)
- assert(stats.sizeInBytes > 1024L)
+ assert(stats.sizeInBytes > 1048584L)
}
}
}
@@ -2036,9 +2037,6 @@ class CometInMemoryCacheSuite extends CometTestBase {
assert(
CometCachedBatchHelper.chunkCount(batch) > 1,
"a payload larger than the chunk size must be stored in more than one chunk")
- assert(
- CometCachedBatchHelper.payloadSize(batch) == batch.sizeInBytes,
- "sizeInBytes must report the whole payload across its chunks")
}
projections.zip(expected).foreach { case (cols, rows) =>
@@ -2066,30 +2064,74 @@ class CometInMemoryCacheSuite extends CometTestBase {
}
}
- test("Comet in-memory cache records per-column sizes in its statistics") {
- // SimpleMetricsCachedBatch reserves a fifth field per column for its size. A column owns a
- // known run of buffers in the payload, so the real stored size is known and must be reported
- // rather than left at zero. Run over the nested relation as well: a nested column's size is
+ test("Comet in-memory cache records per-column decoded sizes in its statistics") {
+ // Each column's size field must be its decoded size (see ArrowCachedBatchSerializer.statsRow),
+ // so each field is compared with its column decoded back out of the payload, not with what the
+ // column occupies compressed. Run over the nested relation as well: a nested column's size is
// the sum of its whole subtree, so this is also where a size attributed to the wrong column
- // surfaces.
+ // surfaces. And over the dictionary relation, whose columns reach the writer dictionary
+ // encoded: a size measured before they are decoded would be the size of their indices.
def checkSizes(
relation: org.apache.spark.sql.execution.columnar.InMemoryRelation,
batches: Array[CachedBatch]): Unit = {
val cacheSchema = Utils.fromAttributes(relation.output)
- batches.foreach { batch =>
- val sizes = CometCachedBatchHelper.columnSizes(batch, cacheSchema)
- val stats = batch.asInstanceOf[SimpleMetricsCachedBatch].stats
- sizes.zipWithIndex.foreach { case (size, i) =>
- assert(
- stats.getLong(i * 5 + 4) == size,
- s"column ${relation.output(i).name} should report the stored size of its own " +
- "buffers in the statistics row")
+ val allocator = CometArrowAllocator.newChildAllocator("decoded-sizes", 0, Long.MaxValue)
+ try {
+ batches.foreach { batch =>
+ val sizes = CometCachedBatchHelper.decodedColumnSizes(batch, cacheSchema, allocator)
+ val stats = batch.asInstanceOf[SimpleMetricsCachedBatch].stats
+ sizes.zipWithIndex.foreach { case (size, i) =>
+ assert(
+ stats.getLong(i * 5 + 4) == size,
+ s"column ${relation.output(i).name} should report its decoded size in the " +
+ "statistics row")
+ }
+ assert(batch.sizeInBytes == sizes.sum)
}
+ } finally {
+ allocator.close()
}
}
withProjectionCache(checkSizes _)
withNestedProjectionCache(checkSizes _)
+ withDictionaryCache(relation =>
+ checkSizes(relation, relation.cacheBuilder.cachedColumnBuffers.collect()))
+ }
+
+ test("Comet in-memory cache reports the same relation size under every codec") {
+ // The planner compares this size with broadcast thresholds and with the other side of a
+ // shuffled hash join, so it must not depend on the codec (see
+ // ArrowCachedBatchSerializer.statsRow).
+ val measured = Seq.newBuilder[(String, Long, Long)]
+ Seq("zstd", "none").foreach { codec =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC.key -> codec) {
+ spark.catalog.clearCache()
+ spark
+ .range(0, 20000, 1, 2)
+ .selectExpr("id", "id % 7 AS k", "concat('v_', cast(id % 100 AS string)) AS s")
+ .createOrReplaceTempView("codec_size_cache")
+ spark.catalog.cacheTable("codec_size_cache")
+ assert(spark.table("codec_size_cache").count() == 20000)
+ val relation = spark.sharedState.cacheManager
+ .lookupCachedData(spark.table("codec_size_cache"))
+ .get
+ .cachedRepresentation
+ val batches = relation.cacheBuilder.cachedColumnBuffers.collect()
+ val payload = batches.map(CometCachedBatchHelper.payloadSize).sum
+ measured += ((codec, relation.computeStats().sizeInBytes.toLong, payload))
+ spark.catalog.clearCache()
+ }
+ }
+
+ val sizes = measured.result()
+ val (_, zstdSize, zstdPayload) = sizes(0)
+ val (_, plainSize, _) = sizes(1)
+ assert(zstdPayload * 2 < zstdSize, s"zstd should compress this relation: $sizes")
+ assert(zstdSize == plainSize, s"the relation size should not depend on the codec: $sizes")
}
test("Comet in-memory cache scans no columns for a row-count-only query") {
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 1b0d6eff78..9b0f652a19 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
@@ -26,7 +26,7 @@ import org.apache.spark.benchmark.Benchmark
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.expressions.AttributeReference
import org.apache.spark.sql.comet.CometInMemoryTableScanExec
-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, DefaultCachedBatchSerializer, InMemoryRelation, InMemoryTableScanExec}
import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf}
@@ -452,10 +452,12 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase {
.optimizedPlan
.collectFirst { case r: InMemoryRelation => r }
.getOrElse(sys.error(s"$view is not cached"))
- // computeStats rather than the builder's size accumulator, which Spark 4.2 replaced. Before
- // the buffers load it falls back to the plan's estimate, so insist they have.
+ // The stored payloads rather than computeStats, which reports the relation's decoded size and
+ // so is the same for every codec.
assert(relation.cacheBuilder.isCachedColumnBuffersLoaded, s"$view is not materialized")
- relation.computeStats().sizeInBytes.toLong
+ relation.cacheBuilder.cachedColumnBuffers
+ .map(CometCachedBatchHelper.payloadSize)
+ .fold(0L)(_ + _)
}
private def runStatsBenchmark(): Unit = {
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 14975cdd21..47e876b5f8 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
@@ -127,7 +127,7 @@ object CometCachedBatchHelper {
/** A payload [[serialize]] wrote, as the cached batch the writer would have stored it in. */
def cachedBatch(payload: ChunkedByteBuffer, numRows: Int): CachedBatch =
- CometCachedBatch(numRows, payload.size, InternalRow.empty, payload)
+ CometCachedBatch(numRows, InternalRow.empty, payload)
/**
* Decode the `selected` columns of a cached batch the way a scan does, into a root the caller
@@ -197,6 +197,20 @@ object CometCachedBatchHelper {
def columnSizes(batch: CachedBatch, cacheSchema: StructType): Seq[Long] =
columnBufferRanges(batch, cacheSchema).map(_.map(_._2).sum)
+ /**
+ * Decoded size of each top-level column: the column read back out of the payload the way a scan
+ * reads it, measured with `getBufferSize`.
+ */
+ def decodedColumnSizes(
+ batch: CachedBatch,
+ cacheSchema: StructType,
+ allocator: BufferAllocator): Seq[Long] =
+ cacheSchema.indices.map { i =>
+ val root = load(batch, cacheSchema, Array(i), allocator)
+ try root.getVector(0).getBufferSize.toLong
+ finally root.close()
+ }
+
/**
* Whether any of a column's buffers is actually stored compressed.
*
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org