andygrove commented on code in PR #5543: URL: https://github.com/apache/datafusion-comet/pull/5543#discussion_r4018236048
########## 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) Review Comment: You're right, and the `FixedSizeBinaryVector` case is real: `toArrowType(BinaryType)` is `Binary` at three buffers against that vector's two, so every buffer index from such a column on shifts. The length check is in, built the way you suggested. `selectedRange` now returns `starts.last` alongside the indices, and `load` compares both totals against `nodesLength()` and `buffersLength()` before it touches `batch.buffers(j)`. I mutated the check away to confirm it earns its place: the new test for it then passes with no exception at all, which is exactly the silent-wrong-answer failure you described. I went wider than dropping `FixedSizeBinaryVector` from `isArrowBacked`, because that only closes the top-level case. `isArrowBacked` answers for the top-level vector and never looks at children, so a struct whose child is a `LargeVarCharVector` passes it today and is stored with 64-bit offsets and read back with 32-bit. The length check can't catch that one either, since `LargeUtf8` and `Utf8` are both three buffers. So the write path now asks the direct question: do this batch's vectors already carry the Arrow types the reader will rebuild, recursively? That's `CachedBatchIpc.matchesReaderLayout`, and a batch that disagrees takes the conversion path it was already taking for non-Arrow input rather than being written unreadable. Names, nullability and a timestamp's timezone are excluded from the comparison — the last because a Comet scan labels with the session zone where the reader rebuilds UTC, and that's a label rather than a layout, so comparing it would send every timestam p column down the conversion path for nothing. One wrinkle: Arrow-Java puts the index type on a dictionary-encoded vector's own field, so the field to compare is the dictionary's, resolved through the same `Utils.lookupDictionary` the writer uses. On coverage, a unit test builds a `FixedSizeBinaryVector`-backed batch and asserts `isArrowBacked` accepts it while the write path declines it, with a `VarBinaryVector` beside it as the control so the predicate can't pass by refusing everything. The reader's check gets its own test that hands the reader one more attribute than the writer stored. -- 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]
