feat: add experimental native support for in-memory cache, disabled by default - #5051
feat: add experimental native support for in-memory cache, disabled by default#5051andygrove wants to merge 24 commits into
Conversation
# Conflicts: # spark/src/main/scala/org/apache/spark/Plugins.scala
InMemoryRelation resolves spark.sql.cache.serializer once per JVM and memoizes the instance in a static field. Test suites share a forked JVM, so whichever suite caches a table first pins the serializer for every suite that follows. CometInMemoryCacheSuite runs after CometExecSuite in the exec group, so its configured serializer was ignored, every cached batch was a DefaultCachedBatch, and 9 of its 11 tests failed on all Spark profiles. Reset the memoized serializer around the suite, via a test-only shim for the private[columnar] clearSerializer. Also address review feedback: - Fall through to the SparkToColumnar path when the native cache is enabled but the relation was cached by a foreign serializer, so enabling the feature is never worse for a scan than leaving it off. Covered by a new CometExecSuite test. - Explain on CometInMemoryTableScanExec.output why the declared output can be narrower than the emitted batch width for an empty projection. - Drop the import made redundant by the org.apache.spark.sql.comet._ wildcard.
andygrove
changed the title
|
I found some issues with data type coverage and will push some fixes soon. |
…serializer Three defects found while reviewing type coverage, each with a regression test that fails without the corresponding fix. Pruning dropped every batch for columns without bounds. `tracksBounds` only computes bounds for the types it lists, so a collated `StringType` on Spark 4.x left them null. Spark still builds a partition filter for such a column because a collated string literal is an `AtomicType`, and comparing against null bounds yields null, which the generated predicate treats as false. A filtered query over a collated string column returned zero rows instead of the matching ones, with no error. `buildFilter` now drops predicates over columns that have no bounds, keeping `IsNull` and `IsNotNull` which only read null and row counts. Interval columns broke `cache()`. `Utils.getFieldVector` has no case for `DurationVector` or `IntervalYearVector`, so materializing a cached relation containing a `DayTimeIntervalType` or `YearMonthIntervalType` column threw `Unsupported Arrow Vector for serialize`. The serializer now reports the schemas its Arrow writer supports and delegates the rest to Spark's default cache serializer. Reading a `DefaultCachedBatch` failed for any non-primitive column. `supportsColumnarOutput` returned true unconditionally while `decodeDefaultCachedBatch` decoded through `ColumnAccessor.decompress`, whose `PassThrough` decoder only handles the seven primitive physical types. Spark's own serializer gates `supportsColumnarOutput` on exactly those types to avoid this. Reading a cached string column threw `scala.MatchError: PhysicalStringType`. The cached format is now decided by schema alone rather than by a runtime config, so a relation is either entirely Comet format or entirely delegated to Spark, and the hand-rolled `DefaultCachedBatch` decoder is gone. `spark.sql.cache.serializer` is a static conf, so a format that could flip mid-session was never readable back reliably. `CometExecRule` checks the schema before choosing the native scan. Also add coverage for all supported types, the row read path over `CometCachedBatch`, and a reordered full-width projection.
|
@0lai0 @manuzhang @sandugood @peterxcli could you help with reviews on this PR? |
|
lgtm👍 |
|
Tried out on a pipeline that relies heavily on caching and got an error: Note that I've tried to enable the suggested config insertions, but it didn't work either. cc @andygrove |
Thanks @sandugood. I'm looking into it now. |
| // Comet's Arrow writer only handles the types listed in supportsSchema. Reporting false here | ||
| // sends the relation down the row path, where it is delegated to Spark's default serializer, | ||
| // instead of failing at cache materialization inside Utils.serializeBatches. | ||
| override def supportsColumnarInput(schema: Seq[Attribute]): Boolean = supportsSchema(schema) |
There was a problem hiding this comment.
supportsColumnarInput promises support for any columnar producer whose schema is supported, but convertColumnarBatchToCachedBatch eventually calls Utils.getBatchFieldVectors, which only accepts CometVector. Spark can expose vectorized Parquet/ORC plans here, causing otherwise supported caches to receive OnHeapColumnVector/OffHeapColumnVector and throw. Please either return false here and use the already-working row-to-Arrow path—the smallest safe fix—or add generic Spark-vector conversion. A vectorized Parquet/ORC cache-materialization regression with Comet execution disabled would cover this.
Assist with LLM
There was a problem hiding this comment.
During initial cache materialization, Spark calls convertColumnarBatchToCachedBatch with the cached plan’s original columnar output. If that plan is a Spark vectorized Parquet/ORC scan rather than a Comet scan, its vectors are not CometVector, so Utils.getBatchFieldVectors throws before any CometCachedBatch is created. Please either force the row-input path or convert generic Spark column vectors to Arrow, with a vectorized Parquet cache-fill regression.
Assist with LLM
There was a problem hiding this comment.
probably the cause of #5051 (comment)
looks like its directly related
There was a problem hiding this comment.
Confirmed, and thanks for catching this — your diagnosis is exactly right, including that it is the cause of the report in #5051 (comment).
I reproduced it: with spark.comet.scan.enabled=false and spark.comet.sparkToColumnar.enabled=false, caching a Parquet table throws the same exception from CachedRDDBuilder$$anon$2.next(InMemoryRelation.scala:277). The mechanism is slightly worse than "Spark can expose vectorized plans here": because supportsColumnarInput returns true, InMemoryRelation.apply actively calls convertToColumnarIfPossible and strips the ColumnarToRow above the cached plan (InMemoryRelation.scala:324-328), so any columnar leaf underneath becomes the serializer input.
Of your two options I took the second (generic Spark-vector conversion) rather than returning false. Returning false is safe but gives up the columnar fast path for Comet-native cached plans, which is most of the point of the feature — the row path would force columnar -> row -> Arrow. Instead encodeBatches now checks Utils.isArrowBacked per batch and copies only foreign vectors into Arrow.
The regression you asked for is in 9dce0d6, and there is a second one covering nullable array/struct/map/binary through the nested vectorized reader. I also documented the contract at supportsColumnarInput itself so the conversion does not read as deletable defensive code.
`ArrowCachedBatchSerializer.supportsColumnarInput` decides the cache input format from the schema alone, which is all Spark gives it. Returning true makes `InMemoryRelation` strip the `ColumnarToRow` above the cached plan and feed `cachedPlan.executeColumnar()` directly to `convertColumnarBatchToCachedBatch`. Any Spark or third-party columnar leaf under that transition (Spark's vectorized Parquet/ORC reader, a connector's own vectors) therefore reaches the serializer with `On/OffHeapColumnVector`-style columns, and `Utils.serializeBatches` fails with "Comet execution only takes Arrow Arrays". The serializer cannot detect this from the schema, so copy non-Arrow batches into Arrow at write time. Comet-native cached plans keep the existing zero-copy path.
| * it. Values are copied element-wise, since Spark's `ColumnVector` implementations do not | ||
| * expose Arrow buffers. | ||
| */ | ||
| def columnarBatchToArrowBatch( |
There was a problem hiding this comment.
I think we already have one?
There was a problem hiding this comment.
checking which type of spark column vector does the SparkColumnarArrowReader support...
There was a problem hiding this comment.
A ColumnarBatch is a container:
class ColumnarBatch {
ColumnVector[] columns;
}Spark 4.1.2’s relevant ColumnVector implementations are:
| Producer | Vectors inside ColumnarBatch |
|---|---|
| Vectorized Parquet | OnHeapColumnVector, OffHeapColumnVector, ConstantColumnVector |
| Vectorized ORC | OrcAtomicColumnVector, OrcArrayColumnVector, OrcMapColumnVector, OrcStructColumnVector; sometimes heap/off-heap or constant vectors |
| Spark Arrow/Python paths | ArrowColumnVector |
| Comet | CometVector |
| External columnar sources | Any implementation extending Spark’s ColumnVector |
SparkColumnarArrowReader does not inspect those concrete classes. It wraps each column in ColumnarArray, and ArrowWriter reads through Spark’s standard accessors such as getLong, getUTF8String, getArray, and getStruct. Therefore, it supports any implementation that correctly follows the spark's ColumnVector
ColumnVector contract for the declared data type. The cache writer could reuse it approximately like this:
override def convertColumnarBatchToCachedBatch(
input: RDD[ColumnarBatch],
schema: Seq[Attribute],
storageLevel: StorageLevel,
conf: SQLConf): RDD[CachedBatch] = {
val sparkSchema = toStructType(schema)
val batchSize = conf.columnBatchSize
val timeZoneId = conf.sessionLocalTimeZone
input.mapPartitions { batches =>;
// Arrow Schema is not serializable, so construct it inside the task.
val arrowSchema = Utils.toArrowSchema(sparkSchema, timeZoneId)
val arrowBatches = CometArrowStream.readerBatchIter(
"CometCacheWrite",
allocator =>;
new SparkColumnarArrowReader(
allocator,
arrowSchema,
batches,
batchSize))
encodeBatches(arrowBatches, schema)
}
}That gives this path:
Spark ColumnarBatch
-> SparkColumnarArrowReader
-> Arrow-backed ColumnarBatch containing CometVector
-> existing encodeBatches/getBatchFieldVectors
-> CometCachedBatch</code>
Assist with LLM
There was a problem hiding this comment.
You are right on both counts, and your analysis of the ColumnVector contract matches what I found — ColumnarArray + the standard accessors means it works for any correctly-implemented ColumnVector, which is why the fix also covers Iceberg and other external columnar sources without naming them.
The duplication is fixed in c7bc49e: the copy loop is now CometArrowConverters.writeColumns(root, batch, startRow, numRows), and SparkColumnarArrowReader.loadNextBatch calls it too. That also collapses the two copies of the zero-column setRowCount workaround back into one.
On your readerBatchIter sketch — I went a slightly different way and want to flag why, since it is a real tradeoff:
- It converts unconditionally, so a Comet-native cached plan (the primary case for this feature) would get a full element-wise copy through the generic accessors instead of the current zero-copy Arrow IPC write. That is the regression I most wanted to avoid, so
encodeBatcheskeeps a per-batchUtils.isArrowBackedcheck and only converts foreign vectors. - The reader wraps the whole partition iterator, so the Arrow-vs-not decision becomes per-partition rather than per-batch. In practice a partition is homogeneous, but making it per-batch costs nothing.
Your version does have one genuine advantage mine gives up: SparkColumnarArrowReader reuses one stable VSR across batches, whereas columnarBatchToArrowBatch allocates a fresh root per batch. If that shows up in cache-materialization profiles it would be worth revisiting — happy to file a follow-up.
One thing your sketch surfaced that I did not change: your code uses conf.sessionLocalTimeZone for the Arrow schema, which matches the pre-existing row path (convertInternalRowToCachedBatch), while my columnar path uses CometArrowStream.NATIVE_TIMEZONE (UTC) to match CometSparkToColumnarExec. So the same CometCachedBatch format is currently written with two timezone conventions depending on the input path. Timestamps are UTC micros either way so this is metadata-only today, but it is inconsistent and I would rather fix it deliberately. Do you have a preference for which one the cache format should standardise on?
There was a problem hiding this comment.
Do you have a preference for which one the cache format should standardise on?
Good catch. I checked Spark’s cache implementations. Released Spark’s default cache does not store timezone metadata: both TimestampType and TimestampNTZType are stored as raw longs, with LTZ versus NTZ preserved by the relation schema.
Current Spark's Arrow cache constructs TimestampType using conf.sessionLocalTimeZone and TimestampNTZType with timezone=null, but its cached RecordBatch is deliberately schema-less. The timezone label is therefore not persisted and is reconstructed on read.
CometCachedBatch differs because it stores a complete IPC stream including the schema. I would therefore standardize its physical format on CometArrowStream.NATIVE_TIMEZONE (UTC) for TimestampType, while keeping TimestampNTZType timezone-free. This avoids embedding the write session’s mutable timezone, matches Comet’s native schema, and avoids an unnecessary native-boundary cast. The row path can simply pass NATIVE_TIMEZONE; no timestamp values are converted.
There was a problem hiding this comment.
Thanks for digging into Spark's implementations — that answers it, and I have taken your recommendation in c25c7b3.
Your framing of the difference is the key point and I had missed it: Spark's Arrow cache can construct TimestampType from conf.sessionLocalTimeZone safely because its RecordBatch is schema-less, so the label is never persisted and gets reconstructed on read. CometCachedBatch stores a full IPC stream including the schema, so for us that label really is written down — which makes using a mutable session value materially worse than it looks by analogy with Spark.
So the row path now passes CometArrowStream.NATIVE_TIMEZONE, matching the columnar path, which already did. Confirmed the two things you predicted:
- No values move. Spark's internal representation is micros since the Unix epoch regardless of session timezone, so this is a label-only change. The test checks both the values and their
CAST(ts AS STRING)rendering against Spark under two non-UTC session timezones. TimestampNTZTypeneeds nothing.Utils.toArrowTypemaps it toTimestamp(MICROSECOND, null)whatever timezone is passed, so it was already timezone-free.
A small bonus: the closure no longer needs the session timezone, so it no longer captures anything derived from conf at all. Worth noting because the original had to hoist val sessionTz outside mapPartitions — when I moved the read inside while testing, it NPEd in ConfigEntry.readString on the executor. That hazard is now gone.
The test caches a row-based plan (local Seq, so convertInternalRowToCachedBatch rather than the columnar path) under America/Los_Angeles and Asia/Kolkata, decodes the cached batches back through the serializer, and asserts the Arrow field's timezone. I checked it is not vacuous: reverting the one-line change fails it with got [America/Los_Angeles].
There was a problem hiding this comment.
Agreed, and that is what I ended up doing — SparkColumnarArrowReader.loadNextBatch and columnarBatchToArrowBatch now share CometArrowConverters.writeColumns (c7bc49e), so there is one copy loop rather than two. Marking this thread as settled; the remaining timezone question from it is answered on the sibling thread.
There was a problem hiding this comment.
Following up here because merging latest apache/main changed the outcome of this thread.
The shared copy loop we settled on, CometArrowConverters.writeColumns, is gone. #5046 and #5442 landed on main in the meantime and moved that responsibility onto ArrowWriter itself as ArrowWriter.writeColumns(input, startRow, numRows), with a bulk-copy path for fixed-width columns rather than the element-wise loop this PR had. So the deduplication you asked for now lives upstream of both callers, and I resolved the conflict by deleting this PR's version rather than keeping a second one: SparkColumnarArrowReader.loadNextBatch takes main's version wholesale, and columnarBatchToArrowBatch now calls ArrowWriter.create(root, numRows), writeColumns, finish().
One detail worth recording, since it was a deliberate deletion rather than an oversight. This PR's writeColumns ended with an explicit root.setRowCount(numRows), because ArrowWriter derived the root row count from per-column writes and a zero-column batch (Spark's count-from-metadata scan, numRows > 0 with numCols == 0) would otherwise have produced rowCount == 0 and silently dropped the rows. Main's writeColumns sets count = numRows directly, so finish() gets it right with no columns present and the workaround is no longer needed. CometInMemoryCacheSuite has a test for that exact case ("supports empty projection scan"), and it still passes, so this is covered rather than taken on trust.
The leak guard from the other thread is unaffected and still wraps both call sites. All 21 tests in the suite pass on the merged branch.
|
I pushed a potential fix at the same time @peterxcli was reviewing. |
Extract the element-wise Spark `ColumnVector` -> Arrow copy loop that `columnarBatchToArrowBatch` had duplicated from `SparkColumnarArrowReader` into a shared `CometArrowConverters.writeColumns`, so the zero-column rowCount workaround lives in one place. Also: - hoist the Arrow schema out of the per-batch path - replace the `(batch, ownsBatch)` tuple and `.toList` with a `serializeBatch` helper that takes the single element eagerly - move `isArrowBacked` to `Utils`, next to the `getBatchFieldVectors` precondition it describes - drop `toStructType` in favour of the existing `Utils.fromAttributes` - document at `supportsColumnarInput` why its schema-only answer means the conversion in `encodeBatches` is load bearing - route the new tests through the existing `withNativeCache` helper and assert both of them really cached in Comet's format
|
Thanks for the report @sandugood, and sorry for the bad error message — it sent you chasing scan configs when the failure was actually in the new cache serializer. @peterxcli diagnosed it correctly in the review thread above. What happened: the Comet cache serializer decided it could accept columnar input based on the schema alone. That makes Spark strip the Workaround on current main: Fixed in 9dce0d6 + c7bc49e on this branch: non-Arrow columnar batches are now copied into Arrow at cache write time, while Comet-native cached plans keep the existing zero-copy path. Two regression tests cover it — primitives/string/decimal/date/timestamp, and nullable array/struct/map/binary. Since Iceberg MOR is the one input I could not reproduce locally, it would be really helpful if you could retest your pipeline against the branch. Iceberg's vectors go through the same generic |
| * (e.g. Comet's cache serializer, which Spark hands the cached plan's columnar output) use this | ||
| * to convert foreign vectors to Arrow instead of tripping the exception below. | ||
| */ | ||
| def isArrowBacked(batch: ColumnarBatch): Boolean = |
There was a problem hiding this comment.
Is this possible to be false?
There was a problem hiding this comment.
Yes — and this is the case the PR exists to handle.
isArrowBacked is false whenever Comet is handed a columnar batch from a plan it did not build. The concrete reproduction is the one from the report on this PR: with spark.comet.scan.enabled=false, Spark's own vectorized Parquet reader produces OnHeapColumnVectors, supportsColumnarInput returns true so InMemoryRelation.apply strips the ColumnarToRow above that scan, and the cache serializer receives non-Comet vectors. Before the converter path those batches hit the exception in getBatchFieldVectors.
It is also false for any external columnar source — an Iceberg or other connector's ColumnVector implementation — which is why the conversion is written against the ColumnVector contract rather than against specific classes.
The two tests at the end of CometInMemoryCacheSuite ("cache a Spark columnar plan whose vectors are not Arrow-backed" and the complex-types variant) cover exactly this, so the false branch is exercised.
| * it. Values are copied element-wise, since Spark's `ColumnVector` implementations do not | ||
| * expose Arrow buffers. | ||
| */ | ||
| def columnarBatchToArrowBatch( |
There was a problem hiding this comment.
Do you have a preference for which one the cache format should standardise on?
Good catch. I checked Spark’s cache implementations. Released Spark’s default cache does not store timezone metadata: both TimestampType and TimestampNTZType are stored as raw longs, with LTZ versus NTZ preserved by the relation schema.
Current Spark's Arrow cache constructs TimestampType using conf.sessionLocalTimeZone and TimestampNTZType with timezone=null, but its cached RecordBatch is deliberately schema-less. The timezone label is therefore not persisted and is reconstructed on read.
CometCachedBatch differs because it stores a complete IPC stream including the schema. I would therefore standardize its physical format on CometArrowStream.NATIVE_TIMEZONE (UTC) for TimestampType, while keeping TimestampNTZType timezone-free. This avoids embedding the write session’s mutable timezone, matches Comet’s native schema, and avoids an unnecessary native-boundary cast. The row path can simply pass NATIVE_TIMEZONE; no timestamp values are converted.
|
When the PR is sound I could test it again. Is it ready to be tested @andygrove? Last time I've built it, got an error while writing to an iceberg table: (note that I obfuscated column names, but kept their order) |
peterxcli
left a comment
peterxcli
left a comment
There was a problem hiding this comment.
@andygrove Thanks for the update on vector converting!
I have four new comments on this, but they all can be address in followup.
Other LGTM!
| // Apply Spark's cache batch filter before decoding. Spark's InMemoryTableScanExec does this in | ||
| // filteredCachedBatches(), but that method is private. Reusing the serializer's buildFilter here | ||
| // keeps Comet on the same stats-based pruning path instead of decoding every cached batch. | ||
| override def doExecuteColumnar(): RDD[ColumnarBatch] = { |
There was a problem hiding this comment.
pruning configuration ignored: still prunes when spark.sql.inMemoryColumnarStorage.partitionPruning=false.
not sure enable partitionPruning is always better than disable? if true, then nvm.
There was a problem hiding this comment.
Real bug, fixed in 1367665 — thanks. Spark's InMemoryTableScanExec.filteredCachedBatches gates on conf.inMemoryPartitionPruning and Comet was applying the filter unconditionally, so the config was silently ignored.
To your "not sure enable partitionPruning is always better than disable": pruning is normally the win, but that is not really the point — the config exists so a user can turn it off, typically to rule out a stats bug when results look wrong. Ignoring it means the one knob they are reaching for does nothing, and Comet diverges from Spark precisely when someone is debugging. So honoring it is right regardless of which default is faster.
Covered by a new test that observes the scan's numOutputRows, which counts rows in the batches actually decoded: 100 with pruning on (a single 100-row batch) and 1000 with it off (every batch decoded). One wrinkle worth noting for anyone writing a similar test: the metric has to be read after forcing that exact df to run, because checkSparkAnswer executes its own copies and leaves the collected plan instance at zero — my first version passed vacuously with 0 on both sides.
| * Reads of `CometCachedBatch` keep working when the native scan is disabled, because Spark then | ||
| * reads the same cached data through the SparkToColumnar fallback path. | ||
| */ | ||
| class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { |
There was a problem hiding this comment.
add early termination handling: eg. LIMIT, take, or any other cancellation
ref: spark's TaskCompletionListener of ArrowCachedBatchSerializer
There was a problem hiding this comment.
Real leak, fixed in 1367665. Confirmed the mechanism: ArrowReaderIterator.close() is only reached from hasNext when the stream is exhausted, so LIMIT, take(), or a cancelled task leaves the reader it was part-way through open.
convertCachedBatchToColumnarBatch now registers a TaskCompletionListener, following the Spark implementation you linked. One simplification: flatMap consumes each inner iterator fully before constructing the next, so at most one reader is open at a time and tracking the current one is enough rather than a collection. close() is already idempotent and synchronized, so closing one that exhausted itself is a no-op and the listener is safe to run from another thread.
| * `relationOutput` is the full schema stored in the cache. `scanOutput` is the subset requested | ||
| * by this scan after pruning. | ||
| */ | ||
| case class CometInMemoryTableScanExec( |
There was a problem hiding this comment.
Add AQE test:
- SPARK-42101: leaves a cached join cold, first-touches it through an AQE aggregation, and checks cold/warm materialization and plan rewrites.
- Table-cache stage in an AQE join: verifies TableCacheQueryStageExec and shuffle behavior.
- SPARK-37742: verifies AQE does not choose joins using invalid cache runtime statistics.
There was a problem hiding this comment.
Not done, and I would rather flag that than half-do it. Porting the three AQE scenarios you linked (SPARK-42101 cold/warm materialization, the TableCacheQueryStageExec join, SPARK-37742) is a meaningful chunk of work and touches AQE plan-shape assertions rather than the cache format this PR changes, so I have left it out of this round rather than rushing it alongside the four correctness fixes.
You said your comments could be follow-ups, so unless you would rather block on it, my suggestion is a tracking issue for AQE coverage of CometInMemoryTableScanExec. I have not filed one yet — say the word and I will, or feel free to open it yourself if you have a preference for how it should be scoped.
There was a problem hiding this comment.
Filed as #5245, scoped to the three scenarios you linked, run with spark.comet.exec.inMemoryCache.enabled=true and asserting the cached scan is CometInMemoryTableScanExec.
I called out SPARK-37742 as the one to weight most: the Comet serializer reports its own sizes for CometCachedBatch, so the runtime statistics AQE reads when picking a join strategy come from Comet code rather than Spark's DefaultCachedBatch accounting. That is the scenario where a divergence would be silent and would change plans rather than results.
| override val numRows: Int, | ||
| override val sizeInBytes: Long, | ||
| override val stats: InternalRow, | ||
| bytes: ChunkedByteBuffer) |
There was a problem hiding this comment.
ChunkedByteBuffer is Externalizable, so BlockManager can serialize it to DiskStore. so both storage level (DISK_ONLY and MEMORY_AND_DISK) should already supported? Could we add a StorageLevel.DISK_ONLY regression test?
A deterministic DISK_ONLY test should:
- Materialize a
CometCachedBatchusingStorageLevel.DISK_ONLY. - Assert
memSize == 0,diskSize > 0, and all partitions cached. - Run a second query and verify the answer plus
CometInMemoryTableScan, proving disk deserialization and decoding work.
Spark uses the same size assertions in its BlockManager regression.
There was a problem hiding this comment.
Added in 1367665, following your spec: "supports DISK_ONLY storage level" materializes the cache at StorageLevel.DISK_ONLY and asserts memSize == 0, diskSize > 0, numCachedPartitions == numPartitions, that the payload is still CometCachedBatch, and that a second query returns the right answer through CometInMemoryTableScan.
Your reasoning was right — ChunkedByteBuffer being Externalizable means BlockManager spills it like any other block, and it worked first time with no production change needed. Worth having the regression test precisely because nothing in Comet's code makes that true, so it could break without notice.
| batch: ColumnarBatch, | ||
| arrowSchema: Schema, | ||
| allocator: BufferAllocator): ColumnarBatch = { | ||
| val root = VectorSchemaRoot.create(arrowSchema, allocator) |
There was a problem hiding this comment.
Minor Nit: columnarBatchToArrowBatch allocates root before NativeUtil.rootAsBatch wraps it.
If writeColumns or the wrap throws, root leaks, since the caller only closes the returned batch.
Consider guarding the body with `try { ... } catch { case NonFatal(e) => root.close(); throw e }.
But normal path will not lead to
There was a problem hiding this comment.
Good catch, fixed in 1367665. Guarded with try/catch NonFatal releasing the root before rethrowing, as you suggested.
You are right that the normal path cannot hit it, but writeColumns walks arbitrary ColumnVector implementations, so a badly-behaved external vector is exactly the case that would throw here — and that is the path this PR added.
rowToArrowBatchIter just above has the same shape (root allocated, then writer.write(rowIter.next()) before rootAsBatch takes ownership), so I gave it the same guard rather than leaving one of the two fixed.
Three review fixes on the cache path, plus tests for two of them. Honor spark.sql.inMemoryColumnarStorage.partitionPruning. CometInMemoryTableScanExec applied the serializer's stats filter unconditionally, so setting the config to false still pruned. Spark's InMemoryTableScanExec.filteredCachedBatches gates on it, and a user reaching for that knob -- to rule out a stats bug, say -- is specifically trying to stop pruning, so silently ignoring it makes Comet diverge on the one thing they are controlling. Close Arrow readers when a consumer stops early. ArrowReaderIterator closes its reader only on exhaustion, so LIMIT, take() or a cancelled task left the reader it was part-way through open. convertCachedBatchToColumnarBatch now registers a TaskCompletionListener, as Spark's own ArrowCachedBatchSerializer does. flatMap consumes each inner iterator fully before building the next, so at most one reader is open at a time and tracking the current one suffices; close() is already idempotent and synchronized, so closing an exhausted one is a no-op. Release the VectorSchemaRoot if conversion throws. columnarBatchToArrowBatch allocated a root before NativeUtil.rootAsBatch wrapped it, and the caller only owns the returned batch, so a throw from writeColumns leaked the allocation. rowToArrowBatchIter had the same shape and gets the same guard. Tests: - "honors inMemoryColumnarStorage.partitionPruning=false" observes the scan's numOutputRows, which counts rows in the batches actually decoded: 100 with pruning on (one batch), 1000 with it off. Note the metric has to be read after forcing this df to run, since checkSparkAnswer executes its own copies and leaves the collected plan instance at zero. - "supports DISK_ONLY storage level" asserts memSize == 0, diskSize > 0, all partitions cached, the payload is still CometCachedBatch, and the cache reads back through CometInMemoryTableScan. CometCachedBatch holds a ChunkedByteBuffer, which is Externalizable, so BlockManager can spill it like any other block. All 19 tests in CometInMemoryCacheSuite pass on spark-3.5.
|
@peterxcli @0lai0 thanks both — five of the six comments are addressed in 1367665, and two were real bugs. Pruning config was silently ignored (@peterxcli). Readers leaked on early termination (@peterxcli). Confirmed: Root leaked if conversion threw (@0lai0). Fixed with the DISK_ONLY test added to your spec:
AQE tests: not done. Porting the three scenarios you linked is a real chunk of work on AQE plan-shape assertions rather than the cache format this PR changes, so I left it rather than rushing it. Since you said these could be follow-ups, I suggest a tracking issue for AQE coverage of One thing worth passing on from writing the pruning test: reading a metric off All 19 tests in |
The row write path passed conf.sessionLocalTimeZone into rowToArrowBatchIter
while the columnar path already encoded with CometArrowStream.NATIVE_TIMEZONE, so
the same logical cache had two physical formats depending on which path filled
it, and the row one persisted the writing session's timezone into cached data.
Unlike Spark's Arrow cache, whose RecordBatch is deliberately schema-less and
reconstructs the timezone on read, CometCachedBatch stores a full IPC stream
including the schema, so that label really is written down. Standardise on
NATIVE_TIMEZONE ("UTC") for TimestampType, per the analysis on the review thread.
This is a label only. Spark's internal timestamp representation is micros since
the Unix epoch regardless of session timezone, so no values are converted, and
matching Comet's native schema also avoids a cast at the native boundary.
TimestampNTZType already had no timezone: Utils.toArrowType maps it to
Timestamp(MICROSECOND, null) whatever is passed in.
The closure no longer needs the session timezone at all, so it no longer captures
anything derived from `conf`.
Test: "stores timestamps with a UTC schema label" caches a row-based (local Seq)
plan under two non-UTC session timezones, decodes the cached batches back through
the serializer, and asserts the Arrow field carries "UTC" -- then checks the
values and their string rendering still match Spark. Reverting the one-line change
fails it with `got [America/Los_Angeles]`, so it pins the format rather than
passing vacuously.
All 20 tests in CometInMemoryCacheSuite pass on spark-3.5.
|
@peterxcli the timezone question is settled in c25c7b3 — thanks for checking Spark's implementations, your recommendation was the right call and for a reason I had missed. The distinction that matters: Spark's Arrow cache can build So the row write path now passes
Small bonus: the closure no longer needs the session timezone, so it captures nothing derived from The new test caches a row-based plan (local All 20 tests in That clears every inline thread on this PR except the AQE test coverage, which we agreed can be a follow-up. |
…hat they project A CometCachedBatch held one compressed Arrow IPC stream covering every cached column, so convertCachedBatchToColumnarBatch inflated all of them and projected afterwards. Read cost was flat in the width of the projection, where Spark's per-column DefaultCachedBatch falls away as it narrows, and a narrow read of a wide cached relation was several times slower than Spark's cache despite materializing faster. Store one stream per column and decode only the selected ones. An empty selection now stays empty rather than expanding to every column, and the scan asks for a single cheap column instead of the whole schema when a query needs only the row count, since the native plan still requires a non-empty scan schema. Per-column sizes are now known, so the statistics field Spark reserves for them holds the real value. Measured on 5M rows and 6 columns, against Spark's cache format: read shape before after Spark count(*) 241 ms 77 ms 62 ms 1 of 6 59 ms 57 ms 105 ms 3 of 6 205 ms 204 ms 339 ms 6 of 6 448 ms 448 ms 375 ms Framing costs a schema block and compression framing per column per batch, and loses cross-column compression: footprint grows 2.5% at 6 columns and 32% at 60, where the cached relation is still 22x smaller than Spark's format.
|
Fixed #5484 here rather than leaving it as a follow-up. A cache format that loses to Spark's on the reads people actually do is not worth shipping, even off by default, so it seemed better to fix it than to document it. What changed. A Per-column sizes are also known now, so the fifth statistics field Spark reserves for them holds the real value instead of zero. Results, 5M rows and 6 columns, min of 5 runs after a discarded warm-up, against a real
Narrow and medium projections now beat Spark's cache by 1.7x to 1.8x rather than losing to it. The cost, since I said I would measure it rather than assume it. Per-column framing adds a schema block and compression framing per column per batch, and gives up cross-column compression. Footprint grows 2.5% at 6 columns and 32% at 60 columns. At 60 columns the cached relation is still 107 MiB against Spark's 2369 MiB, so I took the tradeoff, but it is a real one and it scales with width, so it is worth knowing about. What is still slower. Reads that feed Spark operators rather than Comet ones remain 1.7x to 2.5x behind Spark's cache. Those pay a row conversion that Spark's format avoids with a codegen'd path over its own layout. That is a different cost from the projection problem and I have not touched it here. @sandugood this should help your pipeline on the read side, but that fallback gap is exactly the shape you were running, so I would not promise it closes the whole 10-to-15-minute difference. If you do retest, the two things worth separating are cache materialization and the repeated reads, since they now move in opposite directions. Testing. Five new tests in @peterxcli this changes the cached payload and the read path, so your approval no longer covers it. Sorry to move the target after you had signed off. The lifetime handling around the readers is the part I would most want a second pair of eyes on: a cached batch is now several independent Arrow streams, so the |
scalafix RemoveUnused rejects the named relation parameter, which this test does not use. Spotless does not catch it, so it only surfaced in the Lint Java jobs.
|
Will take a look! It looks similar to apache/spark#56334 too although seems in Comet we store compressed IPC stream and can decode only selected columns. cc @viirya too |
Thanks @sunchao, that is a useful pointer and I had not seen it. I went through it and filed #5487 with the comparison so it does not get lost. A few things worth saying here. On the version. It is in On the difference you spotted, which is real but slightly the other way round from how it sounds. Both formats decode only the selected columns. Comet gets there by splitting the payload into one compressed stream per column, which is the change I pushed a few commits ago. Spark keeps a single RecordBatch and instead parses the IPC message flatbuffer, which lists every buffer's offset and length within the body, copying out only the byte ranges belonging to the selected columns and letting Spark's is the better design point. Per-column streams pay framing per column per batch and give up cross-column compression: measured against the previous single-stream format, footprint grows 2.5% at 6 columns and 32% at 60. Spark's approach gets the same projection-proportional decode with none of that. Part of the gap is simply that their Two smaller things from their implementation, both in #5487: they register One point of convergence worth noting, since it suggests neither of us is off in the weeds: they landed the same empty-projection optimization I did, emitting row counts without touching the payload so @viirya if you have a moment, the question I would most value your view on is whether the buffer-span projection reader was worth its complexity in practice, or whether you would take the simpler split-by-column route knowing what you know now. I also checked our collation handling against yours and there is no bug on our side, but you do prune collated string columns where we decline to: a collated |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Reviewed commit 86a91ea7dd0fbc9765ea74340bc460c5e74ccfc4. Five inline findings are attached: one incorrect-result issue and four additional correctness/resource/execution issues.
Verification: Spark 4.1.3 with JDK 17, freshly compiled JVM classes, and a verified native build from the identical native source tree. All 58 existing tests passed across CometInMemoryCacheSuite, UtilsSuite, NativeUtilSuite, and CometNativeColumnarToRowSuite; separate targeted probes reproduced the reported failures. The large-offset Arrow finding was reproduced at the serializer boundary; the Python producer path was source-traced only. Other Spark versions were not tested.
Review assisted by Codex.
| } else { | ||
| Seq(op.relation.output.minBy(a => decodeCostRank(a.dataType))) |
There was a problem hiding this comment.
[P1] Preserve zero-column output for joins
An empty-output cache scan can feed a join, not only a count-style aggregate. Adding this column changes the native join's column positions while output still declares the original empty schema. With native caching enabled and AQE disabled, I reproduced:
val left = spark.range(10L, 13L).cache()
left.collect()
left.createOrReplaceTempView("cached_left")
spark.sql("""
SELECT /*+ BROADCAST(r) */ sum(r.id)
FROM cached_left l JOIN range(2) r ON true
""").collect()The result is 66, versus the correct 3 when native cache scanning is disabled. The native join receives the hidden left column before r.id, so the aggregate's bound ordinal 0 reads the wrong values. Selecting r.id directly also fails with Output column count mismatch: expected 1, got 2. Please preserve the zero-column scan schema/batches instead of introducing an unreported column. A genuinely zero-column cached relation already returns the correct count on this head.
There was a problem hiding this comment.
Confirmed and fixed in cb07053. I reproduced your case exactly: 66 with native caching on, 3 with it off, and the direct r.id projection failing with the column count mismatch.
The premise the code rested on turned out to be false. The comment claimed the widening was needed because "the native plan still needs a non-empty scan schema", and I carried that forward when I narrowed the fallback from the whole cache schema to a single cheap column. It isn't true: passing op.output straight through, empty and all, plans and executes fine. So the fix is to delete the fallback rather than to make the extra column agree with output.
Worth noting for the record that this predates my change. The previous code fell back to relation.output, and since spark.range has exactly one column, your repro produced an identical one-column scan before and after. My change altered how wide the leak was on wider relations, not whether it existed.
Two things improve as a result. The join is correct, and count(*) now decodes nothing at all rather than one placeholder column, which is strictly better than what I had: selectedIndices keeps an empty selection empty, so the serializer short-circuits on row count alone.
Regression test added as "joins correctly over an empty-output cache scan", covering both the aggregate and the direct projection, and the existing count-only test now asserts scanOutput.isEmpty rather than a placeholder column.
| val root = new VectorSchemaRoot(Seq(fieldVector).asJava) | ||
| val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out)) | ||
| writer.start() | ||
| writer.writeBatch() |
There was a problem hiding this comment.
[P2] Use each column's dictionary provider when serializing it
The new decoder opens one Arrow reader per column, so its dictionary-backed output columns have independent providers. getBatchFieldVectors returns only the first dictionary column's provider, which is then passed to every writer here. Re-encoding a decoded cache batch cannot resolve the later columns' dictionary IDs.
I reproduced this through normal Spark operations: with Comet caching and spark.comet.shuffle.mode=jvm, cache a repartitioned DataFrame with two low-cardinality string columns. Then disable spark.comet.exec.enabled and spark.comet.exec.inMemoryCache.enabled and run first.union(first).cache().count(). Spark's columnar Union passes the decoded cached batches back to this serializer, and the second cache fails with IllegalArgumentException: Could not find dictionary with ID 1. The original cache reads correctly and its columns have separate provider namespaces. Please obtain the provider associated with each column rather than reusing the first one.
There was a problem hiding this comment.
Correct, and this one is squarely mine: it became reachable when I split the payload into one stream per column. Before that every column came from a single reader and therefore shared a provider, so collapsing to the first one was harmless. With independent readers the ID namespaces collide.
Fixed in cb07053. Utils gains getBatchFieldVectorsWithProviders, which pairs each field vector with the provider its own column was decoded with, and serializeBatchColumns uses that per column. getBatchFieldVectors keeps its existing signature and collapses to the first provider, so the shuffle and broadcast callers are unchanged.
Regression test added, using your reproduction shape: cache a repartitioned frame with two low-cardinality string columns, then with spark.comet.exec.enabled=false and the native cache scan disabled, first.union(first).cache().count(), so Spark's columnar Union hands the decoded batches back to the serializer. I verified it fails without the fix, with the Could not find dictionary with ID 1 you described.
| private val readers: Array[Iterator[ColumnarBatch]] = | ||
| buffers.map(Utils.decodeBatches(_, "CometCache")) |
There was a problem hiding this comment.
[P2] Close previously opened readers when initialization fails
decodeBatches opens and decodes a column eagerly. If opening a later column throws, the readers already created by this map are lost. The task-completion listener cannot release them because the new ColumnReaders instance is assigned to current only after this constructor returns.
I reproduced this with valid cached data containing two integer columns and an allocator limit that permits the first column but rejects the second allocation. After task completion, the first column's 512 bytes remain allocated; three failed tasks leave 512, 1024, and 1536 bytes respectively, starting from zero. This turns allocation failures into persistent off-heap leaks and worsens subsequent memory pressure. Please close every successfully opened reader if constructing the remaining readers fails.
There was a problem hiding this comment.
Confirmed and fixed in cb07053. Your reading of the mechanism is exactly right: current is assigned only after the constructor returns, so the listener has nothing to close, and decodeBatches allocates eagerly because ArrowReaderIterator decodes its first batch in its constructor.
The readers are now opened in a loop guarded by try/NonFatal that closes everything already opened before rethrowing, with a failing close attached via addSuppressed rather than replacing the original failure, matching the guard used in CometArrowConverters.
I tested it without needing an allocator limit, since CometArrowAllocator exposes its accounting: corrupt the second of two selected column streams, record getAllocatedMemory, attempt the read, and assert the allocation returns to the baseline. Without the fix that leaks 2048 bytes per attempt, which matches the growth you saw across repeated tasks.
| relation.cacheBuilder.serializer, | ||
| relation.cacheBuilder.cachedColumnBuffers, |
There was a problem hiding this comment.
[P2] Defer cache RDD construction until execution
Accessing cachedColumnBuffers here is not a passive metadata lookup: Spark's cache builder constructs it by calling cachedPlan.execute/executeColumnar. If that plan is adaptive, this can execute shuffle stages while Comet is still planning the outer query.
With AQE and native caching enabled, I reproduced:
val cached = spark.range(100).repartition(2).cache()
cached.createOrReplaceTempView("cached_adaptive")
spark.sql("SELECT * FROM cached_adaptive").explain()explain() starts one Spark job and changes the cached adaptive plan to isFinalPlan=true; with native cache scanning disabled and the same cache serializer, it starts zero jobs and the cached plan stays unexecuted. Planning/EXPLAIN can therefore perform expensive computation or fail on execution errors. Please retain the relation/cache builder and obtain its RDD inside doExecuteColumnar().
There was a problem hiding this comment.
Fixed in cb07053. CometInMemoryTableScanExec now holds the CachedRDDBuilder instead of an RDD[CachedBatch], and resolves cachedColumnBuffers inside doExecuteColumnar().
Regression test added. One note on how it asserts, because my first attempt was wrong in a way worth flagging: isCachedColumnBuffersLoaded is not the right signal, since it also requires the blocks to be populated, so it stayed false either way and the test passed with the fix reverted. The test now asserts on AQE finalization instead, which is the effect you actually described: the cached plan must still report isFinalPlan=false after explain(), and the query must still materialize normally when run. That version fails without the fix.
This one predates my change as well, since the original constructor took the RDD too, but it is a real problem and worth fixing here.
| val columns = if (Utils.isArrowBacked(batch)) { | ||
| Utils.serializeBatchColumns(batch) | ||
| } else { | ||
| val arrowBatch = | ||
| CometArrowConverters.columnarBatchToArrowBatch(batch, arrowSchema, CometArrowAllocator) | ||
| try Utils.serializeBatchColumns(arrowBatch) | ||
| finally arrowBatch.close() |
There was a problem hiding this comment.
[P2] Handle large-offset Arrow vectors before taking the direct write path
Utils.isArrowBacked accepts any CometVector, including a CometPlainVector wrapping LargeVarCharVector or LargeVarBinaryVector. However, serializeBatchColumns calls Utils.getFieldVector, which rejects both representations. Thus supportsColumnarInput accepts the Spark StringType/BinaryType schema, but materializing the cache fails.
At the serializer boundary, I reproduced both large-vector cases with just hello and world: Spark's default cache round-trips them, while this serializer throws Unsupported Arrow Vector for serialize. Ordinary VarCharVector succeeds on both paths. Accelerated mapInArrow is a relevant producer because its existing runner preserves returned Arrow vectors, including pa.large_string()/pa.large_binary(). That Python integration path was source-traced, not executed end to end; the serializer failure and default-cache comparison were run. Please normalize these physical representations through the conversion path or support them in the direct writer.
There was a problem hiding this comment.
Fixed in cb07053. The mismatch is exactly as you describe: isArrowBacked asked only whether every column is a CometVector, while getFieldVector accepts a specific vector list that excludes both large-offset representations, so supportsColumnarInput accepted the schema and materialization then threw.
Utils now has isSupportedFieldVector, which answers the same question getFieldVector does without throwing to find out, and both use it: getFieldVector for its check, isArrowBacked to report false for a CometVector wrapping something it cannot write. Such batches take the conversion path, which rebuilds them against the Arrow schema derived from the Spark schema and so normalizes large offsets to VarChar/VarBinary.
Regression test added in UtilsSuite at the level the fix lives, covering both LargeVarCharVector and LargeVarBinaryVector. I did not attempt the mapInArrow end-to-end path, matching the scope you tested.
This is pre-existing rather than new here, since the previous single-stream writer went through the same getBatchFieldVectors call, but it is cheap to fix and the fix is contained.
There was a problem hiding this comment.
Filed the durable part of this as #5488.
To be precise about what does and does not outlive this PR, since I called it pre-existing above and that was ambiguous. CometInMemoryTableScanExec and ArrowCachedBatchSerializer are both new files in #5051, so the two findings in the scan node cannot survive it being dropped. This one is different: getFieldVector and both Utils.serializeBatches consumers, getByteArrayRdd in operators.scala and CometBroadcastExchangeExec, are on main today.
And the fix here does not close it. Making isArrowBacked reject those vectors only reroutes the cache write path to conversion; the broadcast and collect paths have no conversion fallback and still throw. isSupportedFieldVector is a useful building block for a real fix, but it goes away with this PR too. #5488 has the producer trail through CometMapInBatchExec, and notes that neither of us ran the mapInArrow path end to end.
I don't look at the details of this yet. I can do it later this week. In Spark, the Arrow cache also stores IPC format data and can decode/decompress the selected columns only. |
"the buffer-span projection reader" is what you implemented in this PR? So your question is about the column pruning (i.e., selection of columns) in the cache? I did the column pruning in Arrow cache by this Spark PR apache/spark#58177. |
Five issues from review of 86a91ea. Emit zero-column batches for an empty-output cache scan. The scan widened an empty projection to a placeholder column so the native plan would have a non-empty scan schema, but that claim was untrue: a zero-column scan plans and runs. Meanwhile the emitted batches disagreed with the declared output, so a join over a cached relation read the wrong column and silently returned a wrong sum rather than failing. Removing the widening fixes the join and makes count(*) decode nothing at all. Serialize each column with the dictionary provider it was decoded with. Columns decoded from separate streams have independent dictionary ID namespaces, so re-encoding a decoded batch with the first column's provider could not resolve the rest. Release readers opened before a later column fails to decode. The holder is published to the task-completion listener only after its constructor returns, so a partial failure leaked off-heap for the life of the executor. Build the cached RDD in doExecuteColumnar rather than at planning time. CachedRDDBuilder.cachedColumnBuffers executes the cached plan, so planning a query over an adaptively cached relation ran a job and finalized that plan during EXPLAIN. Report large-offset Arrow vectors as not directly writable. isArrowBacked accepted any CometVector, including one wrapping LargeVarCharVector or LargeVarBinaryVector, which getFieldVector rejects; such batches now take the conversion path instead of failing at cache materialization.
|
@sunchao thanks, that is a genuinely valuable review and all five findings were real. Fixed in cb07053, with a regression test for each. Replies are on the individual threads; the short version:
The P1 fix is better than a repair. The code widened an empty projection to a placeholder column because a comment asserted the native plan needs a non-empty scan schema. That assertion is simply false: passing Two of the five are mine, both consequences of splitting the payload into one stream per column: independent readers mean independent dictionary namespaces, and a partial failure mid-construction has no owner yet. Worth saying plainly, since the split is a change I made deliberately and these are its costs. One process note in case it is useful to you. My first version of the planning regression test asserted on All 30 tests in Separately, I went through the Spark work you linked and filed #5487 with the comparison, including a couple of things worth taking from it that are unrelated to these findings. |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed commit cb07053e30ba084dd63787bf65ecb88a194f414a. All five original reproductions now pass. Three remaining P2 findings are attached: dictionary handling during native broadcast, cleanup when a reader's own initialization fails, and cache-scan canonicalization/exchange reuse.
Verification: Spark 4.1.3 with JDK 17, freshly compiled JVM classes, and a verified native build from the identical native source tree. All 63 tests passed across CometInMemoryCacheSuite, UtilsSuite, NativeUtilSuite, and CometNativeColumnarToRowSuite. Another 222 correctness checks, 20 type/conversion reads, and 64 configuration checks passed. Each inline finding was independently reproduced. Older Spark versions received source inspection only.
Review assisted by Codex.
| val columns = getBatchFieldVectorsWithProviders(batch) | ||
| (columns.map(_._1), columns.flatMap(_._2).headOption) |
There was a problem hiding this comment.
[P2] Preserve each column's dictionary when broadcasting cached batches
The per-column cache write fix works, but this whole-batch helper still discards every provider except the first. A native broadcast can consume CometInMemoryTableScanExec directly after the wrapper is removed, so CometBroadcastExchangeExec.getByteArrayRdd passes a batch with independent providers to Utils.serializeBatches.
I reproduced this with Comet caching, JVM shuffle and AQE disabled: cache a repartitioned DataFrame with two low-cardinality string columns, then run SELECT /*+ BROADCAST(c) */ c.a, c.b FROM range(1) r JOIN dictionary_cache c ON true. It fails with Could not find dictionary with ID 1. Disabling only spark.comet.exec.inMemoryCache.enabled lets the same query over the same cache return all 10,000 rows. Please normalize or safely combine the dictionaries for whole-batch serialization as well as the per-column writer.
There was a problem hiding this comment.
Fixed in 70f046a. getBatchFieldVectors now combines the providers instead of taking the first: each dictionary-encoded column's dictionary is looked up under the provider that column was decoded with, and the result is one MapDictionaryProvider covering the whole batch.
Your repro is the regression test, Comet in-memory cache broadcasts a batch whose columns have separate dictionaries. With the fix reverted it fails with exactly IllegalArgumentException: Could not find dictionary with ID 1.
On normalize versus combine: combine. Renumbering would mean rewriting each vector's Field dictionary ID, which is not settable without copying the vector, so a same-ID/different-dictionary clash now raises rather than resolving one column against another's dictionary. That case is unreachable today, since every provider reaching this path descends either from a single NativeUtil importer or from a stream whose IDs that importer wrote, but silently decoding one column with the wrong dictionary was the worse of the two failure modes to leave available.
| var i = 0 | ||
| try { | ||
| while (i < buffers.length) { | ||
| opened(i) = Utils.decodeBatches(buffers(i), "CometCache") |
There was a problem hiding this comment.
[P2] Close the reader whose own initialization fails
The new guard correctly closes previously returned readers, but decodeBatches can allocate a dictionary and then throw while loading that same column's first record batch. In that case it never returns a reader to opened(i), and neither this catch nor the task-completion listener can close it.
With one valid dictionary-encoded cached string column, I limited the allocator to allow its dictionary message but reject the following indices buffer. After three completed failed tasks, retained Arrow allocations grew from zero to 64, 128 and 192 bytes. The original two-plain-column reproduction now returns to zero. Please make reader initialization exception-safe inside ArrowReaderIterator/the reader factory too, so a failed first batch releases that reader's own allocations.
There was a problem hiding this comment.
Fixed in 70f046a. ArrowReaderIterator now closes the reader if its eager first decode throws, so a column that loads its dictionary and then fails on the record batch releases it. I also guarded StreamReader's own getVectorSchemaRoot, which allocates the root's vectors under the same condition: nothing holds the reader until the constructor returns, so nothing else can close it.
Regression test: Comet in-memory cache releases a reader whose own first batch fails to decode. Instead of an allocator limit it truncates the decoded Arrow stream, dropping the end-of-stream marker plus part of the record batch body, so the dictionary loads and the read that follows runs out of input. Reverted, it leaks 64 bytes per attempt, matching the growth you measured; with the fix the allocator returns to baseline.
The existing corruptColumnStream helper was not usable for this. A small column compresses to a single LZ4 block, so truncating the compressed bytes fails the decompressor before Arrow reads anything at all, which is the case that was already covered. The new helper cuts the decoded bytes and re-compresses.
| case class CometInMemoryTableScanExec( | ||
| originalPlan: InMemoryTableScanExec, | ||
| serializer: CachedBatchSerializer, | ||
| cacheBuilder: CachedRDDBuilder, |
There was a problem hiding this comment.
[P2] Canonicalize the embedded scan to preserve exchange reuse
This leaf extends CometExec, not CometNativeExec, so it does not get the latter's canonicalization of embedded plan fields. originalPlan is not a child and retains its expression IDs even after the Comet scan is canonicalized. Equivalent references to the same cache therefore fail sameResult.
With AQE disabled, a UNION of two identical grouped aggregates over one cache produces one shuffle exchange plus ReusedExchange when native cache scanning is disabled, but two independent shuffle exchanges when enabled. Repeating a broadcast join has the same one-versus-two behavior. In the canonical scans, cacheBuilder, relationOutput and scanOutput compare equal; only originalPlan differs, and canonicalizing that field restores equality. Please canonicalize the wrapped scan while preserving its predicates and source semantics to avoid repeating scans, aggregation, shuffle and broadcast work.
There was a problem hiding this comment.
Fixed in 70f046a. CometInMemoryTableScanExec now overrides doCanonicalize and canonicalizes originalPlan by deferring to InMemoryTableScanExec, which normalizes its own attributes and predicates against the relation's output. relationOutput and scanOutput were already normalized by QueryPlan, and cacheBuilder and serializer are the same instances, so that one field was the whole difference.
Regression test: Comet in-memory cache scans of one cache canonicalize equal, so exchanges are reused, using your UNION-of-two-aggregates shape. It asserts one exchange plus one ReusedExchangeExec; reverted, it finds two exchanges and no reuse.
It also asserts that two scans differing only in their pushed predicates still canonicalize unequal. That is the reason for deferring rather than nulling the field the way CometNativeExec.canonicalizePlans does for embedded plans, which would have equated those as well.
Three findings from review, each with a regression test that fails without the fix: - A broadcast of a cache scan re-serializes each decoded batch as one stream covering every column, but `getBatchFieldVectors` handed the writer only the first column's dictionary provider. The columns are decoded from separate streams, so they arrive with separate providers and the write failed with "Could not find dictionary with ID 1". Combine the providers instead, and refuse a genuine ID clash rather than resolve one column against another's dictionary. - Opening an Arrow reader decodes the column's first batch, and a dictionary encoded column loads its dictionary before the record batch that indexes into it. A failure in between left that dictionary owned by a reader nobody holds: the constructor never returned. Close the reader on the way out, in `ArrowReaderIterator` and in `StreamReader`'s own schema read. - `CometInMemoryTableScanExec` carries the wrapped Spark scan as a plan-typed field rather than a child, so canonicalization walked past it and left its expression IDs in place. Two scans of one cache compared unequal, and a UNION of two identical aggregates over it ran two shuffles where Spark's cache scan runs one and reuses it. Canonicalize the wrapped scan, which keeps scans with different pushed predicates distinct. Also fix the Spark 3.4 failure in "does not build the cached RDD while planning": 3.4 defaults `canChangeCachedPlanOutputPartitioning` to false, which force-disables AQE inside the cached plan, so there is no finalization for the test to observe. Set the conf so the relation stays adaptive on every version.
|
Reviewed the current HEAD ( One thing I'd like to see addressed before merge, and one nit.
This is the Kryo-registration gap you already recorded in #5487 ("caching fails outright under #5487 (and the natural reading of it) frames this as only affecting explicitly serialized storage levels, with DISK_ONLY as the path that exercises it. But a
So under The existing DISK_ONLY test doesn't cover it — the suite never sets a Kryo serializer, so that test runs on the default Java serializer. Suggestion, in order of preference:
I'd lean against leaving it purely as #5487 given that plain Nit: the Overall I think this is mergeable as an experimental, default-off feature along the "merge and iterate" lines you proposed, with the Kryo item as the one I'd want handled (fix or document) first. The AQE gap (#5245) — especially the SPARK-37742 stats-divergence you called out, where Comet's own size accounting feeds AQE's join-strategy choice — is the follow-up I'd weight most after this lands. |
|
One small accuracy point on the benchmark, separate from the above.
So unless I'm misreading the plan, the two cases differ only at the cache-scan boundary, not in the execution engine above it:
The Worth tightening the |
…cache benchmark framing Under spark.kryo.registrationRequired=true, Kryo rejects any class it has not been told about. Two Comet payloads reach it: * CometCachedBatch, whenever Spark serializes a cached block. That is not only DISK_ONLY: the _SER levels, replication, cross-executor fetches, and the disk half of the default MEMORY_AND_DISK all serialize, so a plain df.cache() that spills is enough. * The Array[ChunkedByteBuffer] a native broadcast broadcasts. Spark registers ChunkedByteBuffer but not an array of them, so CometBroadcastExchangeExec fails here regardless of which Comet features are enabled. That one predates the cache work. Spark registers its own ArrowCachedBatch in KryoSerializer.loadableSparkClasses; Comet cannot add to that list, and spark.kryo.registrator is read when SparkEnv builds the serializer, before any plugin runs, so Comet cannot set it either. So: provide CometKryoRegistrator, document it on the cache config, and have CometDriverPlugin warn at startup when the combination is unsafe. DefaultCachedBatch is registered too, because a schema this serializer cannot store is delegated to Spark's, and Spark only registers that class itself from 4.1 onwards. Separately, correct what CometInMemoryCacheBenchmark claims to measure. Both cases run the aggregation on Comet; only the cache-scan boundary moves. The case labels and comment said "Spark execution over CometCachedBatch", which reads as a Spark-execution baseline it never was. verifyPlan now asserts the CometSparkColumnarToColumnar bridge in the disabled case so the framing is pinned by the benchmark itself.
|
@viirya thanks — you were right on both counts, and the Kryo one turned out to be wider than either of us had it. Fixed in 7e19dff. On the blast radius. Agreed, and your framing is what made me go looking. #5487 framed it as a DISK_ONLY concern and that was under-scoped: the But it is not only the cache. While writing the test I checked whether the registrator needed anything beyond the cached batch, and Same registration fixes it, since the cache write path hands back the same type, so I fixed it here rather than filing it — but it is a pre-existing bug rather than one this PR introduces, so say the word if you would rather it were split out. One thing I could not do, and it changes the shape of the fix. Comet cannot install this for the user. So the fix is three parts rather than one:
Two findings worth recording from building the list.
Testing. I did not test the On the nit — confirmed, that is intended and not a bug. Because Agreed on AQE (#5245) being the follow-up to weight most, and SPARK-37742 within it for the reason you and @peterxcli both landed on. |
|
@viirya you read it correctly and I had it wrong. Fixed in 7e19dff. Here is the plan for the disabled run, which settles it: The aggregation is What changed:
One addition beyond what you asked for: I could not re-run the benchmark itself in my current environment — |
|
Filed the broadcast half of the Kryo work as #5510 so it can be split out if reviewers prefer. To be clear about what is and is not in this PR's scope: The two are entangled only through the shared class list ( @viirya @sunchao happy either way — say which you would rather review and I will restructure, or leave it as is if folding it in here is fine. |
viirya
left a comment
viirya
left a comment
There was a problem hiding this comment.
Both addressed in 7e19dff20, and I went through the fix rather than just the replies — it's complete, and in the Kryo case more thorough than my comment asked for.
Kryo. The registration is transitively complete, which is the part that actually matters under registrationRequired=true — registering CometCachedBatch alone would still fail. arrowBytesKryoClasses covers ChunkedByteBuffer, Array[ChunkedByteBuffer], and the on-heap chunk types inside it (Array[ByteBuffer] and the concrete ByteBuffer class); kryoClasses covers the stats row's reachable types (GenericInternalRow, Array[Any], UTF8String, Decimal, and the BigDecimal/BigInteger a high-precision Decimal carries), which is exactly where I'd expected a gap and there isn't one. Registering DefaultCachedBatch for the pre-4.1 delegated path is a nice catch too. The new suite exercises DISK_ONLY and MEMORY_AND_DISK_SER with a column of each stats-tracked type, so the round-trip is actually covered rather than asserted.
Good find on the broadcast half (#5510). That's the sibling instance the finding pointed at without my naming it: CometBroadcastExchangeExec hands the same Array[ChunkedByteBuffer] to Spark's serializer, so a native broadcast join already fails under registrationRequired=true on main today, cache or no cache. Fixing it here because the registrator covers the same type is reasonable; splitting it out to keep it independent of an experimental feature is also reasonable — either way it shouldn't wait on this PR.
The one thing that stays a caveat is that spark.kryo.registrator can't be auto-set the way spark.sql.cache.serializer can, since it's read before the plugin runs. The startup warning on the unsafe combination plus the doc on the config is the right handling given that ordering — a user can't have it set silently for them, so being told at startup is the next best thing.
Benchmark. Pinning the framing with a verifyPlan assertion on the CometSparkColumnarToColumnar bridge in the disabled case is better than the comment fix I suggested — it keeps the "both cases run on Comet; only the scan boundary moves" claim honest against future edits rather than in prose that can drift. Thanks for the disabled-run plan; it settles it.
Both of my pre-merge items are resolved from my side. The remaining things I'd track are all post-merge and non-blocking: the AQE coverage in #5245 — SPARK-37742 most of all, since Comet's own size accounting feeds AQE's join-strategy choice and a divergence there would change plans silently — and the per-column framing follow-ups in #5487. This looks mergeable to me as an experimental, default-off feature.
Thanks @andygrove , that makes sense. I'm in favor of splitting the broadcast fix into its own PR under #5510, since it affects existing native broadcast joins without the cache feature. The shared registrations, registrator, startup warning, and a broadcast regression test can land there first, with #5051 keeping the cache-specific registrations and tests. That lets us fix the existing bug independently and keeps this PR focused. |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed commit 7e19dff202288330c74bd777ae7e63dd02db0ae7. The earlier cleanup and exchange-reuse regressions pass. Targeted probes reproduced two remaining native cache execution failures and unnecessary AQE shuffles on cold-cache reads. Five inline findings cover those issues, the missing CI suite registration, and the benchmark projection.
Verification: fresh native and JVM builds on Spark 4.1.3/JDK 17; all 260 selected tests across nine suites and 40 strict-Kryo component checks passed. Separate Spark/native probes reproduce the failures, with fallback and isolated code controls confirming both execution regressions. Older Spark versions received source inspection only.
Review assisted by Codex.
| builder: OperatorOuterClass.Operator.Builder, | ||
| childOp: Operator*): Option[Operator] = { | ||
|
|
||
| val scanTypes = op.output.flatMap(attr => serializeDataType(attr.dataType)) |
There was a problem hiding this comment.
[P2] Widen nested nullability in the cache scan schema
This serializes containsNull=false/valueContainsNull=false into the native scan schema, reintroducing the nested-type failures already fixed for CometLocalTableScanExec. On Spark 4.1.3, caching Seq(Seq(1, 2, 3), Seq(4, 5)) and evaluating slice(x, 2, 2) fails because spark_array_slice returns List(non-null Int32) while its declared result is List(Int32). map_entries over a cached Map[Int, Int] similarly panics on the value field's nullability. Both work uncached and with native cache scanning disabled over the same payload. An isolated change to serializeDataType(attr.dataType.asNullable) fixes both. Please preserve that normalization at this scan boundary and add cache regressions.
| case Some(existing) if existing.getVector ne dictionary.getVector => | ||
| throw new SparkException( | ||
| s"Columns of the same batch carry different dictionaries under ID $id") |
There was a problem hiding this comment.
[P2] Handle equivalent dictionaries from repeated cached columns
Selecting one cached dictionary column twice opens two readers for the same IPC stream. Their dictionary IDs and values match, but their vector objects differ, so this identity check rejects valid data. With a dictionary-encoded cache produced through JVM shuffle and AQE disabled, spark.range(1).join(broadcast(cached.select(cached("s1"), cached("s1"))), lit(true), "inner").collect() fails with 'Columns of the same batch carry different dictionaries under ID 0'. Both plans use CometBroadcastExchange; disabling only native cache scanning returns all 2,000 rows correctly. Please support equivalent dictionaries or normalize IDs instead of requiring object identity.
| extends CometExec | ||
| with LeafExecNode { |
There was a problem hiding this comment.
[P2] Preserve AQE cache materialization and reoptimization
Replacing Spark's InMemoryTableScanLike with an ordinary leaf during AQE preparation prevents TableCacheQueryStageExec creation. On Spark 4.1.3, first reading an adaptive cached join through a grouped aggregate executes an extra outer shuffle with native caching enabled: native disabled materializes the cache and removes that shuffle, while enabled retains it. A join against a cold cache already partitioned by its join key similarly runs two outer shuffles instead of one. Warm-cache queries avoid the difference. Results are correct, but the cold path unnecessarily reshuffles data. Please preserve materialization and reoptimization at the cache boundary. This is an executed failure of the scenarios requested in #5245, not just a coverage concern.
| * This needs its own suite because `spark.serializer` and `spark.kryo.registrator` are read when | ||
| * `SparkEnv` builds the serializer, so they cannot be changed per test. | ||
| */ | ||
| class CometInMemoryCacheKryoSuite extends CometTestBase { |
There was a problem hiding this comment.
[P2] Add the new Kryo suite to both CI test matrices
This new suite is absent from the suite lists in both .github/workflows/pr_build_linux.yml and .github/workflows/pr_build_macos.yml. python3 dev/ci/check-suites.py exits 255 with 'Suite not found in workflow ... CometInMemoryCacheKryoSuite'; the current GitHub Preflight job fails for exactly this reason and the platform/Spark test jobs are skipped. Add the suite to both matrices so CI can validate the PR and keep exercising the new Kryo regressions.
| runCacheBenchmark( | ||
| "in-memory cache full projection (6 of 6 columns)", | ||
| s"SELECT count(id), count(k), count(v), count(s1), count(s2), count(s3) FROM $cacheTable") |
There was a problem hiding this comment.
[P3] Make the full projection benchmark consume all six columns
The Range-derived id, v, and s3 columns are non-nullable, so Catalyst rewrites count(id), count(v), and count(s3) to count(1). Both the native and fallback benchmark plans therefore scan only k, s1, and s2, despite this case being labeled 6 of 6 columns. This also omits the highest-cardinality string column s3. Use expressions that consume all six columns and assert the actual cache-scan projection, otherwise the benchmark does not measure the full-width per-column decoding cost it claims.
9 participants
Which issue does this PR close?
Closes #2391.
Rationale for this change
This PR continues the work started by @pchintar in #4591, who is no longer able to work on it. All credit for the original implementation goes to them. This branch is based on their branch with latest
apache/mainmerged in, and picks up the remaining review feedback and CI failures.Comet currently has limited support for Spark's in-memory cache.
When a table is cached and later read, the cached data cannot be consumed directly by Comet operators. Instead, the execution plan falls back to Spark's cache scan path and introduces an additional
CometSparkColumnarToColumnarconversion before execution can continue in Comet.This extra conversion adds overhead to cached table scans and prevents cached data from remaining on a native Comet execution path.
What changes are included in this PR?
A native cache path for in-memory cached tables behind a new configuration:
spark.comet.exec.inMemoryCache.enabledWhen enabled:
CometCachedBatch.CometInMemoryTableScanExec.CometSparkColumnarToColumnarconversion.SimpleMetricsCachedBatchSerializer, so Spark'sbuildFiltercan prune cached batches before they are decoded.When disabled:
How are these changes tested?
CometInMemoryCacheSuitecovers:CometCachedBatchSELECT count(*))DefaultCachedBatchThe two projection tests corrupt the streams a read must not touch, so they assert on what was decoded rather than on timings, and each is checked to fail against the previous single-stream format.
CometInMemoryCacheKryoSuitecoversspark.kryo.registrationRequired=true, which makes Kryo reject any unregistered class. It runs withKryoSerializerandspark.kryo.registrator=org.apache.comet.CometKryoRegistratorand asserts:CometCachedBatchround-trips atDISK_ONLYandMEMORY_AND_DISK_SER, over a relation carrying every type whose bounds the statistics row records, read back through a predicate so the bounds themselves go through Kryomaintoday independently of this feature:CometBroadcastExchangeExecbroadcasts anArray[ChunkedByteBuffer], and Spark registersChunkedByteBufferbut not an array of themDefaultCachedBatchpath survives, which needs Comet to register that class on Spark 3.4, 3.5 and 4.0 because Spark only registers it itself from 4.1spark.kryo.registratoris read whenSparkEnvbuilds the serializer, which is beforeCometDriverPluginruns, so Comet cannot set it for the user the way it setsspark.sql.cache.serializer. It is documented onspark.comet.exec.inMemoryCache.enabledandCometDriverPluginwarns at startup when Kryo,registrationRequired, and a missing registrator are combined.CometInMemoryCacheBenchmarkcompares the native cache scan against the fallback read path over the same Comet-written cache, for a repeated full scan, a selective filter exercising cache pruning, and row-count-only (0 of 6 columns), narrow (1 of 6) and full (6 of 6) projections.Note what this benchmark does and does not isolate. Comet execution is on in both cases, so the aggregation runs on Comet either way and only the cache-scan boundary moves: disabled, Spark's
InMemoryTableScanExecfeeds those same Comet operators through aCometSparkColumnarToColumnarbridge; enabled,CometInMemoryTableScanfeeds them directly. So these numbers measure "keep the cached scan native" against "fall back to a Spark cache scan and convert", not Comet execution against Spark execution.Neither case is a baseline for Spark's own cache format either.
spark.sql.cache.serializeris a static config, so a single session cannot also materialize aDefaultCachedBatchto compare against; both cases read the same Comet-writtenCometCachedBatch. Earlier revisions of this description presented these numbers as a speedup over Spark's cache, which they are not, and then described the disabled case as "Spark execution over CometCachedBatch", which reads as a Spark-execution baseline it also never was.Benchmark results (release build, Apple M3 Ultra, JDK 17, Spark 3.5 profile, 5M-row cached table):
Repeated full scan (
SELECT sum(id), sum(k), sum(v))Selective filter (
WHERE id >= 4500000 AND id < 4750000)Comparison against Spark's cache format
Measured outside the benchmark harness, using two sessions so that one materializes a
DefaultCachedBatchand the other aCometCachedBatch. 5M rows, 6 columns (3 longs, 3 strings), min of 5 runs after a discarded warm-up:count(*)(0 of 6 columns)The first revision of this PR stored each cached batch as a single Arrow IPC stream covering every column, which had to be inflated in full before any projection could be applied. Narrow reads were then several times slower than Spark's cache, which stores columns separately. Reported in #5051 (comment) and tracked as #5484, now fixed here:
count(*)Per-column framing costs an Arrow schema block and compression framing per column per batch, and gives up cross-column compression. Footprint grows 2.5% at 6 columns and 32% at 60, where the cached relation is still 22x smaller than Spark's format (107 MiB versus 2369 MiB).
Remaining gap
Reads that feed Spark operators rather than Comet ones are still slower than Spark's cache, by roughly 1.7x to 2.5x depending on projection width. Those pay a row conversion that Spark's cache format avoids with a codegen'd path over its own layout. That is a separate cost from the projection issue fixed here and is not addressed in this PR. It is why the feature stays off by default and why the limitation is documented on
spark.comet.exec.inMemoryCache.enabled.Performance follow-ups are tracked in #4781. AQE test coverage is tracked in #5245.