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..efd949ff92c --- /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` | 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`; +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) | 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 +`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..5d8ef9d4c69 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} @@ -630,6 +634,16 @@ under the License. ${comet.shade.packageName}.guava.thirdparty + + + + 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..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,12 +20,12 @@ 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 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} @@ -33,25 +33,26 @@ 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.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 /** @@ -75,12 +76,29 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { 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) @@ -91,16 +109,17 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { 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 } } @@ -130,9 +149,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 +163,14 @@ 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. 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 => + _: DecimalType | _: StringType | DateType | TimestampType | TimestampNTZType => true case _ => false } @@ -160,36 +186,10 @@ 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 } - // 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 StringType => - ByteArray.compareBinary( - left.asInstanceOf[UTF8String].getBytes, - right.asInstanceOf[UTF8String].getBytes) - 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. @@ -197,34 +197,47 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { // 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]): Iterator[CachedBatch] = { + attrs: Seq[Attribute], + codecSetting: (String, Int)): Iterator[CachedBatch] = { val arrowSchema = Utils.toArrowSchema(Utils.fromAttributes(attrs), CometArrowStream.NATIVE_TIMEZONE) + 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, which serializing then clears, so - // they have to be gathered first. The row is only assembled once the per-column sizes are - // known. - val (lower, upper, nulls) = gatherColumnStats(batch, attrs) + // 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, orderings) 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 +306,19 @@ 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] = { + val codec = codecSettings(conf) + input.mapPartitions { batches => - encodeBatches(batches, schema) + encodeBatches(batches, schema, codec) } } @@ -320,24 +336,36 @@ 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. + // 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 + .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 + // 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 +376,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, projection) + current = projected + projected.batches } case other => @@ -360,78 +388,36 @@ 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 - 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 - } - opened - } + /** + * Owns the Arrow vectors decoded for one cached batch. + * + * 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, 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 + // 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 { 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 +437,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { throw new NoSuchElementException } emitted = true - assemble() + NativeUtil.rootAsBatch(root) } } } @@ -467,24 +453,25 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { fallback.convertInternalRowToCachedBatch(input, schema, storageLevel, conf) } else { val batchSize = conf.columnBatchSize + val codec = codecSettings(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, codec) } } } @@ -551,6 +538,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..43a78ab858b --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala @@ -0,0 +1,456 @@ +/* + * 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, 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, 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 + +/** + * 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") + } + + // 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. + * + * 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() + + // 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)) + } 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(_) => () }) + } + } + + /** + * Everything about reading one projection of this format that does not change between batches. + * + * 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. + * + * 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. + */ + final class Projection(arrowFields: Seq[Field], selectedIndices: Array[Int]) { + + private val schema = new Schema(selectedIndices.map(arrowFields).toSeq.asJava) + + // 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 + + 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 variadicCounts = new java.util.ArrayList[java.lang.Long](variadicIndices.length) + if (batch.variadicBufferCountsLength() > 0) { + variadicIndices.foreach(j => variadicCounts.add(batch.variadicBufferCounts(j))) + } + + 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 + } + 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 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) => + try root.close() + catch { case NonFatal(closeError) => e.addSuppressed(closeError) } + throw e + } finally { + plainBatch.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)) + } + + /** + * 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 + // 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 { + batch.getBuffers.asScala.foreach { buffer => + buffer.getReferenceManager.retain() + val plain = + try { + // An empty buffer carries no compressed length prefix to read. + codec match { + case Some(c) if buffer.writerIndex() > 0 => c.decompress(allocator, buffer) + case _ => 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 + } + } + + /** + * 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) => + if (vector.getField.getDictionary == null) { + vector + } else { + val dictionary = Utils.lookupDictionary(vector, providerOpt) + 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..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 @@ -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], @@ -479,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] = { @@ -494,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. @@ -517,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/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index a7a9a6c8590..3e154094a90 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,33 +1382,36 @@ 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") } } /** - * 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, 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 = { + private def withDictionaryCache(f: InMemoryRelation => Unit): Unit = { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", @@ -1301,52 +1438,64 @@ class CometInMemoryCacheSuite extends CometTestBase { .cachedRepresentation try { - f(relation, relation.cacheBuilder.cachedColumnBuffers.collect()) + f(relation) } finally { spark.catalog.clearCache() } } } - 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 /*+ BROADCAST(c) */ c.s1, c.s2 FROM range(1) r JOIN dictionary_cache c ON true") + val df = spark.sql("SELECT s1, s2 FROM dictionary_cache") checkSparkAnswer(df) - assert(df.count() == 2000) + + 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 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. - 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)) + 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 before = CometArrowAllocator.getAllocatedMemory - intercept[Exception] { - decodedRowCount(relation, batches, Seq(relation.output.head)) - } - assert( - CometArrowAllocator.getAllocatedMemory == before, - "a reader that fails while opening must release what it already allocated") + val df = spark.sql( + "SELECT /*+ BROADCAST(c) */ c.s1, c.s2 FROM range(1) r JOIN dictionary_cache c ON true") + checkSparkAnswer(df) + assert(df.count() == 2000) } } 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..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 @@ -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.{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 +import org.apache.spark.sql.types.StructType /** * Test-only access to the internals of `CometCachedBatch`. @@ -37,83 +39,177 @@ 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 - /** Serialized size of each column stream, in column order. */ - def columnStreamSizes(batch: CachedBatch): Seq[Long] = - batch.asInstanceOf[CometCachedBatch].columns.map(_.size).toSeq + /** + * 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 = + 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] = + 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 = + compressedRanges(batch, cacheSchema, index).nonEmpty + + /** + * 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. + * + * Requires the column to have a genuinely compressed buffer -- see [[columnIsCompressed]]. + */ + def corruptColumn(batch: CachedBatch, cacheSchema: StructType, index: Int): Unit = { + val ranges = compressedRanges(batch, cacheSchema, index) + require( + 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) } } - /** 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 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) + } + + /** + * The absolute (start, length) of each of a column's buffers that Arrow actually compressed. + * + * `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. + */ + private def compressedRanges( + batch: CachedBatch, + cacheSchema: StructType, + index: Int): Seq[(Long, Long)] = { + val data = payload(batch) + 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) } + } /** - * Drop the last `dropBytes` of one column's decoded Arrow stream, in place. + * Overwrite a compressed buffer's payload, leaving its uncompressed-length prefix intact. * - * [[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. + * 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 + 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. */ - 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) - } - buffer.toByteArray - } finally { - decodedStream.close() + 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()) } - 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() } - 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 + + /** 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") + 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: Long): Long = { + var value = 0L + var i = 7 + while (i >= 0) { + value = (value << 8) | (data(bufferStart.toInt + i) & 0xffL) + i -= 1 + } + value + } }