andygrove commented on code in PR #5543: URL: https://github.com/apache/datafusion-comet/pull/5543#discussion_r4020819764
########## 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: You're right, and the repro made this quick to confirm. `appendNodes` retains each input buffer and then does `buffers.add(codec.compress(...))` into a list that lives inside `getRecordBatch`, so when a compress throws, that retain is stranded and every buffer compressed before it is unreachable from anywhere a caller can get to. Closing the input batch afterwards releases the input's own reference, not the extra retain, and it never sees the compressed buffers at all. Fixed by not handing the codec to the unloader. `serialize` now unloads with `NoCompressionCodec`, which retains and hands each buffer straight back without allocating, and a new `CachedBatchIpc.compressed` does the compression in the same shape `decompressed` already uses on the read side: retain, compress inside a try that releases on a throw, and an outer catch that closes whatever is already in the list. Peak memory is unchanged, since the unloader was doing the same thing one buffer at a time. The retain is worth a note because it earns its keep differently in the two cases. A codec that allocates consumes the retained reference and returns a buffer of its own; `NoCompressionCodec` returns the input, and the retain is then the reference the result batch ends up owning. Either way the input buffer is back at its original count once the result is closed. There is a regression test, and I mutated the fix away to check it isn't passing for free: it then fails with 2,176 bytes still allocated after the input is closed, at 256 rows, which is the same thing you measured at 17.5 MiB with a real workload. The codec it uses is a real zstd codec that throws on the third buffer, so the int column's two are genuine allocations by the time it does. Failing at the first buffer would have passed with no cleanup at all. ########## 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: Right, and the page contradicted itself: the paragraph immediately below the example said the config is read at startup. It is now a startup `--conf` on `spark-shell`, with a sentence saying why -- the driver plugin picks `spark.sql.cache.serializer` while the `SparkContext` is initializing, so a session that started with the default goes on using Spark's format however the config is set afterwards. The "is read at startup" sentence in the next section came out, since it was then saying the same thing twice. -- 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]
