sunchao commented on code in PR #5543: URL: https://github.com/apache/datafusion-comet/pull/5543#discussion_r4018641740
########## docs/source/user-guide/latest/in-memory-cache.md: ########## @@ -0,0 +1,170 @@ +<!--- + 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**. + +```scala +spark.conf.set("spark.comet.exec.inMemoryCache.enabled", "true") Review Comment: **[P3] Configure the cache before SparkContext startup** Starting from the documented default of `false`, this runtime setting cannot install Comet's cache serializer. `CometDriverPlugin` chooses `spark.sql.cache.serializer` once during SparkContext initialization, so subsequent cached relations still use Spark's default format and cannot take the native cache path. I checked this with the current plugin on Spark 4.1.3: startup `false` followed by this runtime setting kept `DefaultCachedBatchSerializer`; startup `true` installed `ArrowCachedBatchSerializer`. Please replace this activation example with `--conf spark.comet.exec.inMemoryCache.enabled=true` on startup, or builder configuration before creating a fresh SparkContext. ########## spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala: ########## @@ -0,0 +1,525 @@ +/* + * 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 + +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") + } + + // 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. + * + * 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)) + } + + /** + * Whether `batch`'s vectors can be unloaded as they stand, or have to be converted first. + * + * The payload records no schema, so [[Projection]] rebuilds the fields from the cached + * relation's Spark attributes and reads the body against them. The direct write path unloads + * whatever vectors the cached plan produced, and one Spark type can arrive as more than one + * Arrow type: `BinaryType` is a `VarBinaryVector` from Comet's own scans but a + * `FixedSizeBinaryVector` from an accelerated `mapInArrow` or an Iceberg `fixed[N]` read, and + * those occupy three buffers and two. Writing one and reading the other shifts every buffer + * from that column on, which is wrong values rather than an error, so a batch that does not + * already carry the reader's types is converted instead. + * + * The same holds inside a nested column, which `Utils.isArrowBacked` does not look at: it + * answers for the top-level vector only, so a struct of large strings passes it while its child + * is stored with 64-bit offsets and read with 32-bit ones. + * + * Names, nullability and a timestamp's timezone are not compared. None of them changes how the + * reader interprets the body, and the writer's legitimately differ -- a Comet scan labels + * timestamps with the session's zone where the reader rebuilds them as UTC, which is a label + * only: Spark's representation is micros since the epoch either way. + */ + def matchesReaderLayout(batch: ColumnarBatch, readerFields: Seq[Field]): Boolean = + batch.numCols() == readerFields.length && + (0 until batch.numCols()).forall { i => + batch.column(i) match { + case v: CometVector => sameLayout(writtenField(v), readerFields(i)) + case _ => false + } + } + + /** + * The field a column reaches the body as. + * + * A dictionary-encoded vector's own field carries the index type, not the values', because + * [[decodeDictionaries]] replaces it with the decoded form before anything is unloaded. + * Resolved through the same `lookupDictionary` the write path uses, so a batch missing its + * dictionary fails here exactly as it would there. + */ + private def writtenField(column: CometVector): Field = { + val vector = column.getValueVector + if (vector.getField.getDictionary == null) { + vector.getField + } else { + Utils + .lookupDictionary(vector.asInstanceOf[FieldVector], Option(column.getDictionaryProvider)) + .getVector + .getField + } + } + + private def sameLayout(written: Field, read: Field): Boolean = + layoutType(written.getType) == layoutType(read.getType) && { + val writtenChildren = written.getChildren + val readChildren = read.getChildren + writtenChildren.size == readChildren.size && + (0 until writtenChildren.size).forall(i => + sameLayout(writtenChildren.get(i), readChildren.get(i))) + } + + private def layoutType(t: ArrowType): ArrowType = t match { + case ts: ArrowType.Timestamp if ts.getTimezone != null => + new ArrowType.Timestamp(ts.getUnit, "UTC") + case other => other + } + + /** + * 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, decoded) = decodeDictionaries(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 Projection.load reproduces when it + // repacks the selected buffers. + val unloader = new VectorUnloader(root, true, codec, true) + val recordBatch = unloader.getRecordBatch Review Comment: **[P2] Release partial Arrow allocations when compression fails** `VectorUnloader.getRecordBatch` retains each input buffer and accumulates compressed buffers without cleaning them up if a later compression fails. That failure happens before this method's `recordBatch.close()` finally becomes active. Closing the input batch afterwards does not release those extra references or the earlier compressed buffers, so a failed cache materialization can leave batch-sized off-heap allocations behind. I reproduced this against the current `CachedBatchIpc.serialize` and Arrow 18.3.0 with 4,194,304 integer values, an unlimited `RootAllocator`, and valid zstd level 22. After warming the child JVM, limiting its virtual address space to its current size plus 64 MiB caused zstd's native workspace allocation to fail with an ordinary `RuntimeException: Error compressing: Allocation error : not enough memory`. After closing the input batch, **18,350,080 bytes (17.5 MiB) remained allocated**. The JVM remained usable. Under equivalent pressure, a component using the previous single-column stream-writing sequence threw `ZstdIOException` and returned Arrow allocation to zero after input close. Please make write-side compression explicitly own and release partial buffers on failure, similar to the guarded decompression path, and add a failure-path regression test. Merely closing the input/root in a finally will not undo the unloader's retained references. -- 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]
