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 12 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.
|
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.
| 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. |
There was a problem hiding this comment.
@andygrove thanks for updating the patch!
AQE tests: not done. ..... — I have not filed one, happy to if you want it.
sure, would be appreciate if you could help to file one, I can help to land it.
the timezone question is settled in c25c7b3 ...
thanks for the update, tz stuff lgtm!
| } catch { | ||
| case NonFatal(e) => | ||
| root.close() | ||
| throw e |
There was a problem hiding this comment.
- https://github.com/apache/spark/blob/78147abefb29096930346c3f9abaf92bdd5de4bb/sql/core/src/main/scala/org/apache/spark/sql/execution/r/ArrowRRunner.scala#L95-L117
- https://github.com/apache/spark/blob/78147abefb29096930346c3f9abaf92bdd5de4bb/common/utils/src/main/scala/org/apache/spark/util/SparkErrorUtils.scala#L73-L104
looked through spark's arrow writer usage, FYI.
| // scalastyle:off println | ||
| println( | ||
| "DIAG rows=" + df.collect().length + " metrics=" + scans.head.metrics | ||
| .map { case (k, v) => k + "=" + v.value } | ||
| .mkString(",")) | ||
| println("DIAG plan=" + df.queryExecution.executedPlan.getClass.getName) | ||
| // scalastyle:on println |
There was a problem hiding this comment.
reading a metric off
df.queryExecution.executedPlanaftercheckSparkAnswer(df)gives zero, because checkSparkAnswer executes its own copies. My first version passed vacuously with 0 on both sides of the comparison; it needs the df forced explicitly. The real numbers are 100 rows decoded with pruning on versus 1000 with it off.
Remove the println and add TODO or new issue link would be easier to followup
6 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(*))DefaultCachedBatchCometInMemoryCacheBenchmarkcompares the native cache path against the existing fallback path for repeated cached table scans and selective-filter scans exercising cache pruning.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)Performance follow-ups are tracked in #4781.