From d7ae3bb2438e97c5d0ea70bf47b77413867a6c7d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 28 Aug 2026 17:51:57 -0600 Subject: [PATCH 1/5] perf: project cached batches by buffer selection, prune on collated strings Follow-up to #5051, applying items from #5487. Replace the per-column Arrow IPC stream layout of `CometCachedBatch` with a single encapsulated IPC record batch message per cached batch, carrying no Schema message and no end-of-stream marker. The reader rebuilds the schema from the cached relation's attributes, so a wide relation no longer repeats the same schema bytes once per cached batch. Compression moves from a whole-payload Spark codec to Arrow's per-buffer IPC compression. That is what makes projection cheap: the message metadata records every buffer's offset and length in the body, so `CachedBatchIpc.readProjected` copies out only the byte ranges of the columns a scan selected and decompresses just those. This subsumes the separate "drop the schema message" item, since there is no longer a per-column stream to frame. Dictionary-encoded columns are decoded before being stored: a payload with no schema message cannot describe a dictionary encoding. The codec defaults to zstd, and lz4 is deliberately not offered. Arrow's lz4 is commons-compress's pure-Java implementation, unrelated to the JNI-accelerated lz4-java behind `spark.io.compression.codec`. Over a 200k-row six-column relation it measured 205s to write against 347ms for zstd, while also producing larger output, so no workload prefers it. zstd also beats storing batches uncompressed on both axes (347ms and 2 MiB against 1743ms and 13 MiB), because the bytes it saves cost more to copy and store than compressing them costs. Decompression is done here rather than left to `VectorLoader`, which leaks: `VectorLoader.loadBuffers` collects a field's decompressed buffers into a local list and releases them only after the whole field loads, so a buffer that fails to decompress strands every buffer of that field decompressed before it. A string column reaches this, its offsets buffer decompressing before its data buffer throws. Also track statistics bounds for collated string columns, comparing with the collation's own ordering through a new `CometTypeShim.compareStrings`. Matching the bare `StringType` object excluded collated columns, which then got null bounds and no pruning. Benchmark over a 5M-row six-column relation, keeping the cached scan native against falling back to a Spark cache scan and converting: 1.3x on a repeated scan, 1.3x on a narrow projection and 2.3x on a full projection. --- .../user-guide/latest/in-memory-cache.md | 148 +++++++ docs/source/user-guide/latest/index.rst | 1 + pom.xml | 26 ++ spark/pom.xml | 4 + .../scala/org/apache/comet/CometConf.scala | 55 ++- .../arrow/ArrowCachedBatchSerializer.scala | 245 ++++++----- .../execution/arrow/CachedBatchIpc.scala | 409 ++++++++++++++++++ .../apache/spark/sql/comet/util/Utils.scala | 39 +- .../apache/comet/shims/CometTypeShim.scala | 12 + .../apache/comet/shims/CometTypeShim.scala | 11 + .../comet/exec/CometInMemoryCacheSuite.scala | 303 ++++++++++--- .../CometInMemoryCacheBenchmark.scala | 7 +- .../arrow/CometCachedBatchHelper.scala | 259 ++++++++--- 13 files changed, 1216 insertions(+), 303 deletions(-) create mode 100644 docs/source/user-guide/latest/in-memory-cache.md create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md new file mode 100644 index 00000000000..67f0b0e9755 --- /dev/null +++ b/docs/source/user-guide/latest/in-memory-cache.md @@ -0,0 +1,148 @@ + + +# In-Memory Cache + +Comet can store Spark's in-memory cache (`CACHE TABLE`, `df.cache()`, `df.persist()`) in an Arrow +format that Comet operators read directly. Without it, a cached table is stored in Spark's own +format and every scan of it has to convert each batch before Comet can continue, which shows up in +the plan as a `CometSparkColumnarToColumnar` above the cache scan. + +This feature is **experimental and disabled by default**. + +```scala +spark.conf.set("spark.comet.exec.inMemoryCache.enabled", "true") +``` + +## What changes when it is enabled + +`spark.comet.exec.inMemoryCache.enabled` is read at startup, and its value then decides whether +Comet installs its cache serializer as `spark.sql.cache.serializer`. When it is installed: + +- Cached data is stored as `CometCachedBatch` rather than Spark's `DefaultCachedBatch`. +- 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. + +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. Nothing about the format depends +on a runtime config, because `spark.sql.cache.serializer` is a static setting and a relation whose +format could change mid-session could not be read back reliably. Turning +`spark.comet.exec.inMemoryCache.enabled` off at runtime only sends cached scans back to Spark's +execution path; the cached data stays readable either way. + +## Storage format + +Each cached batch is stored as a single Arrow IPC record batch message and its body. + +The message carries **no Arrow schema**. The reader already has one: `InMemoryRelation` knows the +cached relation's attributes, and Comet maps them to exactly the Arrow fields the writer produced. +Storing a schema in every batch would repeat the same bytes once per cached batch — for a wide +relation cached in many batches, a large share of a payload that is not data. + +Compression is applied by Arrow to **each buffer separately**, rather than by wrapping the whole +payload in a Spark compression codec. That is what makes a projected read cheap: the message +metadata records every buffer's offset and length within the body, so a scan copies out only the +byte ranges belonging to the columns it selected, and only those are decompressed. A read of one +column out of six does roughly a sixth of the decompression work, and a `SELECT count(*)`, which +selects no columns at all, answers from the row count stored beside the payload without touching +it. + +Compression defaults to `zstd`, which is faster than storing cached batches uncompressed: the +bytes it saves cost more to copy and store than compressing them costs. Measured over a 200k-row, +six-column relation: + +| Codec | Materialize | Footprint | Read 1 of 6 | Read 6 of 6 | +| ------ | ----------: | --------: | ----------: | ----------: | +| `zstd` | 347 ms | 2 MiB | 52 ms | 63 ms | +| `none` | 1743 ms | 13 MiB | 74 ms | 79 ms | + +Arrow's other IPC codec, LZ4, is deliberately not offered. It is commons-compress's pure-Java +implementation and is unrelated to the JNI-accelerated lz4-java behind `spark.io.compression.codec`; +it measured three orders of magnitude slower to write than `zstd` while also producing larger +output, so no workload prefers it. + +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` | `false` | Whether to store and scan Spark's in-memory cache in Comet's format. Read at startup. | +| `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. | + +## Performance + +Measured with `CometInMemoryCacheBenchmark` on a 5M-row, six-column relation (Apple M3 Ultra, +JDK 17, Spark 4.1, release build). Regenerate with: + +```sh +SPARK_GENERATE_BENCHMARK_FILES=1 \ + make benchmark-org.apache.spark.sql.benchmark.CometInMemoryCacheBenchmark +``` + +| Query shape | Spark cache scan + convert | `CometInMemoryTableScan` | Relative | +| ------------------------------ | -------------------------: | -----------------------: | -------: | +| Repeated scan (3 of 6 columns) | 156 ms | 118 ms | 1.3x | +| Selective filter | 44 ms | 39 ms | 1.1x | +| Row count only (0 of 6) | 32 ms | 28 ms | 1.1x | +| Narrow projection (1 of 6) | 50 ms | 39 ms | 1.3x | +| Full projection (6 of 6) | 316 ms | 135 ms | 2.3x | + +Read what this compares carefully. Comet execution is on in both columns, so the aggregation runs +on Comet either way and only the cache-scan boundary moves: on the left, Spark's +`InMemoryTableScanExec` feeds those same Comet operators through a `CometSparkColumnarToColumnar` +bridge; on the right, `CometInMemoryTableScan` feeds them directly. Both columns read the same +Comet-written `CometCachedBatch` — `spark.sql.cache.serializer` is static, so one session cannot +also materialize Spark's format to compare against. These numbers are therefore "keep the cached +scan native" against "fall back to a Spark cache scan and convert", not Comet against Spark +execution, and not a comparison with Spark's own cache format. + +## Kryo + +Spark serializes a cached batch with `spark.serializer` whenever the block leaves the heap: the +`_SER` storage levels, replication, cross-executor fetches, and the disk half of the default +`MEMORY_AND_DISK`. So an ordinary `df.cache()` that spills is enough to reach it. + +If you run with `spark.kryo.registrationRequired=true`, register Comet's classes: + +``` +spark.serializer=org.apache.spark.serializer.KryoSerializer +spark.kryo.registrationRequired=true +spark.kryo.registrator=org.apache.comet.CometKryoRegistrator +``` + +Comet cannot set `spark.kryo.registrator` for you the way it sets `spark.sql.cache.serializer`: +`KryoSerializer` reads it when `SparkEnv` builds the serializer, which happens before any plugin +runs. Without it, caching fails with a "Class is not registered" error that does not name this +feature. Comet's driver plugin warns at startup when it sees Kryo, `registrationRequired`, and no +registrator. + +## Limitations + +Reads that feed **Spark** operators rather than Comet ones are still slower than Spark's own cache +format, by roughly 1.7x to 2.5x depending on how wide the projection is. Those reads pay a row +conversion that Spark's format avoids with generated code over its own layout. This is why the +feature is off by default. + +Comet's serializer exists because Spark's own Arrow cache format +([SPARK-57268](https://issues.apache.org/jira/browse/SPARK-57268)) is only available from Spark +4.3, which Comet does not yet support. diff --git a/docs/source/user-guide/latest/index.rst b/docs/source/user-guide/latest/index.rst index 815e12289c7..063a5581a04 100644 --- a/docs/source/user-guide/latest/index.rst +++ b/docs/source/user-guide/latest/index.rst @@ -74,6 +74,7 @@ to read more. Understanding Comet Plans Tuning Guide Metrics Guide + In-Memory Cache PyArrow UDF Acceleration .. toctree:: diff --git a/pom.xml b/pom.xml index 11494665755..2fa63d71103 100644 --- a/pom.xml +++ b/pom.xml @@ -225,6 +225,32 @@ under the License. arrow-c-data ${arrow.version} + + org.apache.arrow + arrow-compression + ${arrow.version} + + + + org.apache.commons + commons-compress + + + com.github.luben + zstd-jni + + + io.netty + netty-common + + + com.google.code.findbugs + jsr305 + + + diff --git a/spark/pom.xml b/spark/pom.xml index a257415dd3f..d30afc75914 100644 --- a/spark/pom.xml +++ b/spark/pom.xml @@ -60,6 +60,10 @@ under the License. org.apache.arrow arrow-vector + + org.apache.arrow + arrow-compression + org.scala-lang.modules scala-collection-compat_${scala.binary.version} diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index f5b7c7c09d0..4183ffef457 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -261,23 +261,50 @@ object CometConf extends ShimCometConf { val COMET_EXEC_IN_MEMORY_CACHE_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.exec.inMemoryCache.enabled") .category(CATEGORY_EXEC) - .doc( - "Whether to enable Comet native execution for in-memory cached tables. Its value at " + - "startup also decides whether CometDriverPlugin installs Comet's cache serializer, " + - "which stores cached data in Arrow format. Because spark.sql.cache.serializer is a " + - "static config, the cached format is fixed for the application, and disabling this " + - "at runtime only sends cached scans back to Spark's execution path. Relations whose " + - "schema Comet's Arrow writer does not support are always cached in Spark's default " + - "format. Each cached column is stored as its own compressed Arrow IPC stream, so a " + - "scan decodes only the columns it projected. Reads that feed Spark operators rather " + - "than Comet ones still pay a row conversion the default format avoids, and can be " + - "slower than Spark's cache. With spark.kryo.registrationRequired=true, also set " + - "spark.kryo.registrator=org.apache.comet.CometKryoRegistrator before creating the " + - "SparkContext, otherwise caching fails as soon as a block is serialized, including " + - "the disk half of the default MEMORY_AND_DISK storage level.") + .doc("Whether to enable Comet native execution for in-memory cached tables. Its value at " + + "startup also decides whether CometDriverPlugin installs Comet's cache serializer, " + + "which stores cached data in Arrow format. Because spark.sql.cache.serializer is a " + + "static config, the cached format is fixed for the application, and disabling this " + + "at runtime only sends cached scans back to Spark's execution path. 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 projected, " + + "so the unselected ones are never decompressed. Reads that feed Spark operators rather " + + "than Comet ones still pay a row conversion the default format avoids, and can be " + + "slower than Spark's cache. With spark.kryo.registrationRequired=true, also set " + + "spark.kryo.registrator=org.apache.comet.CometKryoRegistrator before creating the " + + "SparkContext, otherwise caching fails as soon as a block is serialized, including " + + "the disk half of the default MEMORY_AND_DISK storage level.") .booleanConf .createWithDefault(false) + val COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC: ConfigEntry[String] = + conf("spark.comet.exec.inMemoryCache.compression.codec") + .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. Set to none to store cached batches " + + "uncompressed, which is both slower to write and larger than zstd because the extra " + + "bytes cost more to move and store than compressing them costs. 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.") + .stringConf + .checkValues(Set("none", "zstd")) + .createWithDefault("zstd") + + val COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_ZSTD_LEVEL: ConfigEntry[Int] = + conf("spark.comet.exec.inMemoryCache.compression.zstd.level") + .category(CATEGORY_EXEC) + .doc("The compression level to use when Comet's cache serializer compresses cached data " + + "with zstd. Ignored for other codecs.") + .intConf + .createWithDefault(1) + val COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.columnarToRow.native.enabled") .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 46a51ad8775..24eba56f737 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 @@ -22,6 +22,8 @@ package org.apache.spark.sql.comet.execution.arrow import scala.collection.JavaConverters._ import scala.util.control.NonFatal +import org.apache.arrow.vector.VectorSchemaRoot +import org.apache.arrow.vector.types.pojo.{Field, Schema} import org.apache.spark.TaskContext import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow @@ -33,25 +35,27 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.spark.storage.StorageLevel -import org.apache.spark.unsafe.types.{ByteArray, UTF8String} -import org.apache.spark.util.io.ChunkedByteBuffer +import org.apache.spark.unsafe.types.UTF8String -import org.apache.comet.CometArrowAllocator +import org.apache.comet.{CometArrowAllocator, CometConf} +import org.apache.comet.shims.CometTypeShim +import org.apache.comet.vector.NativeUtil /** * Cached batch format used when Comet writes Spark in-memory cache data. * - * `columns` holds one compressed Arrow stream per cached column, in cache-schema order, produced - * by `Utils.serializeBatchColumns`. Storing columns separately is what lets a scan decode only - * the ones it projected; a single stream covering the whole batch would have to be inflated in - * full before any projection could be applied. The cache manager still owns storage and eviction; - * this class only changes the cached payload. + * `bytes` is one encapsulated Arrow IPC RecordBatch message and its body, with no Schema message + * and no end-of-stream marker, produced by `CachedBatchIpc.serialize`. Compression is applied per + * Arrow buffer rather than over the payload as a whole, which is what lets a scan decompress only + * the columns it projected: the message records every buffer's offset and length, so + * `CachedBatchIpc.readProjected` copies out just the selected columns' byte ranges. The cache + * manager still owns storage and eviction; this class only changes the cached payload. */ private case class CometCachedBatch( override val numRows: Int, override val sizeInBytes: Long, override val stats: InternalRow, - columns: Array[ChunkedByteBuffer]) + bytes: Array[Byte]) extends SimpleMetricsCachedBatch /** @@ -69,7 +73,7 @@ private case class CometCachedBatch( * Reads of `CometCachedBatch` keep working when the native scan is disabled, because Spark then * reads the same cached data through the SparkToColumnar fallback path. */ -class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { +class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer with CometTypeShim { import ArrowCachedBatchSerializer.supportsSchema @@ -130,9 +134,11 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { values(base + 1) = upper(c) values(base + 2) = nulls(c) values(base + 3) = numRows - // Each column is its own compressed stream, so its size is known exactly. 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 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. values(base + 4) = columnSizes(c) c += 1 } @@ -142,9 +148,15 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // Spark can prune cache batches only for types whose bounds can be compared. // Other types still report null count and row count but leave bounds as null. + // + // Every StringType qualifies, collated ones included: bounds are recorded with the collation's + // own comparison, which is the same ordering the predicate Spark generates over that column + // uses. Matching the bare `StringType` object instead would exclude collated columns, since a + // collated StringType is not equal to the default one, and they would then get null bounds and + // no pruning at all. private def tracksBounds(dt: DataType): Boolean = dt match { case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | - _: DecimalType | StringType | DateType | TimestampType | TimestampNTZType => + _: DecimalType | _: StringType | DateType | TimestampType | TimestampNTZType => true case _ => false } @@ -160,7 +172,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { case FloatType => col.getFloat(rowId) case DoubleType => col.getDouble(rowId) case d: DecimalType => col.getDecimal(rowId, d.precision, d.scale) - case StringType => col.getUTF8String(rowId).copy() + case _: StringType => col.getUTF8String(rowId).copy() case _ => null } @@ -182,10 +194,8 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { java.lang.Double.compare(left.asInstanceOf[Double], right.asInstanceOf[Double]) case _: DecimalType => left.asInstanceOf[Decimal].compare(right.asInstanceOf[Decimal]) - case StringType => - ByteArray.compareBinary( - left.asInstanceOf[UTF8String].getBytes, - right.asInstanceOf[UTF8String].getBytes) + case st: StringType => + compareStrings(left.asInstanceOf[UTF8String], right.asInstanceOf[UTF8String], st) case other => throw new IllegalStateException(s"compare called for unsupported type $other") } @@ -199,32 +209,33 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // writes CometVector columns. private def encodeBatches( batches: Iterator[ColumnarBatch], - attrs: Seq[Attribute]): Iterator[CachedBatch] = { + attrs: Seq[Attribute], + codecName: String, + zstdLevel: Int): Iterator[CachedBatch] = { val arrowSchema = Utils.toArrowSchema(Utils.fromAttributes(attrs), CometArrowStream.NATIVE_TIMEZONE) + val codec = CachedBatchIpc.compressionCodec(codecName, zstdLevel) batches.map { batch => - // Bounds and null counts are read from the input batch, which serializing then clears, so - // they have to be gathered first. The row is only assembled once the per-column sizes are - // known. + // 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. val (lower, upper, nulls) = gatherColumnStats(batch, attrs) val numRows = batch.numRows() - val columns = if (Utils.isArrowBacked(batch)) { - Utils.serializeBatchColumns(batch) + val (bytes, columnSizes) = if (Utils.isArrowBacked(batch)) { + CachedBatchIpc.serialize(batch, codec, CometArrowAllocator) } else { val arrowBatch = CometArrowConverters.columnarBatchToArrowBatch(batch, arrowSchema, CometArrowAllocator) - try Utils.serializeBatchColumns(arrowBatch) + try CachedBatchIpc.serialize(arrowBatch, codec, CometArrowAllocator) finally arrowBatch.close() } - val columnSizes = columns.map(_.size) CometCachedBatch( numRows = numRows, - sizeInBytes = columnSizes.sum, + sizeInBytes = bytes.length.toLong, stats = statsRow(lower, upper, nulls, numRows, columnSizes), - columns = columns) + bytes = bytes) } } @@ -293,16 +304,22 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { } } - // Columnar Comet output is stored as compressed Arrow stream bytes. Spark only calls this when - // supportsColumnarInput returned true, so the schema is known to be Comet-writable here. + // Columnar Comet output is stored as one Arrow IPC record batch message per cached batch. Spark + // only calls this when supportsColumnarInput returned true, so the schema is known to be + // Comet-writable here. override def convertColumnarBatchToCachedBatch( input: RDD[ColumnarBatch], schema: Seq[Attribute], storageLevel: StorageLevel, conf: SQLConf): RDD[CachedBatch] = { + // Read on the driver: the closure ships to the executors, where CometConf would resolve + // against whatever SQLConf happens to be current on that thread rather than this session's. + val codecName = CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC.get(conf) + val zstdLevel = CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_ZSTD_LEVEL.get(conf) + input.mapPartitions { batches => - encodeBatches(batches, schema) + encodeBatches(batches, schema, codecName, zstdLevel) } } @@ -320,24 +337,33 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { } val indices = selectedIndices(cacheAttributes, selectedAttributes) + // Captured as a StructType rather than the attributes themselves: this closure ships to the + // executors, and the Arrow schema is rebuilt there from the same mapping the writer used. + val cacheSchema = Utils.fromAttributes(cacheAttributes) input.mapPartitions { it => - // A ColumnReaders closes its readers (releasing the vectors they are holding) only when the - // batch it produced has been consumed. A consumer that stops early -- LIMIT, take(), or a - // cancelled task -- leaves the readers for the batch in flight open, so close them on task - // completion. Spark's own ArrowCachedBatchSerializer registers a listener for the same - // reason. + val arrowFields = + Utils + .toArrowSchema(cacheSchema, CometArrowStream.NATIVE_TIMEZONE) + .getFields + .asScala + .toSeq + + // A ProjectedBatch owns the vectors of the batch it produced, and releases them only when + // that batch has been consumed. A consumer that stops early -- LIMIT, take(), or a + // cancelled task -- leaves the batch in flight open, so close it on task completion. + // Spark's own ArrowCachedBatchSerializer registers a listener for the same reason. // // flatMap consumes each inner iterator fully before building the next, so at most one batch // is open at a time and tracking the current one is enough. close() is idempotent, so // closing one that already released itself is a no-op. - @volatile var current: ColumnReaders = null + @volatile var current: ProjectedBatch = null Option(TaskContext.get()).foreach { tc => tc.addTaskCompletionListener[Unit] { _ => - val readers = current + val open = current current = null - if (readers != null) { - readers.close() + if (open != null) { + open.close() } } } @@ -348,9 +374,9 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // Nothing to decode: the row count is the whole answer, and it is already here. Iterator.single(new ColumnarBatch(Array.empty[ColumnVector], cb.numRows)) } else { - val readers = new ColumnReaders(indices.map(i => cb.columns(i)), cb.numRows) - current = readers - readers.batches + val projected = new ProjectedBatch(cb, arrowFields, indices) + current = projected + projected.batches } case other => @@ -360,78 +386,56 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { } } - // Decodes one selected column stream apiece and stitches the results back into a single batch. - // - // Each stream is self-contained, so the columns a scan did not select are never inflated. The - // decoded vectors stay owned by their readers: closing them releases the batch, which is why - // this yields a single-element iterator that closes on exhaustion, matching what - // ArrowReaderIterator did when the payload was one stream. - private class ColumnReaders(buffers: Array[ChunkedByteBuffer], numRows: Int) { - // decodeBatches opens a reader and eagerly decodes its first batch, so it allocates. If a - // later column throws, the readers already opened here are unreachable: the task-completion - // listener cannot release them because `current` is only assigned once this constructor - // returns, so they would leak off-heap for the life of the executor. - private val readers: Array[Iterator[ColumnarBatch]] = { - val opened = new Array[Iterator[ColumnarBatch]](buffers.length) - var i = 0 + /** + * Loads the projected columns of one cached batch into Arrow vectors. + * + * The schema is rebuilt from the cached relation's attributes rather than read from the + * payload, which stores none. `CachedBatchIpc.readProjected` then materializes only the + * selected columns' buffers, so the rest are never copied out of the cached bytes or + * decompressed. + * + * The decoded vectors stay owned by this object: closing it releases the batch, which is why + * this yields a single-element iterator that closes on exhaustion. + */ + private class ProjectedBatch( + cached: CometCachedBatch, + arrowFields: Seq[Field], + indices: Array[Int]) { + + // Allocated before anything can throw, so that a failure below has a root to release. + private val root = VectorSchemaRoot.create( + new Schema(indices.map(arrowFields).toSeq.asJava), + CometArrowAllocator) + private var closed = false + + // Loading happens during construction, so `batches` below can hand out the root directly. + try { + val recordBatch = + CachedBatchIpc.readProjected(cached.bytes, arrowFields, indices, CometArrowAllocator) try { - while (i < buffers.length) { - opened(i) = Utils.decodeBatches(buffers(i), "CometCache") - i += 1 - } - } catch { - case NonFatal(e) => - var j = 0 - while (j < i) { - opened(j) match { - case reader: ArrowReaderIterator => - try reader.close() - catch { case NonFatal(closeError) => e.addSuppressed(closeError) } - case _ => () - } - j += 1 - } - throw e + CachedBatchIpc.loaderFor(root).load(recordBatch) + } finally { + recordBatch.close() + } + // A cached batch's columns all cover the same rows. Check rather than trust: a mismatch + // would otherwise build a batch whose columns disagree with the row count recorded beside + // them, which reads as corrupt data far from here. + if (root.getRowCount != cached.numRows) { + throw new IllegalStateException( + s"Cached batch decoded ${root.getRowCount} rows, expected ${cached.numRows}") } - opened + } catch { + case NonFatal(e) => + try root.close() + catch { case NonFatal(closeError) => e.addSuppressed(closeError) } + throw e } - private var closed = false def close(): Unit = synchronized { if (!closed) { closed = true - readers.foreach { - case reader: ArrowReaderIterator => reader.close() - case _ => () - } - } - } - - private def assemble(): ColumnarBatch = { - val columns = new Array[ColumnVector](readers.length) - var i = 0 - while (i < readers.length) { - val reader = readers(i) - if (!reader.hasNext) { - throw new IllegalStateException( - s"Cached column stream $i of ${readers.length} decoded to no batch") - } - val decoded = reader.next() - // Each stream holds exactly one single-column record batch, and every column of a cached - // batch covers the same rows. Check rather than trust: a mismatch would otherwise build a - // batch whose columns disagree on length, which reads as corrupt data far from here. - if (decoded.numCols() != 1) { - throw new IllegalStateException( - s"Cached column stream $i decoded to ${decoded.numCols()} columns, expected 1") - } - if (decoded.numRows() != numRows) { - throw new IllegalStateException( - s"Cached column stream $i decoded ${decoded.numRows()} rows, expected $numRows") - } - columns(i) = decoded.column(0) - i += 1 + root.close() } - new ColumnarBatch(columns, numRows) } def batches: Iterator[ColumnarBatch] = new Iterator[ColumnarBatch] { @@ -451,7 +455,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { throw new NoSuchElementException } emitted = true - assemble() + NativeUtil.rootAsBatch(root) } } } @@ -467,24 +471,26 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { fallback.convertInternalRowToCachedBatch(input, schema, storageLevel, conf) } else { val batchSize = conf.columnBatchSize + val codecName = CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC.get(conf) + val zstdLevel = CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_ZSTD_LEVEL.get(conf) input.mapPartitions { rows => val iter = CometArrowConverters.rowToArrowBatchIter( rows, Utils.fromAttributes(schema), batchSize, - // NATIVE_TIMEZONE ("UTC"), not conf.sessionLocalTimeZone, so both write paths produce - // the same physical format: the columnar path above already encodes with - // NATIVE_TIMEZONE. Unlike Spark's Arrow cache, whose RecordBatch is deliberately - // schema-less, CometCachedBatch stores a full IPC stream including the schema, so a - // session-local label would persist the writing session's mutable timezone into cached - // data. This is a label only: Spark's internal timestamp representation is micros since - // the Unix epoch regardless of session timezone, so no values are converted. It also - // matches Comet's native schema, avoiding a cast at the native boundary. + // NATIVE_TIMEZONE ("UTC"), not conf.sessionLocalTimeZone. The payload stores no schema, + // so the read path rebuilds one with toArrowSchema(cacheSchema, NATIVE_TIMEZONE); a + // write that labelled its timestamps with the writing session's timezone would be read + // back under a different label. Both write paths therefore have to agree on this, and + // the columnar path above encodes with NATIVE_TIMEZONE too. This is a label only: + // Spark's internal timestamp representation is micros since the Unix epoch regardless + // of session timezone, so no values are converted. It also matches Comet's native + // schema, avoiding a cast at the native boundary. CometArrowStream.NATIVE_TIMEZONE, CometArrowAllocator) - encodeBatches(iter, schema) + encodeBatches(iter, schema, codecName, zstdLevel) } } } @@ -551,6 +557,9 @@ object ArrowCachedBatchSerializer { */ def kryoClasses: Seq[Class[_]] = Seq( classOf[CometCachedBatch], + // The payload itself. Kryo registers Array[Byte] by default, but registering it here is what + // keeps that true if the payload type ever changes again. + classOf[Array[Byte]], // The statistics row, whose values are bounds in Spark's internal representation: boxed // primitives, which Kryo registers by default, plus UTF8String and Decimal, which it does not. // A Decimal above Long precision holds a scala.math.BigDecimal, which Chill's Scala registrar 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 new file mode 100644 index 00000000000..55215ef548a --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala @@ -0,0 +1,409 @@ +/* + * 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 java.io.{ByteArrayInputStream, ByteArrayOutputStream} +import java.nio.channels.Channels + +import scala.collection.mutable +import scala.jdk.CollectionConverters._ +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, 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} +import org.apache.arrow.vector.ipc.message.{ArrowBodyCompression, ArrowFieldNode, ArrowRecordBatch, MessageSerializer} +import org.apache.arrow.vector.types.pojo.{ArrowType, Field} +import org.apache.spark.SparkException +import org.apache.spark.sql.comet.util.Utils +import org.apache.spark.sql.vectorized.ColumnarBatch + +/** + * The on-disk shape of a `CometCachedBatch` payload, and the two operations over it. + * + * A cached batch is one encapsulated Arrow IPC RecordBatch message followed by its body, with no + * Schema message and no end-of-stream marker. The schema is not stored because the reader already + * has it: `InMemoryRelation` knows the cached relation's attributes, and `Utils.toArrowSchema` + * maps them to exactly the fields the writer unloaded. Leaving it out saves a schema message per + * cached batch, which for a wide relation cached in many batches is a large share of the payload + * that is not data. + * + * Compression is applied by Arrow per buffer rather than by wrapping the whole payload in a Spark + * `CompressionCodec`. That is what makes projection cheap: the message metadata records every + * buffer's offset and length within the body, so [[readProjected]] can copy out only the buffers + * of the columns a scan selected and let `VectorLoader` decompress just those. A whole-payload + * codec would have to inflate everything before any column could be read. + */ +private[comet] object CachedBatchIpc { + + /** + * The Arrow compression codec named by `spark.comet.exec.inMemoryCache.compression.codec`. + * + * Only the write path consults the config. A batch records which codec compressed it, so the + * read path looks the codec up from the batch itself and keeps reading data cached before the + * config changed. + */ + def compressionCodec(codecName: String, zstdLevel: Int): CompressionCodec = codecName match { + case "none" => NoCompressionCodec.INSTANCE + // 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. + case other => + throw new SparkException( + s"Unsupported Arrow compression codec for Comet's cache: $other. " + + "Supported values: none, zstd") + } + + /** + * Serialize `batch` into one encapsulated IPC RecordBatch message. + * + * Returns the message bytes 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. + * + * 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 + * attributes never carries one, so a dictionary-encoded column has nowhere to record either its + * index type or the dictionary itself. Comet's native scans do produce such columns, so this is + * a real path, not a defensive one. + * + * As in `Utils.serializeBatches`, `batch`'s vectors are cleared once written, so callers gather + * anything they need from the batch (statistics, for instance) before calling this. + */ + def serialize( + batch: ColumnarBatch, + codec: CompressionCodec, + allocator: BufferAllocator): (Array[Byte], Array[Long]) = { + val (vectors, hydrated) = hydrateDictionaries(batch, allocator) + try { + val root = new VectorSchemaRoot(vectors.asJava) + // A batch of zero columns carries only a row count, which a VectorSchemaRoot cannot infer + // without vectors to measure. + if (vectors.isEmpty) { + root.setRowCount(batch.numRows()) + } + + // alignBuffers=true matches the 8-byte buffer alignment readProjected reproduces when it + // repacks the selected buffers. + val unloader = new VectorUnloader(root, true, codec, true) + val recordBatch = unloader.getRecordBatch + try { + val fields = vectors.map(_.getField) + // Serializing consumes the batch, as it does in Utils.serializeBatches. The record batch + // holds its own buffers by now -- compressed copies, or retained references when the codec + // is none -- so releasing the vectors here does not touch it. getField still answers + // afterwards: clearing releases buffers, not the schema. + // + // Not load bearing for memory: the plan that produced the batch releases its vectors + // either way, and dropping this line leaks nothing. It is here because serializeBatches + // does the same, so both writers leave a batch they were handed in the same state. + root.clear() + + val out = new ByteArrayOutputStream() + val channel = new WriteChannel(Channels.newChannel(out)) + MessageSerializer.serialize(channel, recordBatch) + (out.toByteArray, columnSizes(fields, recordBatch)) + } finally { + recordBatch.close() + } + } finally { + // Only the vectors this method allocated. The rest belong to the input batch. + hydrated.foreach(v => + try v.close() + catch { case NonFatal(_) => () }) + } + } + + /** + * Read an encapsulated IPC RecordBatch message, materializing off-heap only the buffers of the + * requested top-level columns. + * + * The body is a flat, depth-first sequence of buffers in schema order, so each top-level column + * owns a contiguous run of buffers whose length is [[fieldBufferCount]]; field nodes and + * variadic buffer counts run in the same order. The selected columns' bytes are copied into a + * single off-heap allocation, each buffer 8-byte aligned exactly as Arrow's IPC body lays them + * out, and the returned batch's buffers are windows into it -- one allocation, no per-buffer + * bookkeeping. + * + * Only the selected buffers are ever decompressed. A buffer's recorded (offset, length) covers + * its on-body bytes including the uncompressed-length prefix, so a copied window is exactly + * what the writer emitted; the columns that were not selected are never read, let alone + * inflated. The copied windows are then decompressed in one pass -- see [[decompressed]] for + * why that is not left to `VectorLoader` -- so what comes back is an uncompressed batch. + * + * The returned batch owns its buffers; the caller closes it. + */ + def readProjected( + data: Array[Byte], + schemaFields: Seq[Field], + selectedIndices: Array[Int], + allocator: BufferAllocator): ArrowRecordBatch = { + val readChannel = new ReadChannel(Channels.newChannel(new ByteArrayInputStream(data))) + // Reads the message metadata only. The body stays in `data` and is copied selectively below. + val metadata = MessageSerializer.readMessage(readChannel) + if (metadata == null) { + throw new SparkException("Unexpected end of input reading a Comet cached batch") + } + val batch = + metadata.getMessage.header(new FlatBufRecordBatch()).asInstanceOf[FlatBufRecordBatch] + // serialize writes exactly [encapsulated message][body] and nothing after it, so the body is + // the tail of `data`. + val bodyStart = data.length - metadata.getMessageBodyLength.toInt + + val compression = + if (batch.compression() == null) NoCompressionCodec.DEFAULT_BODY_COMPRESSION + else new ArrowBodyCompression(batch.compression().codec(), batch.compression().method()) + + val nodeStarts = schemaFields.scanLeft(0)(_ + fieldNodeCount(_)).toArray + val bufferStarts = schemaFields.scanLeft(0)(_ + fieldBufferCount(_)).toArray + val variadicStarts = schemaFields.scanLeft(0)(_ + fieldVariadicCount(_)).toArray + val hasVariadic = batch.variadicBufferCountsLength() > 0 + + // The selected columns' field nodes, buffer indices and variadic counts, in output order. + val nodes = new java.util.ArrayList[ArrowFieldNode]() + val bufferIndices = mutable.ArrayBuffer.empty[Int] + val variadicCounts = new java.util.ArrayList[java.lang.Long]() + selectedIndices.foreach { i => + val field = schemaFields(i) + val nodeStart = nodeStarts(i) + (nodeStart until nodeStart + fieldNodeCount(field)).foreach { j => + val node = batch.nodes(j) + nodes.add(new ArrowFieldNode(node.length(), node.nullCount())) + } + val bufferStart = bufferStarts(i) + (bufferStart until bufferStart + fieldBufferCount(field)).foreach(bufferIndices += _) + if (hasVariadic) { + val variadicStart = variadicStarts(i) + (variadicStart until variadicStart + fieldVariadicCount(field)) + .foreach(j => variadicCounts.add(batch.variadicBufferCounts(j))) + } + } + + val layout = bufferIndices.map { j => + val buffer = batch.buffers(j) + (buffer.offset(), buffer.length()) + } + val alignedSizes = layout.map { case (_, length) => ((length + 7) / 8) * 8 } + // allocator.buffer(0) is legal but yields a buffer no window can be sliced from, and an + // all-empty projection (every selected column a NullVector, say) would ask for exactly that. + val body = allocator.buffer(math.max(alignedSizes.sum, 1L)) + val compressedBatch = + try { + val buffers = new java.util.ArrayList[ArrowBuf]() + var position = 0L + layout.indices.foreach { k => + val (sourceOffset, length) = layout(k) + if (length > 0) { + body.setBytes(position, data, bodyStart + sourceOffset.toInt, length.toInt) + } + val window = body.slice(position, length) + window.writerIndex(length) + buffers.add(window) + position += alignedSizes(k) + } + new ArrowRecordBatch( + batch.length().toInt, + nodes, + buffers, + compression, + variadicCounts, + false) + } catch { + case NonFatal(e) => + body.close() + throw e + } + + // The constructor retained each window; slice() alone does not. Dropping `body`'s own + // reference leaves the batch as sole owner of the one allocation, so closing the batch below + // is what frees it -- and closing `body` again here would drive its reference count negative. + body.close() + try decompressed(compressedBatch, allocator) + finally compressedBatch.close() + } + + /** + * The same record batch with every buffer decompressed, as a new batch the caller owns. + * + * `VectorLoader` would do this itself, but arrow-java 18.3.0 leaks on the failure path: + * `VectorLoader.loadBuffers` decompresses a field's buffers into a local list and only releases + * them after the whole field has loaded, so if one buffer of a field fails to decompress, every + * buffer of that field decompressed before it is unreachable and never freed. A string column + * is enough to reach it -- its offsets buffer decompresses, then its data buffer throws -- so a + * single corrupt cached batch leaks off-heap for the life of the executor. Doing the + * decompression here keeps every allocation reachable from this method's own error path. + * + * Buffers are retained before decompressing rather than after, which is the other half of the + * difference. `decompress` consumes a reference to its input on the paths where it allocates, + * so retaining afterwards leaves the reference stranded if it throws -- and, when a batch has a + * single buffer, drops the shared body to zero references and frees it before the retain that + * was meant to protect it. + */ + private def decompressed( + batch: ArrowRecordBatch, + allocator: BufferAllocator): ArrowRecordBatch = { + // getCodec is the raw IPC byte; the factory keys off the enum. Both sides of the comparison + // below have to be CodecType: NoCompressionCodec.COMPRESSION_TYPE is the byte -1, and Scala + // compares a CodecType against it by universal equality, which is quietly always unequal. + val codecType = + CompressionUtil.CodecType.fromCompressionType(batch.getBodyCompression.getCodec) + val compressed = codecType != CompressionUtil.CodecType.NO_COMPRESSION + val codec: CompressionCodec = + if (compressed) CommonsCompressionFactory.INSTANCE.createCodec(codecType) + else NoCompressionCodec.INSTANCE + + val buffers = new java.util.ArrayList[ArrowBuf]() + try { + batch.getBuffers.asScala.foreach { buffer => + buffer.getReferenceManager.retain() + val plain = + try { + // An empty buffer carries no compressed length prefix to read. + if (compressed && buffer.writerIndex() > 0) codec.decompress(allocator, buffer) + else buffer + } catch { + case NonFatal(e) => + buffer.getReferenceManager.release() + throw e + } + buffers.add(plain) + } + + val result = new ArrowRecordBatch( + batch.getLength, + batch.getNodes, + buffers, + NoCompressionCodec.DEFAULT_BODY_COMPRESSION, + batch.getVariadicBufferCounts, + false) + // The constructor retained each buffer, so drop the references held here. + buffers.asScala.foreach(_.close()) + result + } catch { + case NonFatal(e) => + buffers.asScala.foreach { buffer => + try buffer.close() + catch { case NonFatal(closeError) => e.addSuppressed(closeError) } + } + throw e + } + } + + /** + * A `VectorLoader` for what [[readProjected]] returns. + * + * No compression factory: [[readProjected]] has already decompressed every buffer, so the + * loader only ever sees a batch marked uncompressed. + */ + def loaderFor(root: VectorSchemaRoot): org.apache.arrow.vector.VectorLoader = + new org.apache.arrow.vector.VectorLoader(root) + + /** + * The on-body compressed size of each top-level column. + * + * 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. + */ + private def columnSizes(fields: Seq[Field], recordBatch: ArrowRecordBatch): Array[Long] = { + val buffers = recordBatch.getBuffersLayout + val starts = fields.scanLeft(0)(_ + fieldBufferCount(_)).toArray + fields.indices.map { i => + (starts(i) until starts(i) + fieldBufferCount(fields(i))) + .map(j => buffers.get(j).getSize) + .sum + }.toArray + } + + /** + * Replace every dictionary-encoded column of `batch` with its decoded form. + * + * Returns the vectors to write and, separately, the ones allocated here so the caller can close + * exactly those. Columns that needed no decoding are returned as they are and stay owned by + * `batch`. + */ + private def hydrateDictionaries( + batch: ColumnarBatch, + allocator: BufferAllocator): (Seq[FieldVector], Seq[ValueVector]) = { + val hydrated = mutable.ArrayBuffer.empty[ValueVector] + try { + val vectors = + Utils.getBatchFieldVectorsWithProviders(batch).map { case (vector, providerOpt) => + val encoding = vector.getField.getDictionary + if (encoding == null) { + vector + } else { + val dictionary = providerOpt.map(_.lookup(encoding.getId)).orNull + if (dictionary == null) { + throw new SparkException( + s"Column ${vector.getField.getName} is dictionary encoded with ID " + + s"${encoding.getId}, but no dictionary with that ID was provided") + } + val decoded = DictionaryEncoder.decode(vector, dictionary, allocator) + hydrated += decoded + decoded.asInstanceOf[FieldVector] + } + } + (vectors, hydrated.toSeq) + } catch { + case NonFatal(e) => + hydrated.foreach(v => + try v.close() + catch { case NonFatal(closeError) => e.addSuppressed(closeError) }) + throw e + } + } + + /** + * Number of Arrow buffers a field occupies in a RecordBatch body, including every descendant, + * in the depth-first order `VectorLoader` consumes them. The type's own count covers its + * validity and offset/data buffers; each child contributes its whole subtree. + */ + private def fieldBufferCount(field: Field): Int = + TypeLayout.getTypeBufferCount(field.getType) + + field.getChildren.asScala.map(fieldBufferCount).sum + + /** Number of field nodes a field occupies: itself plus every descendant. */ + private def fieldNodeCount(field: Field): Int = + 1 + field.getChildren.asScala.map(fieldNodeCount).sum + + /** + * Number of variadic buffer counts a field contributes, one per view-type buffer, recursively. + * + * Only Utf8View and BinaryView carry one. Comet's cache never writes view vectors today, but + * the span arithmetic above has to stay correct if that changes. + */ + private def fieldVariadicCount(field: Field): Int = { + val own = field.getType match { + case _: ArrowType.Utf8View | _: ArrowType.BinaryView => 1 + case _ => 0 + } + own + field.getChildren.asScala.map(fieldVariadicCount).sum + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index 769d8058de5..f418c31a626 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -270,45 +270,12 @@ object Utils extends CometTypeShim with Logging { } /** - * Serializes each column of `batch` into its own compressed Arrow IPC stream, in column order. - * - * [[serializeBatches]] writes one stream covering every column, so a reader has to inflate all - * of them before it can project. Comet's in-memory cache stores columns separately instead, so - * a scan decodes only the ones it selected. Each stream is self-contained, including its schema - * and any dictionaries the column needs. - * - * The row count is not recoverable from the result when `batch` has no columns, so callers keep - * it alongside. As with [[serializeBatches]], the batch's vectors are cleared once written. - */ - def serializeBatchColumns(batch: ColumnarBatch): Array[ChunkedByteBuffer] = { - val codec = CompressionCodec.createCodec(SparkEnv.get.conf) - - // Each column is written with the provider it was decoded with, not the batch's first one: - // columns decoded from separate streams have independent dictionary ID namespaces. - getBatchFieldVectorsWithProviders(batch).map { case (fieldVector, providerOpt) => - val provider = providerOpt.getOrElse(new CDataDictionaryProvider) - val cbbos = new ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate) - val out = new DataOutputStream(codec.compressedOutputStream(cbbos)) - - val root = new VectorSchemaRoot(Seq(fieldVector).asJava) - val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out)) - writer.start() - writer.writeBatch() - root.clear() - writer.close() - - cbbos.toChunkedByteBuffer - }.toArray - } - - /** - * The classes that carry the output of [[serializeBatches]] and [[serializeBatchColumns]] out - * of Comet, for Kryo registration by [[org.apache.comet.CometKryoRegistrator]]. + * The classes that carry the output of [[serializeBatches]] out of Comet, for Kryo registration + * by [[org.apache.comet.CometKryoRegistrator]]. * * Spark registers `ChunkedByteBuffer` itself but not an array of them, and * `CometBroadcastExchangeExec` broadcasts exactly that array, so a native broadcast fails under - * `spark.kryo.registrationRequired=true` whichever Comet features are enabled. Comet's cache - * format stores one buffer per column and so needs the same registrations. + * `spark.kryo.registrationRequired=true` whichever Comet features are enabled. */ def arrowBytesKryoClasses: Seq[Class[_]] = Seq( classOf[ChunkedByteBuffer], diff --git a/spark/src/main/spark-3.x/org/apache/comet/shims/CometTypeShim.scala b/spark/src/main/spark-3.x/org/apache/comet/shims/CometTypeShim.scala index b71476c3dd1..a6b34b74e7d 100644 --- a/spark/src/main/spark-3.x/org/apache/comet/shims/CometTypeShim.scala +++ b/spark/src/main/spark-3.x/org/apache/comet/shims/CometTypeShim.scala @@ -22,6 +22,7 @@ package org.apache.comet.shims import scala.annotation.nowarn import org.apache.spark.sql.types.{DataType, StructType} +import org.apache.spark.unsafe.types.{ByteArray, UTF8String} trait CometTypeShim { @nowarn // Spark 4 feature; stubbed to false in Spark 3.x for compatibility. @@ -41,4 +42,15 @@ trait CometTypeShim { @nowarn // Spark 4.1 feature; TimeType doesn't exist in Spark 3.x. def isTimeType(dt: DataType): Boolean = false + + /** + * Compare two strings under the collation of `dt`, which must be a `StringType`. + * + * Spark 3.x has no collations, so every string comparison is byte order. Callers that record + * comparable bounds (Comet's cache statistics, for instance) use this so the ordering they + * store is the one Spark's own comparison would produce. + */ + @nowarn // Collation is a Spark 4 feature; on 3.x every StringType compares as bytes. + def compareStrings(left: UTF8String, right: UTF8String, dt: DataType): Int = + ByteArray.compareBinary(left.getBytes, right.getBytes) } diff --git a/spark/src/main/spark-4.x/org/apache/comet/shims/CometTypeShim.scala b/spark/src/main/spark-4.x/org/apache/comet/shims/CometTypeShim.scala index f48955a7da5..71e72dd2de4 100644 --- a/spark/src/main/spark-4.x/org/apache/comet/shims/CometTypeShim.scala +++ b/spark/src/main/spark-4.x/org/apache/comet/shims/CometTypeShim.scala @@ -21,6 +21,7 @@ package org.apache.comet.shims import org.apache.spark.sql.execution.datasources.VariantMetadata import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StringType, StructType, VariantType} +import org.apache.spark.unsafe.types.UTF8String trait CometTypeShim { // A `StringType` carries collation metadata in Spark 4.0. Only non-default (non-UTF8_BINARY) @@ -64,4 +65,14 @@ trait CometTypeShim { dt.getClass.getSimpleName.startsWith("TimeType") def hasCollationSupport: Boolean = true + + /** + * Compare two strings under the collation of `dt`, which must be a `StringType`. + * + * `semanticCompare` is the comparison Spark's own expressions use for the type, so bounds + * recorded with it order the same way a predicate over the column does. For the default + * UTF8_BINARY collation it is byte order, which is what Spark 3.x always does. + */ + def compareStrings(left: UTF8String, right: UTF8String, dt: DataType): Int = + left.semanticCompare(right, dt.asInstanceOf[StringType].collationId) } 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 a7a9a6c8590..3464e5e3ddc 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -29,6 +29,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, Attribute, Expression, Gr import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} import org.apache.spark.sql.comet.CometInMemoryTableScanExec import org.apache.spark.sql.comet.execution.arrow.CometCachedBatchHelper +import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation} import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} @@ -724,13 +725,14 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - test("Comet in-memory cache prunes only on columns that have bounds") { + test("Comet in-memory cache prunes on collated string columns") { assume(isSpark40Plus, "collated string types require Spark 4.0+") withNativeCache { - // A collated StringType does not match `case StringType` in the serializer's bounds - // tracking, so its lower and upper bounds stay null. Spark still builds a partition filter - // for it because a collated string literal is an AtomicType, and comparing against null - // bounds prunes every batch. Without the buildFilter guard this query returns no rows. + // Bounds for a collated column are recorded with that collation's own comparison, which is + // the same ordering the partition filter Spark generates over the column uses. Tracking + // bounds only for the bare `StringType` object would leave a collated column's bounds null, + // and a comparison against null bounds prunes every batch, so the query below would return + // no rows at all rather than merely losing the pruning. spark .sql("SELECT id, CAST(id AS STRING) COLLATE UTF8_LCASE AS s FROM range(100)") .createOrReplaceTempView("collated_cache") @@ -747,12 +749,57 @@ class CometInMemoryCacheSuite extends CometTestBase { assert( spark.sql("SELECT id FROM collated_cache WHERE s >= '5'").collect().length == expected) + // UTF8_LCASE compares case-insensitively, so bounds recorded under it have to as well: a + // batch whose values all sort above 'A' under byte order still contains matches for a + // predicate that is looking for lower-case letters. + spark + .sql( + "SELECT id, CAST(concat('X', cast(id as string)) AS STRING) COLLATE UTF8_LCASE AS s " + + "FROM range(100)") + .createOrReplaceTempView("collated_case_cache") + spark.catalog.cacheTable("collated_case_cache") + spark.table("collated_case_cache").count() + assert( + spark.sql("SELECT id FROM collated_case_cache WHERE s = 'x1'").collect().length == 1, + "a case-insensitive match must survive pruning") + // Null-count based pruning stays available for columns without bounds. assert( spark.sql("SELECT id FROM collated_cache WHERE s IS NOT NULL").collect().length == 100) } } + test("Comet in-memory cache prunes only on columns that have bounds") { + withNativeCache { + // Binary has no bounds recorded, so its lower and upper stay null. Spark would still build + // a partition filter for it, and comparing against null bounds prunes every batch, so + // without the buildFilter guard this query returns no rows. + spark + .sql("SELECT id, CAST(CAST(id AS STRING) AS BINARY) AS b FROM range(100)") + .createOrReplaceTempView("binary_cache") + spark.catalog.cacheTable("binary_cache") + spark.table("binary_cache").count() + + assert( + cachedBatchTypes("binary_cache").sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) + + assert( + spark + .sql("SELECT id FROM binary_cache WHERE b >= CAST('5' AS BINARY)") + .collect() + .length == + spark + .sql("SELECT id FROM range(100) WHERE CAST(CAST(id AS STRING) AS BINARY) >= " + + "CAST('5' AS BINARY)") + .collect() + .length) + + // Null-count based pruning stays available for columns without bounds. + assert(spark.sql("SELECT id FROM binary_cache WHERE b IS NOT NULL").collect().length == 100) + } + } + test("Comet in-memory cache is readable when Comet is disabled") { // spark.sql.cache.serializer is static, so the cached format cannot depend on a runtime // config. Disabling Comet must still leave the cached relation readable, including for @@ -986,8 +1033,12 @@ class CometInMemoryCacheSuite extends CometTestBase { SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true") { spark.catalog.clearCache() + // Wide enough that every column's buffers are big enough for Arrow to actually compress + // them. Arrow stores a buffer verbatim when compressing it would not make it smaller, and a + // boolean column of a few hundred rows is a few dozen bytes, which takes that fallback -- + // leaving the corruption the projection tests rely on with nothing to corrupt. spark - .range(0, 500, 1, 2) + .range(0, 8000, 1, 2) .selectExpr( "id", "id % 100 AS k", @@ -997,7 +1048,7 @@ class CometInMemoryCacheSuite extends CometTestBase { "cast(id % 2 = 0 as boolean) AS flag") .createOrReplaceTempView("projection_cache") spark.catalog.cacheTable("projection_cache") - assert(spark.table("projection_cache").count() == 500) + assert(spark.table("projection_cache").count() == 8000) assert( cachedBatchTypes("projection_cache").sameElements( Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch"))) @@ -1032,67 +1083,149 @@ class CometInMemoryCacheSuite extends CometTestBase { .sum } - test("Comet in-memory cache stores one stream per column") { + /** + * Run `f`, require it to fail, and require the failure to be the decode error itself. + * + * The read path allocates an off-heap body, hands it to a record batch that takes its own + * references, and drops its own. A cleanup path that then releases the body a second time + * drives its reference count negative, and the reference-count error replaces the decode + * failure that caused it -- leaving a plain `intercept[Exception]` green while the user sees a + * error that says nothing about their corrupt cache. + */ + private def interceptDecodeFailure(f: => Unit): Throwable = { + val thrown = intercept[Exception](f) + val chain = + Iterator.iterate(thrown: Throwable)(_.getCause).takeWhile(_ != null).take(20).toSeq + assert( + !chain.exists { t => + t.getClass.getName.contains("IllegalReferenceCount") || + Option(t.getMessage).exists(m => m.contains("RefCnt") || m.contains("refCnt")) + }, + s"the decode failure must surface as itself, not as a reference-count error: $thrown") + thrown + } + + test("Comet in-memory cache round-trips under every compression codec") { + // 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 => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC.key -> codec, + "spark.comet.sparkToColumnar.enabled" -> "true") { + + spark.catalog.clearCache() + val view = s"codec_cache_$codec" + spark + .range(0, 4000, 1, 2) + .selectExpr( + "id", + "cast(id as double) / 3 AS d", + "concat('s_', cast(id as string)) AS s", + "cast(id % 2 = 0 as boolean) AS flag") + .createOrReplaceTempView(view) + spark.catalog.cacheTable(view) + + assert( + cachedBatchTypes(view).sameElements( + Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch")), + s"codec $codec should still store CometCachedBatch") + + // A full read, a projected read (the buffer-selection path), and a row count that decodes + // nothing -- the three shapes the read path distinguishes. + checkSparkAnswer(spark.sql(s"SELECT * FROM $view")) + checkSparkAnswer(spark.sql(s"SELECT s FROM $view WHERE id >= 3990")) + assert(spark.sql(s"SELECT count(*) FROM $view").collect()(0).getLong(0) == 4000) + // Pruning reads the statistics rather than the payload, so exercise it too. + assert(spark.sql(s"SELECT id FROM $view WHERE id >= 3990").collect().length == 10) + + spark.catalog.clearCache() + } + } + } + + test("Comet in-memory cache stores no schema message per cached batch") { + // The reader rebuilds the schema from the cached relation's attributes, so storing one in + // every batch would repeat the same bytes for as many batches as the relation was cached in. withProjectionCache { (relation, batches) => assert(batches.nonEmpty) + val cacheSchema = Utils.fromAttributes(relation.output) batches.foreach { batch => assert( - CometCachedBatchHelper.numColumnStreams(batch) == relation.output.length, - "a cached batch must hold one independently decodable stream per cached column") + !CometCachedBatchHelper.hasSchemaMessage(batch), + "a cached batch must begin with its record batch, not a schema message") + val sizes = CometCachedBatchHelper.columnSizes(batch, cacheSchema) assert( - CometCachedBatchHelper.columnStreamSizes(batch).forall(_ > 0), - "every column stream must carry data") + sizes.length == relation.output.length, + "every cached column must own a run of buffers in the payload") + assert(sizes.forall(_ > 0), "every cached column must carry data") } } } test("Comet in-memory cache decodes only the projected columns") { - // Timings would be a weak assertion here, so this corrupts the streams the read must not - // touch. Reading still has to succeed, which it only can if those streams were never - // inflated. The second half checks the corruption is detectable at all, so that the first - // half cannot pass just because the bad bytes decode silently to nothing. + // Timings would be a weak assertion here, so this scrambles the compressed bytes of the + // columns the read must not touch, leaving every other byte of the payload identical. + // Reading still has to succeed, which it only can if those columns' buffers were never copied + // out of the payload and handed to the decompressor. The second half checks the corruption is + // detectable at all, so the first half cannot pass just because the bad bytes decode silently + // to nothing. withProjectionCache { (relation, batches) => + val cacheSchema = Utils.fromAttributes(relation.output) val selectedIdx = 1 val selected = Seq(relation.output(selectedIdx)) + relation.output.indices.foreach { i => + assert( + batches.forall(b => CometCachedBatchHelper.columnIsCompressed(b, cacheSchema, i)), + s"column $i is not stored compressed, so corrupting it would prove nothing") + } + relation.output.indices.filter(_ != selectedIdx).foreach { i => - batches.foreach(b => CometCachedBatchHelper.corruptColumnStream(b, i)) + batches.foreach(b => CometCachedBatchHelper.corruptColumn(b, cacheSchema, i)) } assert( - decodedRowCount(relation, batches, selected) == 500, - "reading one column must not decode the other five") + decodedRowCount(relation, batches, selected) == 8000, + "reading one column must not decompress the other five") - batches.foreach(b => CometCachedBatchHelper.corruptColumnStream(b, selectedIdx)) - intercept[Exception] { + batches.foreach(b => CometCachedBatchHelper.corruptColumn(b, cacheSchema, selectedIdx)) + interceptDecodeFailure { decodedRowCount(relation, batches, selected) } } } test("Comet in-memory cache decodes no columns for a row-count-only read") { - // SELECT count(*) selects no columns. Every stream is corrupted, so the read can only succeed - // by decoding none of them and answering from the row count the cached batch already carries. + // SELECT count(*) selects no columns. Every column's bytes are corrupted, so the read can + // only succeed by touching none of them and answering from the row count the cached batch + // already carries beside the payload. withProjectionCache { (relation, batches) => + val cacheSchema = Utils.fromAttributes(relation.output) relation.output.indices.foreach { i => - batches.foreach(b => CometCachedBatchHelper.corruptColumnStream(b, i)) + batches.foreach(b => CometCachedBatchHelper.corruptColumn(b, cacheSchema, i)) } - assert(decodedRowCount(relation, batches, Seq.empty) == 500) + assert(decodedRowCount(relation, batches, Seq.empty) == 8000) } } test("Comet in-memory cache records per-column sizes in its statistics") { - // SimpleMetricsCachedBatch reserves a fifth field per column for its size. Each column is now - // its own stream, so the real size is known and must be reported rather than left at zero. - withProjectionCache { (_, batches) => + // 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. + withProjectionCache { (relation, batches) => + val cacheSchema = Utils.fromAttributes(relation.output) batches.foreach { batch => - val sizes = CometCachedBatchHelper.columnStreamSizes(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 $i should report its own stream size in the statistics row") + s"column $i should report the stored size of its own buffers in the statistics row") } } } @@ -1169,11 +1302,12 @@ class CometInMemoryCacheSuite extends CometTestBase { test( "Comet in-memory cache re-encodes a decoded batch whose columns have separate dictionaries") { - // Each cached column is decoded from its own stream, so dictionary-backed columns come back - // with independent providers whose IDs collide. Re-encoding such a batch with only the first - // column's provider cannot resolve the later columns' dictionary IDs. Spark's columnar Union - // hands decoded cached batches straight back to this serializer, so caching a union of a - // cached relation exercises exactly that. + // Spark's columnar Union hands decoded cached batches straight back to this serializer, so + // caching a union of a cached relation re-encodes batches that came out of the cache. The + // cache no longer stores dictionary-encoded columns -- the writer decodes them first -- but + // batches reaching serializeBatches from a shuffle or broadcast still carry independent + // dictionary providers whose IDs collide, and re-encoding one with only the first column's + // provider cannot resolve the later columns' IDs. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", @@ -1248,23 +1382,26 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - test("Comet in-memory cache releases opened readers when a later column fails to decode") { - // A cached batch is several independent Arrow streams and decodeBatches opens each eagerly. - // If a later column throws, the readers already opened are unreachable: the task-completion - // listener cannot release them, because the holder is only published once its constructor - // returns. The failure would then leak off-heap for the life of the executor. + 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 + // that fails to decompress throws between the two, and neither is reachable from anywhere + // else -- the holder is published to the task-completion listener only once its constructor + // returns -- so a failure that does not release them leaks off-heap for the life of the + // executor. withProjectionCache { (relation, batches) => - // Corrupt the second selected column, so the first is opened successfully first. + val cacheSchema = Utils.fromAttributes(relation.output) + // Corrupt the second selected column, so the first is copied out successfully first. val selected = Seq(relation.output(0), relation.output(1)) - batches.foreach(b => CometCachedBatchHelper.corruptColumnStream(b, 1)) + batches.foreach(b => CometCachedBatchHelper.corruptColumn(b, cacheSchema, 1)) val before = CometArrowAllocator.getAllocatedMemory - intercept[Exception] { + interceptDecodeFailure { decodedRowCount(relation, batches, selected) } assert( CometArrowAllocator.getAllocatedMemory == before, - "readers opened before the failure must be released") + "everything allocated before the failure must be released") } } @@ -1272,7 +1409,7 @@ class CometInMemoryCacheSuite extends CometTestBase { * Cache two low-cardinality string columns and hand the test the cached payload. * * The shuffle is what makes this worth its own fixture: its reader hands the cache writer - * dictionary-encoded columns, so each cached column stream carries a dictionary of its own. + * dictionary-encoded columns, which the writer has to decode before storing them. */ private def withDictionaryCache(f: (InMemoryRelation, Array[CachedBatch]) => Unit): Unit = { withSQLConf( @@ -1308,15 +1445,51 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - test("Comet in-memory cache broadcasts a batch whose columns have separate dictionaries") { - // A broadcast of a cache scan re-serializes each decoded batch as one stream covering every - // column, and the writer resolves all of their dictionary IDs against the single provider it - // is handed. The columns were decoded from separate streams, so they arrive carrying separate - // providers: passing any one of them cannot resolve the others. - withDictionaryCache { (relation, batches) => + test("Comet in-memory cache releases its vectors when a column fails after a partial decode") { + // Tighter than the two cases above, and the one that actually catches a leak. A string column + // stores its offsets and its data as separate compressed buffers, so corrupting only the + // second makes the decoder decompress one buffer of the column into a fresh allocation and + // then throw on the next, with the first reachable from nothing the failure path can see. + withProjectionCache { (relation, batches) => + val cacheSchema = Utils.fromAttributes(relation.output) + val stringIdx = 3 + assert(relation.output(stringIdx).dataType.typeName == "string") + batches.foreach(b => + CometCachedBatchHelper.corruptTrailingBuffer(b, cacheSchema, stringIdx)) + + val before = CometArrowAllocator.getAllocatedMemory + interceptDecodeFailure { + decodedRowCount(relation, batches, Seq(relation.output(stringIdx))) + } assert( - CometCachedBatchHelper.columnsAreDictionaryEncoded(batches.head).forall(identity), - "this test is only meaningful over dictionary-encoded cached columns") + CometArrowAllocator.getAllocatedMemory == before, + "a buffer decoded before the failure must be released") + } + } + + test("Comet in-memory cache decodes dictionary-encoded columns before storing them") { + // The payload carries no schema, so it has nowhere to record that a column is dictionary + // encoded, nor the dictionary itself. The reader rebuilds a plain Utf8 field for a string + // column either way, so a writer that stored the index vector as-is would hand the loader + // integer indices to read as strings. Reading the values back correctly is what proves the + // writer decoded them first; a row count alone would not. + withDictionaryCache { (relation, _) => + assert(relation.output.length == 2) + + val df = spark.sql("SELECT s1, s2 FROM dictionary_cache") + checkSparkAnswer(df) + + val distinct = + spark.sql("SELECT DISTINCT s1, s2 FROM dictionary_cache ORDER BY s1, s2").collect() + assert(distinct.length == 12, "3 distinct s1 values by 4 distinct s2 values") + assert(distinct.head.getString(0) == "a_0" && distinct.head.getString(1) == "b_0") + } + } + + test("Comet in-memory cache broadcasts a batch read back from the cache") { + // A broadcast of a cache scan re-serializes each decoded batch through serializeBatches, + // which is a different writer from the one that produced the cached payload. + withDictionaryCache { (relation, _) => assert(relation.output.length == 2) val df = spark.sql( @@ -1326,27 +1499,23 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - test("Comet in-memory cache releases a reader whose own first batch fails to decode") { - // A dictionary-encoded column loads its dictionary before the record batch that indexes into - // it, so a reader can allocate and then fail while opening. Nothing else can release what it - // took: its constructor never returns, so no caller holds the reader it would close, and the - // task-completion listener has not been told about it either. + test("Comet in-memory cache releases its vectors when a column fails part way through") { + // Distinct from the corrupted-column case above: there the compressed bytes are wrong from + // their first byte, so the decompressor rejects them outright. Here the bytes start out + // genuine and only the tail is destroyed, so the failure lands after the loader has already + // begun filling vectors. Nothing else can release them -- the holder's constructor never + // returns, so no caller holds it and the task-completion listener has not been told about it. withDictionaryCache { (relation, batches) => - assert( - CometCachedBatchHelper.columnsAreDictionaryEncoded(batches.head).forall(identity), - "this test is only meaningful over dictionary-encoded cached columns") - - // Enough to take out the end-of-stream marker and bite into the record batch body, so the - // read fails after the dictionary has been loaded rather than before. - batches.foreach(b => CometCachedBatchHelper.truncateColumnStream(b, 0, 64)) + val cacheSchema = Utils.fromAttributes(relation.output) + batches.foreach(b => CometCachedBatchHelper.truncateColumn(b, cacheSchema, 0)) val before = CometArrowAllocator.getAllocatedMemory - intercept[Exception] { + interceptDecodeFailure { decodedRowCount(relation, batches, Seq(relation.output.head)) } assert( CometArrowAllocator.getAllocatedMemory == before, - "a reader that fails while opening must release what it already allocated") + "a read that fails while loading must release what it already allocated") } } 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 959b5590b34..c269dc220c6 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 @@ -85,9 +85,10 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { |WHERE id >= 4500000 AND id < 4750000 """.stripMargin) - // A CometCachedBatch stores each column as its own stream, so a scan decodes only what it - // projected and cost tracks the width of the projection. These three cases span that range - // over one cached relation: no columns, one column, and all six. + // A CometCachedBatch records where each column's buffers sit in its payload, so a scan + // copies out and decompresses only what it projected and cost tracks the width of the + // projection. These three cases span that range over one cached relation: no columns, one + // column, and all six. runCacheBenchmark( "in-memory cache row count only (0 of 6 columns)", s"SELECT count(*) FROM $cacheTable") 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 5548f6dadbd..56c3b949569 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 @@ -19,17 +19,19 @@ package org.apache.spark.sql.comet.execution.arrow -import java.io.{DataInputStream, DataOutputStream} -import java.nio.ByteBuffer +import java.io.ByteArrayInputStream import java.nio.channels.Channels -import org.apache.arrow.vector.ipc.ArrowStreamReader -import org.apache.spark.SparkEnv -import org.apache.spark.io.CompressionCodec -import org.apache.spark.sql.columnar.CachedBatch -import org.apache.spark.util.io.{ChunkedByteBuffer, ChunkedByteBufferOutputStream} +import scala.jdk.CollectionConverters._ -import org.apache.comet.CometArrowAllocator +import org.apache.arrow.flatbuf.{MessageHeader, RecordBatch => FlatBufRecordBatch} +import org.apache.arrow.vector.TypeLayout +import org.apache.arrow.vector.ipc.ReadChannel +import org.apache.arrow.vector.ipc.message.MessageSerializer +import org.apache.arrow.vector.types.pojo.Field +import org.apache.spark.sql.columnar.CachedBatch +import org.apache.spark.sql.comet.util.Utils +import org.apache.spark.sql.types.StructType /** * Test-only access to the internals of `CometCachedBatch`. @@ -37,83 +39,210 @@ import org.apache.comet.CometArrowAllocator * A top-level `private` class in Scala is visible to its own package, so this shim needs no * reflection; it exists so tests outside `org.apache.spark.sql.comet.execution.arrow` can assert * on the cached payload's shape. + * + * The IPC buffer arithmetic below is deliberately re-derived here rather than reused from + * `CachedBatchIpc`. A helper that called into the code under test would inherit any bug in it and + * still agree with itself, so the assertions built on this would pass for the wrong reason. */ object CometCachedBatchHelper { - /** Number of independently decodable column streams in a cached batch. */ - def numColumnStreams(batch: CachedBatch): Int = - batch.asInstanceOf[CometCachedBatch].columns.length + /** The raw cached payload: one encapsulated Arrow IPC record batch message and its body. */ + private def payload(batch: CachedBatch): Array[Byte] = + batch.asInstanceOf[CometCachedBatch].bytes + + /** Stored size of the whole cached batch, in bytes. */ + def payloadSize(batch: CachedBatch): Long = payload(batch).length.toLong + + /** + * Whether the payload begins with a Schema message rather than going straight to the record + * batch. + * + * Comet stores no schema per cached batch -- the reader rebuilds it from the cached relation's + * attributes -- so this is false, and a regression to a self-describing stream would show up + * here rather than only as a footprint number. + */ + def hasSchemaMessage(batch: CachedBatch): Boolean = + readMetadata(payload(batch))._1.headerType() == MessageHeader.Schema + + /** + * The on-body (offset, length) of every Arrow buffer belonging to each top-level column, in + * column order. + */ + def columnBufferRanges(batch: CachedBatch, cacheSchema: StructType): Seq[Seq[(Long, Long)]] = { + val data = payload(batch) + val (_, recordBatch) = readMetadata(data) + val fields = arrowFields(cacheSchema) + val starts = fields.scanLeft(0)(_ + bufferCount(_)).toArray + + fields.indices.map { i => + (starts(i) until starts(i) + bufferCount(fields(i))).map { j => + val buffer = recordBatch.buffers(j) + (buffer.offset(), buffer.length()) + } + } + } - /** Serialized size of each column stream, in column order. */ - def columnStreamSizes(batch: CachedBatch): Seq[Long] = - batch.asInstanceOf[CometCachedBatch].columns.map(_.size).toSeq + /** Stored size of each top-level column: the sum of its buffers' on-body lengths. */ + def columnSizes(batch: CachedBatch, cacheSchema: StructType): Seq[Long] = + columnBufferRanges(batch, cacheSchema).map(_.map(_._2).sum) /** - * Replace one column's stream with bytes that cannot be decoded, in place. + * Whether any of a column's buffers is actually stored compressed. * - * Reading a column this has corrupted fails; reading any other column only succeeds if that - * column's stream was never touched. That is the difference between decoding what was projected - * and decoding everything and projecting afterwards, so it is what the projection tests assert - * on rather than timings. + * Arrow prefixes each compressed buffer with its uncompressed length, and falls back to storing + * a buffer verbatim (length prefix `-1`) when compressing it would not make it smaller. Small + * buffers routinely take that fallback, so [[corruptColumn]] only has something to corrupt when + * this is true; the projection tests assert it as a precondition rather than assuming it. */ - def corruptColumnStream(batch: CachedBatch, index: Int): Unit = { - val columns = batch.asInstanceOf[CometCachedBatch].columns - columns(index) = new ChunkedByteBuffer(Array(ByteBuffer.wrap(Array[Byte](1, 2, 3, 4)))) + def columnIsCompressed(batch: CachedBatch, cacheSchema: StructType, index: Int): Boolean = { + val data = payload(batch) + val start = bodyStart(data) + columnBufferRanges(batch, cacheSchema)(index).exists { case (offset, length) => + length > 8 && uncompressedLength(data, start + offset.toInt) > 0 + } } - /** Whether each column's stream stores that column dictionary encoded, in column order. */ - def columnsAreDictionaryEncoded(batch: CachedBatch): Seq[Boolean] = - batch.asInstanceOf[CometCachedBatch].columns.toSeq.map { buffer => - val in = new DataInputStream(codec.compressedInputStream(buffer.toInputStream())) - val reader = new ArrowStreamReader(Channels.newChannel(in), CometArrowAllocator) - try { - reader.getVectorSchemaRoot.getSchema.getFields.get(0).getDictionary != null - } finally { - reader.close() + /** + * Scramble one column's compressed bytes in place, leaving every other column byte-identical. + * + * Reading a column this has corrupted fails in the decompressor; reading any other column only + * succeeds if this column's buffers were never copied out of the payload. That is the + * difference between decoding what was projected and decoding everything and projecting + * afterwards, so it is what the projection tests assert on rather than timings. + * + * Each buffer's 8-byte uncompressed-length prefix is left intact and only the compressed bytes + * after it are overwritten, so a read of this column fails while decompressing rather than by + * trying to allocate a nonsense length. Requires the column to have a genuinely compressed + * buffer -- see [[columnIsCompressed]]. + */ + def corruptColumn(batch: CachedBatch, cacheSchema: StructType, index: Int): Unit = { + val data = payload(batch) + val start = bodyStart(data) + var corrupted = false + + columnBufferRanges(batch, cacheSchema)(index).foreach { case (offset, length) => + val bufferStart = start + offset.toInt + if (length > 8 && uncompressedLength(data, bufferStart) > 0) { + var i = bufferStart + 8 + while (i < bufferStart + length.toInt) { + // A fixed pattern rather than random bytes, so a failure reproduces exactly. + data(i) = (0xa5 ^ i).toByte + i += 1 + } + corrupted = true } } + require( + corrupted, + s"column $index of the cached batch has no compressed buffer to corrupt; " + + "the test needs data that Arrow actually compresses") + } + /** - * Drop the last `dropBytes` of one column's decoded Arrow stream, in place. + * Truncate the tail of one column's compressed bytes, in place, padding with zeros so every + * other column keeps its offset. * - * [[corruptColumnStream]] replaces the stream outright, so a reader over it fails on the very - * first message, before it has allocated anything. This keeps the stream genuine up to the cut: - * the reader parses the schema and loads the column's dictionary, and only then runs out of - * input part way through the record batch that indexes into it. The cut is made on the decoded - * bytes rather than the compressed ones because a small column compresses to a single block, - * and truncating that fails the decompressor before Arrow reads anything at all. + * [[corruptColumn]] rewrites the whole compressed payload, which fails as soon as the + * decompressor looks at it. This keeps the leading bytes genuine, so a decoder gets a stream + * that starts out valid and then runs out, exercising a failure part way through a column + * rather than at its first byte. */ - def truncateColumnStream(batch: CachedBatch, index: Int, dropBytes: Int): Unit = { - val columns = batch.asInstanceOf[CometCachedBatch].columns - - val decodedStream = new DataInputStream( - codec.compressedInputStream(columns(index).toInputStream())) - val decoded = - try { - val buffer = new java.io.ByteArrayOutputStream() - val chunk = new Array[Byte](8192) - var read = decodedStream.read(chunk) - while (read >= 0) { - buffer.write(chunk, 0, read) - read = decodedStream.read(chunk) + def truncateColumn(batch: CachedBatch, cacheSchema: StructType, index: Int): Unit = { + val data = payload(batch) + val start = bodyStart(data) + var truncated = false + + columnBufferRanges(batch, cacheSchema)(index).foreach { case (offset, length) => + val bufferStart = start + offset.toInt + if (length > 32 && uncompressedLength(data, bufferStart) > 0 && !truncated) { + var i = bufferStart + length.toInt - 16 + while (i < bufferStart + length.toInt) { + data(i) = 0 + i += 1 } - buffer.toByteArray - } finally { - decodedStream.close() + truncated = true } + } + + require( + truncated, + s"column $index of the cached batch has no compressed buffer long enough to truncate") + } + + /** + * Scramble only the last compressed buffer of one column, leaving its earlier buffers genuine. + * + * A string column stores offsets and data as separate compressed buffers, so this makes the + * decoder decompress one buffer of the column successfully and then fail on the next. That is a + * different failure point from [[corruptColumn]], which takes out a column's first buffer and + * so fails before anything of it has been decompressed. + */ + def corruptTrailingBuffer(batch: CachedBatch, cacheSchema: StructType, index: Int): Unit = { + val data = payload(batch) + val start = bodyStart(data) + val compressed = columnBufferRanges(batch, cacheSchema)(index).filter { + case (offset, length) => + length > 8 && uncompressedLength(data, start + offset.toInt) > 0 + } require( - decoded.length > dropBytes, - s"column $index decodes to ${decoded.length} bytes, too few to drop $dropBytes") - - val cbbos = new ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate) - val out = new DataOutputStream(codec.compressedOutputStream(cbbos)) - try { - out.write(decoded, 0, decoded.length - dropBytes) - } finally { - out.close() + compressed.length > 1, + s"column $index has ${compressed.length} compressed buffers; this needs at least two so a " + + "decode can succeed on one and then fail on the next") + + val (offset, length) = compressed.last + val bufferStart = start + offset.toInt + var i = bufferStart + 8 + while (i < bufferStart + length.toInt) { + data(i) = (0xa5 ^ i).toByte + i += 1 } - columns(index) = cbbos.toChunkedByteBuffer } - private def codec: CompressionCodec = CompressionCodec.createCodec(SparkEnv.get.conf) + /** The Arrow fields the read path rebuilds for `cacheSchema`. */ + private def arrowFields(cacheSchema: StructType): Seq[Field] = + Utils + .toArrowSchema(cacheSchema, CometArrowStream.NATIVE_TIMEZONE) + .getFields + .asScala + .toSeq + + /** + * Buffers a field occupies in the record batch body, including every descendant, in the + * depth-first order the body lays them out. + */ + private def bufferCount(field: Field): Int = + TypeLayout.getTypeBufferCount(field.getType) + + field.getChildren.asScala.map(bufferCount).sum + + private def readMetadata(data: Array[Byte]) = { + val channel = new ReadChannel(Channels.newChannel(new ByteArrayInputStream(data))) + val metadata = MessageSerializer.readMessage(channel) + require(metadata != null, "cached payload holds no IPC message") + ( + metadata.getMessage, + metadata.getMessage.header(new FlatBufRecordBatch()).asInstanceOf[FlatBufRecordBatch]) + } + + /** Offset of the record batch body within the payload; the body is its tail. */ + private def bodyStart(data: Array[Byte]): Int = { + val channel = new ReadChannel(Channels.newChannel(new ByteArrayInputStream(data))) + val metadata = MessageSerializer.readMessage(channel) + require(metadata != null, "cached payload holds no IPC message") + data.length - metadata.getMessageBodyLength.toInt + } + + /** + * The uncompressed-length prefix Arrow writes ahead of a compressed buffer, little-endian. A + * value of -1 means the buffer was stored verbatim because compressing it did not pay. + */ + private def uncompressedLength(data: Array[Byte], bufferStart: Int): Long = { + var value = 0L + var i = 7 + while (i >= 0) { + value = (value << 8) | (data(bufferStart + i) & 0xffL) + i -= 1 + } + value + } } From f59c9dc6743486cef696e0d5e52af07038d6311c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 29 Aug 2026 08:52:31 -0600 Subject: [PATCH 2/5] refactor: use Spark's interpreted ordering for bounds, hoist projection layout Cleanup pass over the cache format change. No behaviour change. Drop the `compareStrings` shim in favour of `TypeUtils.getInterpretedOrdering`. That method is public with the same signature on every supported Spark version, and on Spark 4 it resolves a `StringType` through `CollationFactory.fetchCollation(collationId).comparator` -- the comparison the shim was reaching for. So the collation awareness comes from Spark itself and the shim, its Spark 3.x stub and the hand-rolled per-type `compare` all go. The ordering is now resolved once per column per partition rather than being re-dispatched on the `DataType` twice per row. Build the projection's index layout once per partition instead of per batch. The node, buffer and variadic index arithmetic is a pure function of the cached schema and the selected columns, but it walks every field of the relation, so recomputing it per batch made the bookkeeping O(total columns) against O(selected columns) of useful work -- worst in the wide-relation, narrow-projection case the format exists for. `CachedBatchIpc.Projection` now holds that layout and the projected schema, and owns the whole decode; `ProjectedBatch` is left with ownership only. This also puts the projected schema next to the code that packs buffers in the same order, an invariant that previously spanned two files unstated. Smaller cleanups: use Arrow's `DataSizeRoundingUtil.roundUpTo8Multiple` rather than open-coding IPC body alignment; size the serialization buffer from the record batch's known body length instead of growing from 32 bytes; resolve decompressors once instead of per batch; share the dictionary lookup guard between `Utils.combineDictionaryProviders` and the cache writer; read the codec config through one helper carrying the driver-vs-executor rationale; and collapse the duplicated compressed-buffer predicate and scramble loop in the test helper. Corrects two `Utils` scaladocs that still described the per-column stream format this change replaced. Benchmark and codec figures in the docs re-measured against the current code. --- .../user-guide/latest/in-memory-cache.md | 14 +- .../arrow/ArrowCachedBatchSerializer.scala | 162 +++++----- .../execution/arrow/CachedBatchIpc.scala | 281 ++++++++++-------- .../apache/spark/sql/comet/util/Utils.scala | 44 ++- .../apache/comet/shims/CometTypeShim.scala | 12 - .../apache/comet/shims/CometTypeShim.scala | 11 - .../arrow/CometCachedBatchHelper.scala | 182 +++++------- 7 files changed, 353 insertions(+), 353 deletions(-) diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md index 67f0b0e9755..efd949ff92c 100644 --- a/docs/source/user-guide/latest/in-memory-cache.md +++ b/docs/source/user-guide/latest/in-memory-cache.md @@ -70,8 +70,8 @@ six-column relation: | Codec | Materialize | Footprint | Read 1 of 6 | Read 6 of 6 | | ------ | ----------: | --------: | ----------: | ----------: | -| `zstd` | 347 ms | 2 MiB | 52 ms | 63 ms | -| `none` | 1743 ms | 13 MiB | 74 ms | 79 ms | +| `zstd` | 363 ms | 2 MiB | 56 ms | 62 ms | +| `none` | 1776 ms | 13 MiB | 78 ms | 81 ms | Arrow's other IPC codec, LZ4, is deliberately not offered. It is commons-compress's pure-Java implementation and is unrelated to the JNI-accelerated lz4-java behind `spark.io.compression.codec`; @@ -101,11 +101,11 @@ SPARK_GENERATE_BENCHMARK_FILES=1 \ | Query shape | Spark cache scan + convert | `CometInMemoryTableScan` | Relative | | ------------------------------ | -------------------------: | -----------------------: | -------: | -| Repeated scan (3 of 6 columns) | 156 ms | 118 ms | 1.3x | -| Selective filter | 44 ms | 39 ms | 1.1x | -| Row count only (0 of 6) | 32 ms | 28 ms | 1.1x | -| Narrow projection (1 of 6) | 50 ms | 39 ms | 1.3x | -| Full projection (6 of 6) | 316 ms | 135 ms | 2.3x | +| Repeated scan (3 of 6 columns) | 157 ms | 116 ms | 1.4x | +| Selective filter | 44 ms | 38 ms | 1.1x | +| Row count only (0 of 6) | 30 ms | 28 ms | 1.1x | +| Narrow projection (1 of 6) | 49 ms | 39 ms | 1.3x | +| Full projection (6 of 6) | 299 ms | 135 ms | 2.2x | Read what this compares carefully. Comet execution is on in both columns, so the aggregation runs on Comet either way and only the cache-scan boundary moves: on the left, Spark's 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 24eba56f737..343d4f68c85 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 @@ -22,12 +22,11 @@ package org.apache.spark.sql.comet.execution.arrow import scala.collection.JavaConverters._ import scala.util.control.NonFatal -import org.apache.arrow.vector.VectorSchemaRoot -import org.apache.arrow.vector.types.pojo.{Field, Schema} import org.apache.spark.TaskContext import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, GenericInternalRow, IsNotNull, IsNull, UnsafeProjection} +import org.apache.spark.sql.catalyst.util.TypeUtils import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch, SimpleMetricsCachedBatchSerializer} import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.columnar.{DefaultCachedBatch, DefaultCachedBatchSerializer} @@ -38,7 +37,6 @@ import org.apache.spark.storage.StorageLevel import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.{CometArrowAllocator, CometConf} -import org.apache.comet.shims.CometTypeShim import org.apache.comet.vector.NativeUtil /** @@ -73,18 +71,35 @@ private case class CometCachedBatch( * Reads of `CometCachedBatch` keep working when the native scan is disabled, because Spark then * reads the same cached data through the SparkToColumnar fallback path. */ -class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer with CometTypeShim { +class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { import ArrowCachedBatchSerializer.supportsSchema private val fallback = new DefaultCachedBatchSerializer() + /** + * How each cached column's bounds are compared, or null for a column that records none. + * + * Spark's own interpreted ordering for the type, which is the comparison its expressions use, + * so bounds recorded with it order the same way a predicate over the column does. That matters + * for collated strings, where it resolves to the collation's comparator rather than byte order, + * and it is why this needs no per-Spark-version shim: the collation awareness comes from Spark. + * + * Resolved once per partition rather than per row -- `getInterpretedOrdering` walks the type + * and, for a collated string, looks the collation up by id. + */ + private def boundsOrderings(attrs: Seq[Attribute]): Array[Ordering[Any]] = + attrs.map { attr => + if (tracksBounds(attr.dataType)) TypeUtils.getInterpretedOrdering(attr.dataType) else null + }.toArray + // Bounds and null counts per column, gathered before the batch is serialized: serializing // clears the batch's vectors, and the per-column byte sizes that complete the statistics row // are only known afterwards. See statsRow. private def gatherColumnStats( batch: ColumnarBatch, - attrs: Seq[Attribute]): (Array[Any], Array[Any], Array[Int]) = { + attrs: Seq[Attribute], + orderings: Array[Ordering[Any]]): (Array[Any], Array[Any], Array[Int]) = { val numCols = attrs.length val lower = new Array[Any](numCols) val upper = new Array[Any](numCols) @@ -95,16 +110,17 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer with while (c < numCols) { val dt = attrs(c).dataType val col = batch.column(c) + val ordering = orderings(c) var r = 0 while (r < numRows) { if (col.isNullAt(r)) { nulls(c) += 1 - } else if (tracksBounds(dt)) { + } else if (ordering != null) { val value = readValue(col, dt, r) - if (lower(c) == null || compare(dt, value, lower(c)) < 0) { + if (lower(c) == null || ordering.compare(value, lower(c)) < 0) { lower(c) = value } - if (upper(c) == null || compare(dt, value, upper(c)) > 0) { + if (upper(c) == null || ordering.compare(value, upper(c)) > 0) { upper(c) = value } } @@ -149,11 +165,10 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer with // Spark can prune cache batches only for types whose bounds can be compared. // Other types still report null count and row count but leave bounds as null. // - // Every StringType qualifies, collated ones included: bounds are recorded with the collation's - // own comparison, which is the same ordering the predicate Spark generates over that column - // uses. Matching the bare `StringType` object instead would exclude collated columns, since a - // collated StringType is not equal to the default one, and they would then get null bounds and - // no pruning at all. + // Every StringType qualifies, collated ones included. Matching the bare `StringType` object + // instead would exclude them, since a collated StringType is not equal to the default one, and + // they would then get null bounds and no pruning at all. See boundsOrderings for how a collated + // column's bounds are compared. private def tracksBounds(dt: DataType): Boolean = dt match { case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | _: DecimalType | _: StringType | DateType | TimestampType | TimestampNTZType => @@ -176,30 +191,6 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer with case _ => null } - // Compare values using the same physical representation used in the stats row. - private def compare(dt: DataType, left: Any, right: Any): Int = dt match { - case BooleanType => - java.lang.Boolean.compare(left.asInstanceOf[Boolean], right.asInstanceOf[Boolean]) - case ByteType => - java.lang.Byte.compare(left.asInstanceOf[Byte], right.asInstanceOf[Byte]) - case ShortType => - java.lang.Short.compare(left.asInstanceOf[Short], right.asInstanceOf[Short]) - case IntegerType | DateType => - java.lang.Integer.compare(left.asInstanceOf[Int], right.asInstanceOf[Int]) - case LongType | TimestampType | TimestampNTZType => - java.lang.Long.compare(left.asInstanceOf[Long], right.asInstanceOf[Long]) - case FloatType => - java.lang.Float.compare(left.asInstanceOf[Float], right.asInstanceOf[Float]) - case DoubleType => - java.lang.Double.compare(left.asInstanceOf[Double], right.asInstanceOf[Double]) - case _: DecimalType => - left.asInstanceOf[Decimal].compare(right.asInstanceOf[Decimal]) - case st: StringType => - compareStrings(left.asInstanceOf[UTF8String], right.asInstanceOf[UTF8String], st) - case other => - throw new IllegalStateException(s"compare called for unsupported type $other") - } - // Compute Spark-compatible cache stats before serializing each batch to Arrow. // The stats are stored beside the Arrow bytes so Spark's cache filter can prune // CometCachedBatch without decoding the batch first. @@ -207,19 +198,31 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer with // A columnar input batch is not guaranteed to be Arrow-backed; see supportsColumnarInput for // why. Batches that are not get copied into Arrow first, since Utils.serializeBatches only // writes CometVector columns. + /** + * The configured write codec, read on the driver. + * + * Both write paths resolve this here rather than inside their `mapPartitions` closure: the + * closure ships to the executors, where `CometConf` would resolve against whatever `SQLConf` + * happens to be current on that thread rather than against this session's. + */ + private def codecSettings(conf: SQLConf): (String, Int) = + ( + CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC.get(conf), + CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_ZSTD_LEVEL.get(conf)) + private def encodeBatches( batches: Iterator[ColumnarBatch], attrs: Seq[Attribute], - codecName: String, - zstdLevel: Int): Iterator[CachedBatch] = { + codecSetting: (String, Int)): Iterator[CachedBatch] = { val arrowSchema = Utils.toArrowSchema(Utils.fromAttributes(attrs), CometArrowStream.NATIVE_TIMEZONE) - val codec = CachedBatchIpc.compressionCodec(codecName, zstdLevel) + val codec = CachedBatchIpc.compressionCodec(codecSetting._1, codecSetting._2) + val orderings = boundsOrderings(attrs) 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. - val (lower, upper, nulls) = gatherColumnStats(batch, attrs) + val (lower, upper, nulls) = gatherColumnStats(batch, attrs, orderings) val numRows = batch.numRows() val (bytes, columnSizes) = if (Utils.isArrowBacked(batch)) { @@ -313,13 +316,10 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer with storageLevel: StorageLevel, conf: SQLConf): RDD[CachedBatch] = { - // Read on the driver: the closure ships to the executors, where CometConf would resolve - // against whatever SQLConf happens to be current on that thread rather than this session's. - val codecName = CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC.get(conf) - val zstdLevel = CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_ZSTD_LEVEL.get(conf) + val codec = codecSettings(conf) input.mapPartitions { batches => - encodeBatches(batches, schema, codecName, zstdLevel) + encodeBatches(batches, schema, codec) } } @@ -342,12 +342,15 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer with val cacheSchema = Utils.fromAttributes(cacheAttributes) input.mapPartitions { it => - val arrowFields = + // Built once per partition: resolving the Arrow schema and the projection's buffer layout + // walks every field of the cached relation, which would otherwise be paid per batch. + val projection = new CachedBatchIpc.Projection( Utils .toArrowSchema(cacheSchema, CometArrowStream.NATIVE_TIMEZONE) .getFields .asScala - .toSeq + .toIndexedSeq, + indices) // A ProjectedBatch owns the vectors of the batch it produced, and releases them only when // that batch has been consumed. A consumer that stops early -- LIMIT, take(), or a @@ -374,7 +377,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer with // Nothing to decode: the row count is the whole answer, and it is already here. Iterator.single(new ColumnarBatch(Array.empty[ColumnVector], cb.numRows)) } else { - val projected = new ProjectedBatch(cb, arrowFields, indices) + val projected = new ProjectedBatch(cb, projection) current = projected projected.batches } @@ -387,48 +390,28 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer with } /** - * Loads the projected columns of one cached batch into Arrow vectors. - * - * The schema is rebuilt from the cached relation's attributes rather than read from the - * payload, which stores none. `CachedBatchIpc.readProjected` then materializes only the - * selected columns' buffers, so the rest are never copied out of the cached bytes or - * decompressed. + * Owns the Arrow vectors decoded for one cached batch. * - * The decoded vectors stay owned by this object: closing it releases the batch, which is why - * this yields a single-element iterator that closes on exhaustion. + * The decode itself belongs to `CachedBatchIpc.Projection`, which is where knowledge of the + * payload format lives; what is left here is ownership. The vectors stay owned by this object + * -- closing it releases them -- which is why this yields a single-element iterator that closes + * on exhaustion. */ - private class ProjectedBatch( - cached: CometCachedBatch, - arrowFields: Seq[Field], - indices: Array[Int]) { - - // Allocated before anything can throw, so that a failure below has a root to release. - private val root = VectorSchemaRoot.create( - new Schema(indices.map(arrowFields).toSeq.asJava), - CometArrowAllocator) + private class ProjectedBatch(cached: CometCachedBatch, projection: CachedBatchIpc.Projection) { + + // 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 var closed = false - // Loading happens during construction, so `batches` below can hand out the root directly. - try { - val recordBatch = - CachedBatchIpc.readProjected(cached.bytes, arrowFields, indices, CometArrowAllocator) - try { - CachedBatchIpc.loaderFor(root).load(recordBatch) - } finally { - recordBatch.close() - } - // A cached batch's columns all cover the same rows. Check rather than trust: a mismatch - // would otherwise build a batch whose columns disagree with the row count recorded beside - // them, which reads as corrupt data far from here. - if (root.getRowCount != cached.numRows) { - throw new IllegalStateException( - s"Cached batch decoded ${root.getRowCount} rows, expected ${cached.numRows}") - } - } catch { - case NonFatal(e) => - try root.close() - catch { case NonFatal(closeError) => e.addSuppressed(closeError) } - throw e + // A cached batch's columns all cover the same rows. Check rather than trust: a mismatch would + // otherwise build a batch whose columns disagree with the row count recorded beside them, + // which reads as corrupt data far from here. + if (root.getRowCount != cached.numRows) { + val decoded = root.getRowCount + close() + throw new IllegalStateException( + s"Cached batch decoded $decoded rows, expected ${cached.numRows}") } def close(): Unit = synchronized { @@ -471,8 +454,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer with fallback.convertInternalRowToCachedBatch(input, schema, storageLevel, conf) } else { val batchSize = conf.columnBatchSize - val codecName = CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC.get(conf) - val zstdLevel = CometConf.COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_ZSTD_LEVEL.get(conf) + val codec = codecSettings(conf) input.mapPartitions { rows => val iter = CometArrowConverters.rowToArrowBatchIter( @@ -490,7 +472,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer with CometArrowStream.NATIVE_TIMEZONE, CometArrowAllocator) - encodeBatches(iter, schema, codecName, zstdLevel) + encodeBatches(iter, schema, codec) } } } 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 55215ef548a..43a78ab858b 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,12 +29,13 @@ 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, VectorSchemaRoot, VectorUnloader} +import org.apache.arrow.vector.{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} import org.apache.arrow.vector.ipc.message.{ArrowBodyCompression, ArrowFieldNode, ArrowRecordBatch, MessageSerializer} -import org.apache.arrow.vector.types.pojo.{ArrowType, Field} +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} +import org.apache.arrow.vector.util.DataSizeRoundingUtil import org.apache.spark.SparkException import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.vectorized.ColumnarBatch @@ -80,6 +81,24 @@ private[comet] object CachedBatchIpc { "Supported values: none, zstd") } + // Room for the encapsulated metadata message that precedes the body. The message is a small + // flatbuffer whose size grows with the field count, not the data, so this is a starting size for + // the output buffer rather than a bound -- it grows if a very wide schema needs more. + private val METADATA_SIZE_HINT = 8 * 1024 + + // 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. + private val readCodecs: Map[CompressionUtil.CodecType, CompressionCodec] = + CompressionUtil.CodecType + .values() + .filter(_ != CompressionUtil.CodecType.NO_COMPRESSION) + .map(t => t -> CommonsCompressionFactory.INSTANCE.createCodec(t)) + .toMap + + /** The decompressor for a body-compression byte, or None when the batch is stored plain. */ + private def readCodec(compressionType: Byte): Option[CompressionCodec] = + readCodecs.get(CompressionUtil.CodecType.fromCompressionType(compressionType)) + /** * Serialize `batch` into one encapsulated IPC RecordBatch message. * @@ -125,7 +144,12 @@ private[comet] object CachedBatchIpc { // does the same, so both writers leave a batch they were handed in the same state. root.clear() - val out = new ByteArrayOutputStream() + // Sized up front from the body length the record batch already knows, plus room for the + // metadata message. An unsized ByteArrayOutputStream starts at 32 bytes and doubles, so a + // multi-MiB payload would be reallocated and recopied a dozen-odd times per batch. + val sizeHint = recordBatch.computeBodyLength() + METADATA_SIZE_HINT + val out = new ByteArrayOutputStream( + math.min(math.max(sizeHint, METADATA_SIZE_HINT), Int.MaxValue.toLong).toInt) val channel = new WriteChannel(Channels.newChannel(out)) MessageSerializer.serialize(channel, recordBatch) (out.toByteArray, columnSizes(fields, recordBatch)) @@ -141,111 +165,151 @@ private[comet] object CachedBatchIpc { } /** - * Read an encapsulated IPC RecordBatch message, materializing off-heap only the buffers of the - * requested top-level columns. + * Everything about reading one projection of this format that does not change between batches. * - * The body is a flat, depth-first sequence of buffers in schema order, so each top-level column - * owns a contiguous run of buffers whose length is [[fieldBufferCount]]; field nodes and - * variadic buffer counts run in the same order. The selected columns' bytes are copied into a - * single off-heap allocation, each buffer 8-byte aligned exactly as Arrow's IPC body lays them - * out, and the returned batch's buffers are windows into it -- one allocation, no per-buffer - * bookkeeping. + * The index arithmetic here is a pure function of the cached schema and the selected columns, + * both fixed for the life of a scan, but it walks every field of the whole relation rather than + * just the projected ones. Recomputing it per batch would make the bookkeeping O(total columns) + * while the useful work is O(selected columns) -- worst in exactly the wide-relation, + * narrow-projection case this format exists for. A scan builds one of these per partition. * - * Only the selected buffers are ever decompressed. A buffer's recorded (offset, length) covers - * its on-body bytes including the uncompressed-length prefix, so a copied window is exactly - * what the writer emitted; the columns that were not selected are never read, let alone - * inflated. The copied windows are then decompressed in one pass -- see [[decompressed]] for - * why that is not left to `VectorLoader` -- so what comes back is an uncompressed batch. - * - * The returned batch owns its buffers; the caller closes it. + * Holding the projected `Schema` here too is what keeps it consistent with the buffers: + * [[load]] packs field nodes and buffers by walking `selectedIndices` in order, and the schema + * is built from the same walk, so the two cannot drift apart. */ - def readProjected( - data: Array[Byte], - schemaFields: Seq[Field], - selectedIndices: Array[Int], - allocator: BufferAllocator): ArrowRecordBatch = { - val readChannel = new ReadChannel(Channels.newChannel(new ByteArrayInputStream(data))) - // Reads the message metadata only. The body stays in `data` and is copied selectively below. - val metadata = MessageSerializer.readMessage(readChannel) - if (metadata == null) { - throw new SparkException("Unexpected end of input reading a Comet cached batch") - } - val batch = - metadata.getMessage.header(new FlatBufRecordBatch()).asInstanceOf[FlatBufRecordBatch] - // serialize writes exactly [encapsulated message][body] and nothing after it, so the body is - // the tail of `data`. - val bodyStart = data.length - metadata.getMessageBodyLength.toInt + final class Projection(arrowFields: Seq[Field], selectedIndices: Array[Int]) { - val compression = - if (batch.compression() == null) NoCompressionCodec.DEFAULT_BODY_COMPRESSION - else new ArrowBodyCompression(batch.compression().codec(), batch.compression().method()) + private val schema = new Schema(selectedIndices.map(arrowFields).toSeq.asJava) - val nodeStarts = schemaFields.scanLeft(0)(_ + fieldNodeCount(_)).toArray - val bufferStarts = schemaFields.scanLeft(0)(_ + fieldBufferCount(_)).toArray - val variadicStarts = schemaFields.scanLeft(0)(_ + fieldVariadicCount(_)).toArray - val hasVariadic = batch.variadicBufferCountsLength() > 0 + // A record batch body is a flat, depth-first sequence of buffers in schema order, so each + // top-level column owns a contiguous run of it; field nodes and variadic buffer counts run in + // the same order. + private val nodeIndices = selectedRange(arrowFields, selectedIndices, fieldNodeCount) + private val bufferIndices = selectedRange(arrowFields, selectedIndices, fieldBufferCount) + private val variadicIndices = selectedRange(arrowFields, selectedIndices, fieldVariadicCount) + + /** + * Decode the projected columns of one cached payload into a fresh root the caller owns. + * + * Only the selected buffers are ever materialized off-heap or decompressed. The message + * metadata records every buffer's offset and length within the body, so the selected columns' + * bytes are copied into a single allocation -- each 8-byte aligned exactly as Arrow's IPC + * body lays them out -- and the columns that were not selected are never read, let alone + * inflated. + * + * A buffer's recorded (offset, length) covers its on-body bytes including the + * uncompressed-length prefix, so a copied window is exactly what the writer emitted. The + * windows are then decompressed in one pass; see [[decompressed]] for why that is not left to + * `VectorLoader`. + */ + def load(data: Array[Byte], allocator: BufferAllocator): VectorSchemaRoot = { + val readChannel = new ReadChannel(Channels.newChannel(new ByteArrayInputStream(data))) + // Reads the message metadata only. The body stays in `data` and is copied selectively. + val metadata = MessageSerializer.readMessage(readChannel) + if (metadata == null) { + throw new SparkException("Unexpected end of input reading a Comet cached batch") + } + val batch = + metadata.getMessage.header(new FlatBufRecordBatch()).asInstanceOf[FlatBufRecordBatch] + // serialize writes exactly [encapsulated message][body] and nothing after it, so the body is + // the tail of `data`. + val bodyStart = data.length - metadata.getMessageBodyLength.toInt - // The selected columns' field nodes, buffer indices and variadic counts, in output order. - val nodes = new java.util.ArrayList[ArrowFieldNode]() - val bufferIndices = mutable.ArrayBuffer.empty[Int] - val variadicCounts = new java.util.ArrayList[java.lang.Long]() - selectedIndices.foreach { i => - val field = schemaFields(i) - val nodeStart = nodeStarts(i) - (nodeStart until nodeStart + fieldNodeCount(field)).foreach { j => + val compression = + if (batch.compression() == null) NoCompressionCodec.DEFAULT_BODY_COMPRESSION + else new ArrowBodyCompression(batch.compression().codec(), batch.compression().method()) + + val nodes = new java.util.ArrayList[ArrowFieldNode](nodeIndices.length) + nodeIndices.foreach { j => val node = batch.nodes(j) nodes.add(new ArrowFieldNode(node.length(), node.nullCount())) } - val bufferStart = bufferStarts(i) - (bufferStart until bufferStart + fieldBufferCount(field)).foreach(bufferIndices += _) - if (hasVariadic) { - val variadicStart = variadicStarts(i) - (variadicStart until variadicStart + fieldVariadicCount(field)) - .foreach(j => variadicCounts.add(batch.variadicBufferCounts(j))) + val variadicCounts = new java.util.ArrayList[java.lang.Long](variadicIndices.length) + if (batch.variadicBufferCountsLength() > 0) { + variadicIndices.foreach(j => variadicCounts.add(batch.variadicBufferCounts(j))) } - } - val layout = bufferIndices.map { j => - val buffer = batch.buffers(j) - (buffer.offset(), buffer.length()) - } - val alignedSizes = layout.map { case (_, length) => ((length + 7) / 8) * 8 } - // allocator.buffer(0) is legal but yields a buffer no window can be sliced from, and an - // all-empty projection (every selected column a NullVector, say) would ask for exactly that. - val body = allocator.buffer(math.max(alignedSizes.sum, 1L)) - val compressedBatch = - try { - val buffers = new java.util.ArrayList[ArrowBuf]() - var position = 0L - layout.indices.foreach { k => - val (sourceOffset, length) = layout(k) - if (length > 0) { - body.setBytes(position, data, bodyStart + sourceOffset.toInt, length.toInt) + val offsets = new Array[Long](bufferIndices.length) + val lengths = new Array[Long](bufferIndices.length) + var total = 0L + var k = 0 + while (k < bufferIndices.length) { + val buffer = batch.buffers(bufferIndices(k)) + offsets(k) = buffer.offset() + lengths(k) = buffer.length() + total += DataSizeRoundingUtil.roundUpTo8Multiple(lengths(k)) + k += 1 + } + + // allocator.buffer(0) is legal but yields a buffer no window can be sliced from, and an + // all-empty projection (every selected column a NullVector, say) would ask for exactly that. + val body = allocator.buffer(math.max(total, 1L)) + val compressedBatch = + try { + val buffers = new java.util.ArrayList[ArrowBuf](bufferIndices.length) + var position = 0L + var i = 0 + while (i < bufferIndices.length) { + val length = lengths(i) + if (length > 0) { + body.setBytes(position, data, bodyStart + offsets(i).toInt, length.toInt) + } + val window = body.slice(position, length) + window.writerIndex(length) + buffers.add(window) + position += DataSizeRoundingUtil.roundUpTo8Multiple(length) + i += 1 } - val window = body.slice(position, length) - window.writerIndex(length) - buffers.add(window) - position += alignedSizes(k) + new ArrowRecordBatch( + batch.length().toInt, + nodes, + buffers, + compression, + variadicCounts, + false) + } catch { + case NonFatal(e) => + body.close() + throw e } - new ArrowRecordBatch( - batch.length().toInt, - nodes, - buffers, - compression, - variadicCounts, - false) + + // The constructor retained each window; slice() alone does not. Dropping `body`'s own + // reference leaves the batch as sole owner of the one allocation, so closing the batch is + // what frees it -- and closing `body` again would drive its reference count negative. + body.close() + val plainBatch = + try decompressed(compressedBatch, allocator) + finally compressedBatch.close() + + // The loader needs no compression factory: every buffer is decompressed by this point. + val root = VectorSchemaRoot.create(schema, allocator) + try { + new VectorLoader(root).load(plainBatch) + root } catch { case NonFatal(e) => - body.close() + try root.close() + catch { case NonFatal(closeError) => e.addSuppressed(closeError) } throw e + } finally { + plainBatch.close() } + } + } - // The constructor retained each window; slice() alone does not. Dropping `body`'s own - // reference leaves the batch as sole owner of the one allocation, so closing the batch below - // is what frees it -- and closing `body` again here would drive its reference count negative. - body.close() - try decompressed(compressedBatch, allocator) - finally compressedBatch.close() + /** + * The indices, within a record batch's flat depth-first sequence, that the selected columns + * own. + * + * `count` gives how many entries of the sequence a field occupies including its descendants, so + * a running total over every field turns a column index into its run within the sequence. + */ + private def selectedRange( + arrowFields: Seq[Field], + selectedIndices: Array[Int], + count: Field => Int): Array[Int] = { + val starts = arrowFields.scanLeft(0)(_ + count(_)).toArray + selectedIndices.flatMap(i => starts(i) until starts(i + 1)) } /** @@ -269,14 +333,10 @@ private[comet] object CachedBatchIpc { batch: ArrowRecordBatch, allocator: BufferAllocator): ArrowRecordBatch = { // getCodec is the raw IPC byte; the factory keys off the enum. Both sides of the comparison - // below have to be CodecType: NoCompressionCodec.COMPRESSION_TYPE is the byte -1, and Scala - // compares a CodecType against it by universal equality, which is quietly always unequal. - val codecType = - CompressionUtil.CodecType.fromCompressionType(batch.getBodyCompression.getCodec) - val compressed = codecType != CompressionUtil.CodecType.NO_COMPRESSION - val codec: CompressionCodec = - if (compressed) CommonsCompressionFactory.INSTANCE.createCodec(codecType) - else NoCompressionCodec.INSTANCE + // in readCodec have to be CodecType: NoCompressionCodec.COMPRESSION_TYPE is the byte -1, and + // Scala compares a CodecType against it by universal equality, which is quietly always + // unequal. + val codec = readCodec(batch.getBodyCompression.getCodec) val buffers = new java.util.ArrayList[ArrowBuf]() try { @@ -285,8 +345,10 @@ private[comet] object CachedBatchIpc { val plain = try { // An empty buffer carries no compressed length prefix to read. - if (compressed && buffer.writerIndex() > 0) codec.decompress(allocator, buffer) - else buffer + codec match { + case Some(c) if buffer.writerIndex() > 0 => c.decompress(allocator, buffer) + case _ => buffer + } } catch { case NonFatal(e) => buffer.getReferenceManager.release() @@ -315,15 +377,6 @@ private[comet] object CachedBatchIpc { } } - /** - * A `VectorLoader` for what [[readProjected]] returns. - * - * No compression factory: [[readProjected]] has already decompressed every buffer, so the - * loader only ever sees a batch marked uncompressed. - */ - def loaderFor(root: VectorSchemaRoot): org.apache.arrow.vector.VectorLoader = - new org.apache.arrow.vector.VectorLoader(root) - /** * The on-body compressed size of each top-level column. * @@ -355,16 +408,10 @@ private[comet] object CachedBatchIpc { try { val vectors = Utils.getBatchFieldVectorsWithProviders(batch).map { case (vector, providerOpt) => - val encoding = vector.getField.getDictionary - if (encoding == null) { + if (vector.getField.getDictionary == null) { vector } else { - val dictionary = providerOpt.map(_.lookup(encoding.getId)).orNull - if (dictionary == null) { - throw new SparkException( - s"Column ${vector.getField.getName} is dictionary encoded with ID " + - s"${encoding.getId}, but no dictionary with that ID was provided") - } + val dictionary = Utils.lookupDictionary(vector, providerOpt) val decoded = DictionaryEncoder.decode(vector, dictionary, allocator) hydrated += decoded decoded.asInstanceOf[FieldVector] diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index f418c31a626..97e78f3eaff 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -446,12 +446,12 @@ object Utils extends CometTypeShim with Logging { /** * The dictionaries every dictionary-encoded column of `columns` refers to, as one provider. * - * Columns of a batch need not share a provider. Comet's cache decodes each column from its own - * Arrow stream, so a dictionary-backed column arrives carrying the provider its reader built, - * and a batch that reaches [[serializeBatches]] -- a native broadcast of a cache scan, say -- - * can hold several. Writing the whole batch emits one schema covering every column and resolves - * each column's dictionary ID against the single provider the writer was given, so handing it - * any one column's provider fails with "Could not find dictionary with ID n" for the others. + * Columns of a batch need not share a provider. A batch assembled from several upstream readers + * -- a shuffle reader's output, or a broadcast that coalesces many blocks -- carries a + * dictionary-backed column with whichever provider its own reader built, so one batch can hold + * several. Writing the whole batch emits one schema covering every column and resolves each + * column's dictionary ID against the single provider the writer was given, so handing it any + * one column's provider fails with "Could not find dictionary with ID n" for the others. */ private def combineDictionaryProviders( columns: Seq[(FieldVector, Option[DictionaryProvider])]): Option[DictionaryProvider] = { @@ -461,12 +461,7 @@ object Utils extends CometTypeShim with Logging { val encoding = vector.getField.getDictionary if (encoding != null) { val id = encoding.getId - val dictionary = providerOpt.map(_.lookup(id)).orNull - if (dictionary == null) { - throw new SparkException( - s"Column ${vector.getField.getName} is dictionary encoded with ID $id, but no " + - "dictionary with that ID was provided") - } + val dictionary = lookupDictionary(vector, providerOpt) dictionaries.get(id) match { // Every provider seen here descends from one upstream reader, which numbers the // dictionaries it hands out, so two columns sharing an ID share the dictionary itself. @@ -484,13 +479,32 @@ object Utils extends CometTypeShim with Logging { else Some(new MapDictionaryProvider(dictionaries.values.toSeq: _*)) } + /** + * The dictionary a dictionary-encoded column refers to, or a failure naming the column. + * + * Shared with the cache serializer, which decodes dictionary-encoded columns rather than + * folding their providers together, so that both report a missing dictionary the same way. + */ + def lookupDictionary( + vector: FieldVector, + providerOpt: Option[DictionaryProvider]): Dictionary = { + val id = vector.getField.getDictionary.getId + val dictionary = providerOpt.map(_.lookup(id)).orNull + if (dictionary == null) { + throw new SparkException( + s"Column ${vector.getField.getName} is dictionary encoded with ID $id, but no " + + "dictionary with that ID was provided") + } + dictionary + } + /** * Field vectors of `batch` paired with the dictionary provider each column was decoded with. * * [[getBatchFieldVectors]] folds these into one provider covering the whole batch, which is - * what a single stream over every column needs. Comet's cache decodes each column from its own - * stream and writes it back the same way, so it keeps the pairing instead: each column is - * written with the provider it was decoded with. + * what a single stream over every column needs. Comet's cache serializer keeps the pairing + * instead: its payload has no schema message to describe a dictionary encoding, so it decodes + * each such column against the provider that column arrived with. */ def getBatchFieldVectorsWithProviders( batch: ColumnarBatch): Seq[(FieldVector, Option[DictionaryProvider])] = { diff --git a/spark/src/main/spark-3.x/org/apache/comet/shims/CometTypeShim.scala b/spark/src/main/spark-3.x/org/apache/comet/shims/CometTypeShim.scala index a6b34b74e7d..b71476c3dd1 100644 --- a/spark/src/main/spark-3.x/org/apache/comet/shims/CometTypeShim.scala +++ b/spark/src/main/spark-3.x/org/apache/comet/shims/CometTypeShim.scala @@ -22,7 +22,6 @@ package org.apache.comet.shims import scala.annotation.nowarn import org.apache.spark.sql.types.{DataType, StructType} -import org.apache.spark.unsafe.types.{ByteArray, UTF8String} trait CometTypeShim { @nowarn // Spark 4 feature; stubbed to false in Spark 3.x for compatibility. @@ -42,15 +41,4 @@ trait CometTypeShim { @nowarn // Spark 4.1 feature; TimeType doesn't exist in Spark 3.x. def isTimeType(dt: DataType): Boolean = false - - /** - * Compare two strings under the collation of `dt`, which must be a `StringType`. - * - * Spark 3.x has no collations, so every string comparison is byte order. Callers that record - * comparable bounds (Comet's cache statistics, for instance) use this so the ordering they - * store is the one Spark's own comparison would produce. - */ - @nowarn // Collation is a Spark 4 feature; on 3.x every StringType compares as bytes. - def compareStrings(left: UTF8String, right: UTF8String, dt: DataType): Int = - ByteArray.compareBinary(left.getBytes, right.getBytes) } diff --git a/spark/src/main/spark-4.x/org/apache/comet/shims/CometTypeShim.scala b/spark/src/main/spark-4.x/org/apache/comet/shims/CometTypeShim.scala index 71e72dd2de4..f48955a7da5 100644 --- a/spark/src/main/spark-4.x/org/apache/comet/shims/CometTypeShim.scala +++ b/spark/src/main/spark-4.x/org/apache/comet/shims/CometTypeShim.scala @@ -21,7 +21,6 @@ package org.apache.comet.shims import org.apache.spark.sql.execution.datasources.VariantMetadata import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StringType, StructType, VariantType} -import org.apache.spark.unsafe.types.UTF8String trait CometTypeShim { // A `StringType` carries collation metadata in Spark 4.0. Only non-default (non-UTF8_BINARY) @@ -65,14 +64,4 @@ trait CometTypeShim { dt.getClass.getSimpleName.startsWith("TimeType") def hasCollationSupport: Boolean = true - - /** - * Compare two strings under the collation of `dt`, which must be a `StringType`. - * - * `semanticCompare` is the comparison Spark's own expressions use for the type, so bounds - * recorded with it order the same way a predicate over the column does. For the default - * UTF8_BINARY collation it is byte order, which is what Spark 3.x always does. - */ - def compareStrings(left: UTF8String, right: UTF8String, dt: DataType): Int = - left.semanticCompare(right, dt.asInstanceOf[StringType].collationId) } 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 56c3b949569..e1a089f20f8 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 @@ -27,7 +27,7 @@ import scala.jdk.CollectionConverters._ import org.apache.arrow.flatbuf.{MessageHeader, RecordBatch => FlatBufRecordBatch} import org.apache.arrow.vector.TypeLayout import org.apache.arrow.vector.ipc.ReadChannel -import org.apache.arrow.vector.ipc.message.MessageSerializer +import org.apache.arrow.vector.ipc.message.{MessageMetadataResult, MessageSerializer} import org.apache.arrow.vector.types.pojo.Field import org.apache.spark.sql.columnar.CachedBatch import org.apache.spark.sql.comet.util.Utils @@ -50,9 +50,6 @@ object CometCachedBatchHelper { private def payload(batch: CachedBatch): Array[Byte] = batch.asInstanceOf[CometCachedBatch].bytes - /** Stored size of the whole cached batch, in bytes. */ - def payloadSize(batch: CachedBatch): Long = payload(batch).length.toLong - /** * Whether the payload begins with a Schema message rather than going straight to the record * batch. @@ -62,25 +59,7 @@ object CometCachedBatchHelper { * here rather than only as a footprint number. */ def hasSchemaMessage(batch: CachedBatch): Boolean = - readMetadata(payload(batch))._1.headerType() == MessageHeader.Schema - - /** - * The on-body (offset, length) of every Arrow buffer belonging to each top-level column, in - * column order. - */ - def columnBufferRanges(batch: CachedBatch, cacheSchema: StructType): Seq[Seq[(Long, Long)]] = { - val data = payload(batch) - val (_, recordBatch) = readMetadata(data) - val fields = arrowFields(cacheSchema) - val starts = fields.scanLeft(0)(_ + bufferCount(_)).toArray - - fields.indices.map { i => - (starts(i) until starts(i) + bufferCount(fields(i))).map { j => - val buffer = recordBatch.buffers(j) - (buffer.offset(), buffer.length()) - } - } - } + readMessage(payload(batch)).getMessage.headerType() == MessageHeader.Schema /** Stored size of each top-level column: the sum of its buffers' on-body lengths. */ def columnSizes(batch: CachedBatch, cacheSchema: StructType): Seq[Long] = @@ -94,13 +73,8 @@ object CometCachedBatchHelper { * buffers routinely take that fallback, so [[corruptColumn]] only has something to corrupt when * this is true; the projection tests assert it as a precondition rather than assuming it. */ - def columnIsCompressed(batch: CachedBatch, cacheSchema: StructType, index: Int): Boolean = { - val data = payload(batch) - val start = bodyStart(data) - columnBufferRanges(batch, cacheSchema)(index).exists { case (offset, length) => - length > 8 && uncompressedLength(data, start + offset.toInt) > 0 - } - } + def columnIsCompressed(batch: CachedBatch, cacheSchema: StructType, index: Int): Boolean = + compressedRanges(batch, cacheSchema, index).nonEmpty /** * Scramble one column's compressed bytes in place, leaving every other column byte-identical. @@ -110,38 +84,38 @@ object CometCachedBatchHelper { * difference between decoding what was projected and decoding everything and projecting * afterwards, so it is what the projection tests assert on rather than timings. * - * Each buffer's 8-byte uncompressed-length prefix is left intact and only the compressed bytes - * after it are overwritten, so a read of this column fails while decompressing rather than by - * trying to allocate a nonsense length. Requires the column to have a genuinely compressed - * buffer -- see [[columnIsCompressed]]. + * Requires the column to have a genuinely compressed buffer -- see [[columnIsCompressed]]. */ def corruptColumn(batch: CachedBatch, cacheSchema: StructType, index: Int): Unit = { - val data = payload(batch) - val start = bodyStart(data) - var corrupted = false - - columnBufferRanges(batch, cacheSchema)(index).foreach { case (offset, length) => - val bufferStart = start + offset.toInt - if (length > 8 && uncompressedLength(data, bufferStart) > 0) { - var i = bufferStart + 8 - while (i < bufferStart + length.toInt) { - // A fixed pattern rather than random bytes, so a failure reproduces exactly. - data(i) = (0xa5 ^ i).toByte - i += 1 - } - corrupted = true - } - } - + val ranges = compressedRanges(batch, cacheSchema, index) require( - corrupted, + ranges.nonEmpty, s"column $index of the cached batch has no compressed buffer to corrupt; " + "the test needs data that Arrow actually compresses") + ranges.foreach { case (start, length) => scramble(payload(batch), start, length) } } /** - * Truncate the tail of one column's compressed bytes, in place, padding with zeros so every - * other column keeps its offset. + * Scramble only the last compressed buffer of one column, leaving its earlier buffers genuine. + * + * A string column stores offsets and data as separate compressed buffers, so this makes the + * decoder decompress one buffer of the column successfully and then fail on the next. That is a + * different failure point from [[corruptColumn]], which takes out a column's first buffer and + * so fails before anything of it has been decompressed. + */ + def corruptTrailingBuffer(batch: CachedBatch, cacheSchema: StructType, index: Int): Unit = { + val ranges = compressedRanges(batch, cacheSchema, index) + require( + ranges.length > 1, + s"column $index has ${ranges.length} compressed buffers; this needs at least two so a " + + "decode can succeed on one and then fail on the next") + val (start, length) = ranges.last + scramble(payload(batch), start, length) + } + + /** + * Zero the tail of one column's compressed bytes, in place, leaving every other column's bytes + * and offsets untouched. * * [[corruptColumn]] rewrites the whole compressed payload, which fails as soon as the * decompressor looks at it. This keeps the leading bytes genuine, so a decoder gets a stream @@ -150,55 +124,70 @@ object CometCachedBatchHelper { */ def truncateColumn(batch: CachedBatch, cacheSchema: StructType, index: Int): Unit = { val data = payload(batch) - val start = bodyStart(data) - var truncated = false - - columnBufferRanges(batch, cacheSchema)(index).foreach { case (offset, length) => - val bufferStart = start + offset.toInt - if (length > 32 && uncompressedLength(data, bufferStart) > 0 && !truncated) { - var i = bufferStart + length.toInt - 16 - while (i < bufferStart + length.toInt) { - data(i) = 0 - i += 1 - } - truncated = true - } + val target = compressedRanges(batch, cacheSchema, index).find { case (_, length) => + length > 32 } - require( - truncated, + target.isDefined, s"column $index of the cached batch has no compressed buffer long enough to truncate") + val (start, length) = target.get + java.util.Arrays.fill(data, (start + length - 16).toInt, (start + length).toInt, 0.toByte) } /** - * Scramble only the last compressed buffer of one column, leaving its earlier buffers genuine. + * The absolute (start, length) of each of a column's buffers that Arrow actually compressed. * - * A string column stores offsets and data as separate compressed buffers, so this makes the - * decoder decompress one buffer of the column successfully and then fail on the next. That is a - * different failure point from [[corruptColumn]], which takes out a column's first buffer and - * so fails before anything of it has been decompressed. + * `start` is an index into the payload, not an offset within the body, so callers can write + * through it directly. A buffer shorter than its 8-byte uncompressed-length prefix, or one + * whose prefix reads `-1`, was stored verbatim and is excluded: overwriting it would change the + * values a read returns rather than making the read fail. */ - def corruptTrailingBuffer(batch: CachedBatch, cacheSchema: StructType, index: Int): Unit = { + private def compressedRanges( + batch: CachedBatch, + cacheSchema: StructType, + index: Int): Seq[(Long, Long)] = { val data = payload(batch) - val start = bodyStart(data) - val compressed = columnBufferRanges(batch, cacheSchema)(index).filter { - case (offset, length) => - length > 8 && uncompressedLength(data, start + offset.toInt) > 0 + val bodyStart = data.length - readMessage(data).getMessageBodyLength + columnBufferRanges(batch, cacheSchema)(index).collect { + case (offset, length) if length > 8 && uncompressedLength(data, bodyStart + offset) > 0 => + (bodyStart + offset, length) } - require( - compressed.length > 1, - s"column $index has ${compressed.length} compressed buffers; this needs at least two so a " + - "decode can succeed on one and then fail on the next") + } - val (offset, length) = compressed.last - val bufferStart = start + offset.toInt - var i = bufferStart + 8 - while (i < bufferStart + length.toInt) { + /** Overwrite a compressed buffer's payload, leaving its uncompressed-length prefix intact. */ + private def scramble(data: Array[Byte], start: Long, length: Long): Unit = { + var i = (start + 8).toInt + val end = (start + length).toInt + while (i < end) { + // A fixed pattern rather than random bytes, so a failure reproduces exactly. data(i) = (0xa5 ^ i).toByte i += 1 } } + /** + * The on-body (offset, length) of every Arrow buffer belonging to each top-level column, in + * column order. + */ + private def columnBufferRanges( + batch: CachedBatch, + cacheSchema: StructType): Seq[Seq[(Long, Long)]] = { + val data = payload(batch) + val recordBatch = + readMessage(data).getMessage + .header(new FlatBufRecordBatch()) + .asInstanceOf[FlatBufRecordBatch] + val fields = arrowFields(cacheSchema) + val starts = fields.scanLeft(0)(_ + bufferCount(_)).toArray + + fields.indices.map { i => + (starts(i) until starts(i + 1)).map { j => + val buffer = recordBatch.buffers(j) + (buffer.offset(), buffer.length()) + } + } + } + /** The Arrow fields the read path rebuilds for `cacheSchema`. */ private def arrowFields(cacheSchema: StructType): Seq[Field] = Utils @@ -215,32 +204,23 @@ object CometCachedBatchHelper { TypeLayout.getTypeBufferCount(field.getType) + field.getChildren.asScala.map(bufferCount).sum - private def readMetadata(data: Array[Byte]) = { - val channel = new ReadChannel(Channels.newChannel(new ByteArrayInputStream(data))) - val metadata = MessageSerializer.readMessage(channel) - require(metadata != null, "cached payload holds no IPC message") - ( - metadata.getMessage, - metadata.getMessage.header(new FlatBufRecordBatch()).asInstanceOf[FlatBufRecordBatch]) - } - - /** Offset of the record batch body within the payload; the body is its tail. */ - private def bodyStart(data: Array[Byte]): Int = { + /** The payload's leading IPC message, carrying both its header and its body length. */ + private def readMessage(data: Array[Byte]): MessageMetadataResult = { val channel = new ReadChannel(Channels.newChannel(new ByteArrayInputStream(data))) val metadata = MessageSerializer.readMessage(channel) require(metadata != null, "cached payload holds no IPC message") - data.length - metadata.getMessageBodyLength.toInt + metadata } /** * The uncompressed-length prefix Arrow writes ahead of a compressed buffer, little-endian. A * value of -1 means the buffer was stored verbatim because compressing it did not pay. */ - private def uncompressedLength(data: Array[Byte], bufferStart: Int): Long = { + private def uncompressedLength(data: Array[Byte], bufferStart: Long): Long = { var value = 0L var i = 7 while (i >= 0) { - value = (value << 8) | (data(bufferStart + i) & 0xffL) + value = (value << 8) | (data(bufferStart.toInt + i) & 0xffL) i -= 1 } value From ccd469e47b66227564ceca27069540ae47dbfd79 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 29 Aug 2026 09:18:17 -0600 Subject: [PATCH 3/5] fix: relocate the arrow-compression service file when shading arrow-compression ships META-INF/services/org.apache.arrow.vector.compression.CompressionCodec$Factory. The shade plugin copies it verbatim without a ServicesResourceTransformer, so the jar declared a provider for Spark's own unshaded Arrow interface while naming a class that exists here only under the relocated package. Every ServiceLoader lookup Spark's Arrow made then failed with a ServiceConfigurationError, which took CompressionCodec.Factory's static initializer down with it and broke unrelated Arrow IPC reads, including mapInArrow. Add ServicesResourceTransformer so the service file name and its contents are both relocated. arrow-compression is the only bundled artifact that ships one. Also drop an unused NonFatal import that scalafix flagged. --- spark/pom.xml | 10 ++++++++++ .../execution/arrow/ArrowCachedBatchSerializer.scala | 1 - 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/spark/pom.xml b/spark/pom.xml index d30afc75914..5d8ef9d4c69 100644 --- a/spark/pom.xml +++ b/spark/pom.xml @@ -634,6 +634,16 @@ under the License. ${comet.shade.packageName}.guava.thirdparty + + + + 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 343d4f68c85..6679671a05a 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 @@ -20,7 +20,6 @@ package org.apache.spark.sql.comet.execution.arrow import scala.collection.JavaConverters._ -import scala.util.control.NonFatal import org.apache.spark.TaskContext import org.apache.spark.rdd.RDD From e71d8025c58c443e1d185da9ec89c6e3b5c10524 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 29 Aug 2026 13:41:59 -0600 Subject: [PATCH 4/5] test: drop the cache leak test that depends on zstd corruption detection "releases its vectors when a column fails part way through" zeroed the last 16 bytes of a compressed buffer and required the read to fail. Whether that fails is a property of the zstd runtime, not of Comet: the cached payload is byte-identical across Spark versions, but Comet takes zstd-jni from Spark rather than from arrow-compression, and 1.5.5 (Spark 3.4, 3.5) decodes that frame while 1.5.7 (Spark 4.x) reports it corrupt. So the test passed on 4.x and failed on 3.4 and 3.5. The scenario it claimed to cover is also unreachable: CachedBatchIpc decompresses every selected buffer before VectorLoader runs, so no content corruption can fail part way through the load. The two remaining leak tests corrupt a frame from its header onwards, which every zstd release rejects, and already cover a failure at a column's first buffer and a failure after an earlier buffer of the same column decoded. Records the constraint on scramble so a future test does not reach for a tail-only corruption again, and drops the now unused truncateColumn helper and the dictionary fixture's payload argument. --- .../comet/exec/CometInMemoryCacheSuite.scala | 30 +++--------------- .../arrow/CometCachedBatchHelper.scala | 31 ++++++------------- 2 files changed, 14 insertions(+), 47 deletions(-) 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 3464e5e3ddc..3e154094a90 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -1406,12 +1406,12 @@ class CometInMemoryCacheSuite extends CometTestBase { } /** - * Cache two low-cardinality string columns and hand the test the cached payload. + * Cache two low-cardinality string columns and hand the test the cached relation. * * The shuffle is what makes this worth its own fixture: its reader hands the cache writer * dictionary-encoded columns, which the writer has to decode before storing them. */ - private def withDictionaryCache(f: (InMemoryRelation, Array[CachedBatch]) => Unit): Unit = { + private def withDictionaryCache(f: InMemoryRelation => Unit): Unit = { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", @@ -1438,7 +1438,7 @@ class CometInMemoryCacheSuite extends CometTestBase { .cachedRepresentation try { - f(relation, relation.cacheBuilder.cachedColumnBuffers.collect()) + f(relation) } finally { spark.catalog.clearCache() } @@ -1473,7 +1473,7 @@ class CometInMemoryCacheSuite extends CometTestBase { // column either way, so a writer that stored the index vector as-is would hand the loader // integer indices to read as strings. Reading the values back correctly is what proves the // writer decoded them first; a row count alone would not. - withDictionaryCache { (relation, _) => + withDictionaryCache { relation => assert(relation.output.length == 2) val df = spark.sql("SELECT s1, s2 FROM dictionary_cache") @@ -1489,7 +1489,7 @@ class CometInMemoryCacheSuite extends CometTestBase { test("Comet in-memory cache broadcasts a batch read back from the cache") { // A broadcast of a cache scan re-serializes each decoded batch through serializeBatches, // which is a different writer from the one that produced the cached payload. - withDictionaryCache { (relation, _) => + withDictionaryCache { relation => assert(relation.output.length == 2) val df = spark.sql( @@ -1499,26 +1499,6 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - test("Comet in-memory cache releases its vectors when a column fails part way through") { - // Distinct from the corrupted-column case above: there the compressed bytes are wrong from - // their first byte, so the decompressor rejects them outright. Here the bytes start out - // genuine and only the tail is destroyed, so the failure lands after the loader has already - // begun filling vectors. Nothing else can release them -- the holder's constructor never - // returns, so no caller holds it and the task-completion listener has not been told about it. - withDictionaryCache { (relation, batches) => - val cacheSchema = Utils.fromAttributes(relation.output) - batches.foreach(b => CometCachedBatchHelper.truncateColumn(b, cacheSchema, 0)) - - val before = CometArrowAllocator.getAllocatedMemory - interceptDecodeFailure { - decodedRowCount(relation, batches, Seq(relation.output.head)) - } - assert( - CometArrowAllocator.getAllocatedMemory == before, - "a read that fails while loading must release what it already allocated") - } - } - test("Comet in-memory cache scans of one cache canonicalize equal, so exchanges are reused") { // The wrapped Spark scan is a plan-typed field rather than a child, so canonicalization walks // past it and leaves in place the expression IDs of whichever occurrence of the relation 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 e1a089f20f8..f8f9430e268 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 @@ -113,27 +113,6 @@ object CometCachedBatchHelper { scramble(payload(batch), start, length) } - /** - * Zero the tail of one column's compressed bytes, in place, leaving every other column's bytes - * and offsets untouched. - * - * [[corruptColumn]] rewrites the whole compressed payload, which fails as soon as the - * decompressor looks at it. This keeps the leading bytes genuine, so a decoder gets a stream - * that starts out valid and then runs out, exercising a failure part way through a column - * rather than at its first byte. - */ - def truncateColumn(batch: CachedBatch, cacheSchema: StructType, index: Int): Unit = { - val data = payload(batch) - val target = compressedRanges(batch, cacheSchema, index).find { case (_, length) => - length > 32 - } - require( - target.isDefined, - s"column $index of the cached batch has no compressed buffer long enough to truncate") - val (start, length) = target.get - java.util.Arrays.fill(data, (start + length - 16).toInt, (start + length).toInt, 0.toByte) - } - /** * The absolute (start, length) of each of a column's buffers that Arrow actually compressed. * @@ -154,7 +133,15 @@ object CometCachedBatchHelper { } } - /** Overwrite a compressed buffer's payload, leaving its uncompressed-length prefix intact. */ + /** + * Overwrite a compressed buffer's payload, leaving its uncompressed-length prefix intact. + * + * The whole payload is rewritten, frame header included, so every zstd release rejects it + * outright. Corrupting only a frame's tail is not enough: whether that is detected depends on + * the zstd-jni each Spark version ships -- Comet takes it from Spark rather than from + * arrow-compression, and 1.5.5 (Spark 3.4, 3.5) decodes a frame whose last bytes have been + * zeroed that 1.5.7 (Spark 4.x) reports as corrupt. + */ private def scramble(data: Array[Byte], start: Long, length: Long): Unit = { var i = (start + 8).toInt val end = (start + length).toInt From d8d4196f2e7402194ab903895a8bd3421da252e5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 2 Sep 2026 14:33:40 -0600 Subject: [PATCH 5/5] feat: enable Comet's in-memory cache by default Flip spark.comet.exec.inMemoryCache.enabled to true so cached tables are stored and scanned in Comet's Arrow format without an opt-in. CometDriverPlugin.maybeSetCacheSerializer read the config out of SparkConf with a hardcoded false default, so flipping the ConfigEntry alone would have left the serializer uninstalled unless the user set the key explicitly. It now falls back to the entry's own default, matching how the plugin reads spark.comet.metrics.enabled. Stacked on #5543. --- docs/source/user-guide/latest/in-memory-cache.md | 7 ++++--- spark/src/main/scala/org/apache/comet/CometConf.scala | 2 +- spark/src/main/scala/org/apache/spark/Plugins.scala | 4 +++- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md index efd949ff92c..235262eb71d 100644 --- a/docs/source/user-guide/latest/in-memory-cache.md +++ b/docs/source/user-guide/latest/in-memory-cache.md @@ -24,10 +24,11 @@ format that Comet operators read directly. Without it, a cached table is stored format and every scan of it has to convert each batch before Comet can continue, which shows up in the plan as a `CometSparkColumnarToColumnar` above the cache scan. -This feature is **experimental and disabled by default**. +This feature is **experimental and enabled by default**. To turn it off, set the config before the +`SparkContext` is created: ```scala -spark.conf.set("spark.comet.exec.inMemoryCache.enabled", "true") +spark.conf.set("spark.comet.exec.inMemoryCache.enabled", "false") ``` ## What changes when it is enabled @@ -85,7 +86,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.enabled` | `true` | Whether to store and scan Spark's in-memory cache in Comet's format. Read at startup. | | `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 4183ffef457..92bad31952f 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -276,7 +276,7 @@ object CometConf extends ShimCometConf { "SparkContext, otherwise caching fails as soon as a block is serialized, including " + "the disk half of the default MEMORY_AND_DISK storage level.") .booleanConf - .createWithDefault(false) + .createWithDefault(true) val COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC: ConfigEntry[String] = conf("spark.comet.exec.inMemoryCache.compression.codec") diff --git a/spark/src/main/scala/org/apache/spark/Plugins.scala b/spark/src/main/scala/org/apache/spark/Plugins.scala index eaeac316655..9736d523633 100644 --- a/spark/src/main/scala/org/apache/spark/Plugins.scala +++ b/spark/src/main/scala/org/apache/spark/Plugins.scala @@ -130,7 +130,9 @@ object CometDriverPlugin extends Logging { private[apache] def maybeSetCacheSerializer( conf: SparkConf, extraConfs: ju.HashMap[String, String]): Unit = { - if (conf.getBoolean(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key, false)) { + if (conf.getBoolean( + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key, + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.defaultValue.get)) { val serializerKey = StaticSQLConf.SPARK_CACHE_SERIALIZER.key val serializerValue = "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer"