andygrove commented on code in PR #5543: URL: https://github.com/apache/datafusion-comet/pull/5543#discussion_r4085877008
########## docs/source/user-guide/latest/in-memory-cache.md: ########## @@ -0,0 +1,184 @@ +<!--- + 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. Review Comment: Applied, thanks. ########## docs/source/user-guide/latest/in-memory-cache.md: ########## @@ -0,0 +1,184 @@ +<!--- + 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`, for footprint rather than for speed. Over the same 5M-row, +six-column relation the tables under [Performance](#performance) use — and measured by the same +benchmark — it holds the data in a sixth of the memory and pays for that on both sides: about 40% +longer to materialize, and, on a read wide enough to inflate everything, close to five times +longer. A narrow projection pays far less, because it only inflates the columns it asked for. + +| Codec | Materialize | Footprint | Read 1 of 6 | Read 6 of 6 | +| ------ | ----------: | --------: | ----------: | ----------: | +| `zstd` | 1503 ms | 55 MiB | 47 ms | 296 ms | +| `none` | 1091 ms | 315 MiB | 36 ms | 63 ms | + +`none` is the better setting for a relation that fits in memory uncompressed and is read at close +to full width. The default is the other way round because a cache that does not fit costs more than +one that is slower to read, and Spark's own cache format compresses by default too. + +Arrow's other IPC codec, LZ4, is deliberately not offered and the config rejects it. It is +commons-compress's pure-Java implementation, unrelated to the JNI-accelerated lz4-java behind +`spark.io.compression.codec`, and is orders of magnitude slower to write than `zstd` while also +producing larger output. + +Dictionary-encoded columns are decoded before they are stored. A payload with no schema message has +nowhere to record either that a column is dictionary encoded or the dictionary itself. + +## Configuration + +| Config | Default | Description | +| ------------------------------------------------------- | ------- | ---------------------------------------------------------------------------------------------------------------------------------------------- | +| `spark.comet.exec.inMemoryCache.enabled` | `false` | Whether to store and scan Spark's in-memory cache in Comet's format. Read at startup. | +| `spark.comet.exec.inMemoryCache.compression.codec` | `zstd` | Arrow IPC compression codec for cached data: `zstd` or `none`. Affects newly cached data only — a batch records the codec it was written with. | +| `spark.comet.exec.inMemoryCache.compression.zstd.level` | `1` | Compression level when the codec is `zstd`. Ignored otherwise. | + +## Performance + +Measured with `CometInMemoryCacheBenchmark` (Apple M3 Max, JDK 17, Spark 4.1, release build). +Regenerate with: + +```sh +SPARK_GENERATE_BENCHMARK_FILES=1 \ + make benchmark-org.apache.spark.sql.benchmark.CometInMemoryCacheBenchmark +``` + +On a 5M-row relation of six flat columns: + +| Query shape | Spark cache scan + convert | `CometInMemoryTableScan` | Relative | +| ------------------------------ | -------------------------: | -----------------------: | -------: | +| Repeated scan (3 of 6 columns) | 180 ms | 147 ms | 1.2x | +| Selective filter | 58 ms | 50 ms | 1.2x | +| Row count only (0 of 6) | 37 ms | 34 ms | 1.1x | +| Narrow projection (1 of 6) | 52 ms | 48 ms | 1.1x | +| Full projection (6 of 6) | 547 ms | 291 ms | 1.9x | + +And on a 1M-row relation of six columns whose middle three are structs, one of them nested two +levels deep: + +| Query shape | Spark cache scan + convert | `CometInMemoryTableScan` | Relative | +| -------------------------- | -------------------------: | -----------------------: | -------: | +| Row count only (0 of 6) | 30 ms | 29 ms | 1.0x | +| Narrow projection (1 of 6) | 98 ms | 53 ms | 1.8x | +| Full projection (6 of 6) | 250 ms | 125 ms | 2.0x | + +Both columns read the cache at the default codec, `zstd`. The codec table above shows what `none` +changes, and it is the full projection that moves most: nothing has to be inflated, so it runs +several times faster, at six times the memory. + +The two relations are not comparable to each other — different row counts, and a struct column +carries several values per row. + +Array and map columns are deliberately absent from the benchmark, not from the format — the cache +stores and projects them, and `CometInMemoryCacheSuite` covers them. They cannot be measured _here_ +because the left column would not exist: it needs Spark's cache scan to bridge into Comet operators, +and `CometSparkToColumnarExec` declines `ArrayType` and `MapType`, so a query projecting one falls +back to Spark row execution above the scan and the two columns stop measuring the same boundary. + +Read what this compares carefully. Comet execution is on in both columns, so the aggregation runs +on Comet either way and only the cache-scan boundary moves: on the left, Spark's +`InMemoryTableScanExec` feeds those same Comet operators through a `CometSparkColumnarToColumnar` +bridge; on the right, `CometInMemoryTableScan` feeds them directly. Both columns read the same +Comet-written `CometCachedBatch` — `spark.sql.cache.serializer` is static, so one session cannot +also materialize Spark's format to compare against. These numbers are therefore "keep the cached +scan native" against "fall back to a Spark cache scan and convert", not Comet against Spark +execution, and not a comparison with Spark's own cache format. + +## Kryo + +Spark serializes a cached batch with `spark.serializer` whenever the block leaves the heap: the +`_SER` storage levels, replication, cross-executor fetches, and the disk half of the default +`MEMORY_AND_DISK`. So an ordinary `df.cache()` that spills is enough to reach it. + +If you run with `spark.kryo.registrationRequired=true`, register Comet's classes: + +``` +spark.serializer=org.apache.spark.serializer.KryoSerializer +spark.kryo.registrationRequired=true +spark.kryo.registrator=org.apache.comet.CometKryoRegistrator +``` + +Comet cannot set `spark.kryo.registrator` for you the way it sets `spark.sql.cache.serializer`: +`KryoSerializer` reads it when `SparkEnv` builds the serializer, which happens before any plugin +runs. Without it, caching fails with a "Class is not registered" error that does not name this +feature. Comet's driver plugin warns at startup when it sees Kryo, `registrationRequired`, and no +registrator. + +## Limitations + +Reads that feed **Spark** operators rather than Comet ones are still slower than Spark's own cache +format, by roughly 1.7x to 2.5x depending on how wide the projection is. Those reads pay a row +conversion that Spark's format avoids with generated code over its own layout. This is why the +feature is off by default. Review Comment: It described the old format. The range came from #5485, which measured the per-column format with Spark's whole-stream codec, and the page quoted it without saying so. I've re-measured it on this format, and this time in `CometInMemoryCacheBenchmark` itself: a new axis reads the same 5M-row relation with Comet off, cached once in each format. Both formats can come out of one session because `InMemoryRelation` looks the serializer up in the active conf when a relation is cached, and each case asserts which serializer cached what it reads. The gap is wider than the page said and widest for the narrowest reads: 5.2x for a row count (183 ms against 35 ms), 4.8x for one column, 3.4x for three, and 1.5x for all six. The page now shows that table and links #5485, and no longer gives a cause, since #5485 says it isn't established. ########## 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)) + } + + /** + * 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. + * + * The message is written into heap chunks of `chunkSize` bytes rather than one array, so a + * batch is not bounded by the 2 GiB a single JVM array can hold. Nothing bounds a cached + * batch's bytes upstream: the row path cuts batches by row count alone, and a columnar input + * batch arrives at whatever size the plan above produced. Chunks are appended rather than grown + * and recopied, so the write also never holds the payload twice. + * + * Returns the message 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, + chunkSize: Int): (ChunkedByteBuffer, 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() + + val out = new ChunkedByteBufferOutputStream(chunkSize, ByteBuffer.allocate) + try { + val channel = new WriteChannel(Channels.newChannel(out)) + MessageSerializer.serialize(channel, recordBatch) + } finally { + out.close() + } + (out.toChunkedByteBuffer, columnSizes(fields, recordBatch)) + } finally { + recordBatch.close() + } + } finally { + // Only the vectors this method allocated. The rest belong to the input batch. + decoded.foreach(v => + try v.close() + catch { case NonFatal(_) => () }) + } + } + + /** + * Everything about reading one projection of this format that does not change between batches. + * A scan builds one of these per partition. + * + * The index arithmetic walks every field of the cached relation rather than just the projected + * ones, so recomputing it per batch would make the bookkeeping O(total columns) against + * O(selected columns) of useful work -- worst in exactly the wide-relation, narrow-projection + * case this format exists for. + * + * Holding the projected `Schema` here too is what keeps it consistent with the buffers: + * [[Projection.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: IndexedSeq[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 run in the same order. The totals + // are what a payload is checked against in load. + private val (nodeIndices, totalNodes) = + selectedRange(arrowFields, selectedIndices, fieldNodeCount) + private val (bufferIndices, totalBuffers) = + selectedRange(arrowFields, selectedIndices, fieldBufferCount) + + /** + * Decode the projected columns of one cached payload into a fresh root the caller owns. + * + * A buffer's recorded (offset, length) covers its on-body bytes including the + * uncompressed-length prefix, so a window copied out of the payload is exactly what the + * writer emitted, 8-byte aligned as Arrow's IPC body lays it out. The columns that were not + * selected are never read, let alone inflated. The windows are then decompressed in one pass; + * see [[decompressed]] for why that is not left to `VectorLoader`. + */ + def load(data: ChunkedByteBuffer, allocator: BufferAllocator): VectorSchemaRoot = { + val readChannel = new ReadChannel(Channels.newChannel(data.toInputStream())) + // Reads the message metadata only. The body stays in `data` and is copied selectively. + val metadata = MessageSerializer.readMessage(readChannel) + if (metadata == null) { + throw new SparkException("Unexpected end of input reading a Comet cached batch") + } + val batch = + metadata.getMessage.header(new FlatBufRecordBatch()).asInstanceOf[FlatBufRecordBatch] + + // The payload carries no schema, so nothing in it says the writer laid the body out the way + // these windows read it. batch.buffers(j) is an unchecked flatbuffer accessor, so a + // disagreement would otherwise surface as wrong values, or as an out-of-range read from + // inside the copy below, rather than as an error naming the cause. See matchesReaderLayout + // for how the write path avoids producing one. + if (batch.nodesLength() != totalNodes || batch.buffersLength() != totalBuffers) { + throw new SparkException( + "Comet cached batch does not match the cached schema: the payload holds " + + s"${batch.nodesLength()} field nodes and ${batch.buffersLength()} buffers, but the " + + s"schema describes $totalNodes and $totalBuffers") + } + + // serialize writes exactly [encapsulated message][body] and nothing after it, so the body is + // the tail of `data`. + val bodyStart = data.size - metadata.getMessageBodyLength + val chunks = new PayloadChunks(data) + + val compression = + if (batch.compression() == null) NoCompressionCodec.DEFAULT_BODY_COMPRESSION + else new ArrowBodyCompression(batch.compression().codec(), batch.compression().method()) + + val nodes = new java.util.ArrayList[ArrowFieldNode](nodeIndices.length) + nodeIndices.foreach { j => + val node = batch.nodes(j) + nodes.add(new ArrowFieldNode(node.length(), node.nullCount())) + } + + val offsets = new Array[Long](bufferIndices.length) + val lengths = new Array[Long](bufferIndices.length) + var total = 0L + var k = 0 + while (k < bufferIndices.length) { + val buffer = batch.buffers(bufferIndices(k)) + offsets(k) = buffer.offset() + lengths(k) = buffer.length() + total += DataSizeRoundingUtil.roundUpTo8Multiple(lengths(k)) + k += 1 + } + + // allocator.buffer(0) is legal but yields a buffer no window can be sliced from, and an + // all-empty projection (every selected column a NullVector, say) would ask for exactly that. Review Comment: Added both. The NullType test caches `SELECT id, NULL AS n` under both codecs and reads `SELECT n` and `SELECT id, n` against rows collected before caching. The zero-row test serializes an int and a string column with no rows under both codecs, loads both columns together and each on its own through `Projection.load`, and checks the allocator is back to zero afterwards. Both paths already worked, so these are coverage rather than a fix. The comment you pointed at is gone, though: empty windows now take Arrow's shared empty buffer, which is what its own codec returns for them, so an all-empty projection allocates nothing. -- 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]
