sunchao commented on code in PR #5543: URL: https://github.com/apache/datafusion-comet/pull/5543#discussion_r4066604932
########## spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala: ########## @@ -0,0 +1,586 @@ +/* + * 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()) + } + + // Unloaded plain and compressed afterwards rather than by handing the codec to the unloader; + // see compressed for why. + val unloader = new VectorUnloader(root, true, NoCompressionCodec.INSTANCE, true) + val plainBatch = unloader.getRecordBatch + val recordBatch = + try compressed(plainBatch, codec, allocator) + finally plainBatch.close() + try { + val fields = vectors.map(_.getField) + // Leaves the batch in the state serializeBatches leaves one. The record batch holds its + // own buffers by now, so this does not touch it, and getField still answers afterwards: + // clearing releases buffers, not the schema. + 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) Review Comment: **[P2] Preserve support for batches larger than a single JVM array** Could we retain chunked storage or split oversized batches by bytes before writing this payload? The row converter limits `conf.columnBatchSize` records but has no aggregate byte limit, so valid binary/string-heavy batches can exceed 2 GiB while every individual column still fits its Arrow buffers. The previous per-column `ChunkedByteBuffer` representation supported that case. I checked this using the production row-to-Arrow converter and the exact-head IPC writer. A valid multi-column batch above that aggregate limit failed here with `OutOfMemoryError: Requested array size exceeds VM limit` under both default zstd and `none`. Replaying the previous per-column writer sequence successfully stored the same input in chunks. This was a component comparison, not a full base-revision Spark query. Clamping the capacity to `Int.MaxValue` still requests an array the JVM rejects, regardless of available heap. Reducing only the initial capacity would also leave `toByteArray` unable to represent the payload. Please handle the size boundary before constructing the single-array payload and add regression coverage. ########## 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: Confirmed fixed at `6b8ded1dff647a0554aff007a74cc12123a6c6d0`. I reran the real zstd workspace-allocation failure from the earlier review. The failed write now retains only the original input allocation, and closing the input returns Arrow allocation to zero. The new compression-failure regression test also passed in the focused suite. ########## 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: Confirmed fixed at `6b8ded1dff647a0554aff007a74cc12123a6c6d0`. The example now enables the cache with a startup `--conf`, and the explanation matches the serializer selection in `CometDriverPlugin` during SparkContext initialization. -- 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]
