andygrove commented on code in PR #5543: URL: https://github.com/apache/datafusion-comet/pull/5543#discussion_r4083388327
########## docs/source/user-guide/latest/in-memory-cache.md: ########## @@ -0,0 +1,176 @@ +<!--- + 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. +--> + +# 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**. Turn it on at startup, alongside the rest +of Comet's configuration: + +```shell +$SPARK_HOME/bin/spark-shell \ + ... \ + --conf spark.comet.exec.inMemoryCache.enabled=true +``` + +It has to be set before the `SparkContext` starts. Comet's driver plugin chooses +`spark.sql.cache.serializer` once, while the context is initializing, so a session that started +with the default goes on using Spark's cache format however the config is set afterwards. + +## What changes when it is enabled + +With Comet's serializer installed as `spark.sql.cache.serializer`: + +- 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 | Review Comment: You were right to doubt it. That table didn't come from a committed benchmark, and it doesn't survive one: with each case warmed up, `none` beats zstd on everything except footprint, so it was supporting the default with the wrong argument. `CometInMemoryCacheBenchmark` now has a codec axis over the 5M-row flat relation: materialize time with the timer around the caching alone, footprint from the relation's size accumulator, and one-column and six-column reads. zstd takes 1503 ms to materialize against 1091 ms for `none`, reads one column in 47 ms against 36 ms and all six in 296 ms against 63 ms, and holds 55 MiB against 315 MiB. The ~410 ms materialize gap also matches the write path on its own: timing `CachedBatchIpc.serialize` directly puts zstd about 87 ns/row above `none`, ~435 ms over 5M rows. Every table on the page is now regenerated from one run. So zstd is a footprint default, not a speed one, and the docs and the `compression.codec` doc string now say that, including when `none` is the better choice. I've kept zstd as the default, since a cache that doesn't fit costs more than one that reads slower and Spark's own format compresses by default too, but I'm open to flipping it. The LZ4 claim came from the same kind of one-off, so I checked it with the same direct timing, since the config rejects that codec and the benchmark can't reach it. It holds: 7.9 s to serialize one 10k-row batch against 0.96 ms for zstd, with 2.5x larger output. ########## spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala: ########## @@ -0,0 +1,633 @@ +/* + * 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.nio.ByteBuffer +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 +import org.apache.spark.util.io.{ChunkedByteBuffer, ChunkedByteBufferOutputStream} + +import org.apache.comet.vector.CometVector + +/** + * 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 [[Projection.load]] 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") + } + + // 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. + * + * A byte this build does not recognize is rejected rather than read as plain bytes. + * `CodecType.fromCompressionType` answers `NO_COMPRESSION` for anything outside its enum, so + * taking its word for it would turn a corrupt payload into garbage values instead of an error. + */ + private def readCodec(compressionType: Byte): Option[CompressionCodec] = + if (compressionType == NoCompressionCodec.COMPRESSION_TYPE) { + None + } else { + val codecType = CompressionUtil.CodecType.fromCompressionType(compressionType) + if (codecType == CompressionUtil.CodecType.NO_COMPRESSION) { + throw new SparkException( + s"Comet cached batch records an unknown Arrow compression codec: $compressionType") + } + Some(readCodecs(codecType)) + } Review Comment: Added, built the way you suggested. `CometCachedBatchHelper.cachedBatchWithBodyCompression` unloads a batch plain, tags it with `new ArrowBodyCompression(99, BodyCompressionMethod.BUFFER)` and writes it with `MessageSerializer.serialize` into a `ChunkedByteBuffer`. The layout is exactly what the writer produces, so the read gets past the layout check to the codec lookup. The test reads it through `convertCachedBatchToColumnarBatch` and asserts that the failure names the codec and surfaces as itself rather than as a reference-count error from cleanup. It checks the message rather than the type, because the job abort wraps it in a `SparkException` either way. With `readCodec` put back to trusting `fromCompressionType`, it fails with "no exception was thrown". ########## spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala: ########## @@ -1035,12 +1042,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. Review Comment: Applied, thanks. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
