[SPARK-57268][SQL] Add Apache Arrow as a native cache format for in-memory Dataset caching

[viirya]

What changes were proposed in this pull request?

This PR adds Apache Arrow as a native cache format for Spark in-memory Dataset
caching, available alongside the existing DefaultCachedBatchSerializer. It is
one of the sub-tasks of SPARK-56978
(SPIP: Faster queries in local laptop mode for Apache Spark), specifically the
"Arrow-based df.cache reimplementation" item.

The new ArrowCachedBatchSerializer stores cached data in Apache Arrow IPC
streaming format. It is opt-in via spark.sql.cache.serializer:

spark.conf.set("spark.sql.cache.serializer",
  "org.apache.spark.sql.execution.columnar.ArrowCachedBatchSerializer")

Main components:

• ArrowCachedBatch -- a SimpleMetricsCachedBatch holding numRows, the
  serialized Arrow RecordBatch (IPC streaming format, optionally compressed),
  and per-column statistics for partition pruning.
• ArrowCachedBatchSerializer -- the serializer:
  • Write paths for both InternalRow and ColumnarBatch input, with a
    zero-copy fast path when the input is already backed by ArrowColumnVector.
  • Read paths for both ColumnarBatch output (wrapping Arrow vectors directly)
    and InternalRow output. The row path uses pre-built typed
    ArrowColumnReaders that write directly into an UnsafeRowWriter to avoid
    per-row pattern matching, and falls back to a columnar-to-row path for
    complex types (Array/Struct/Map/UDT/Variant/etc.).
  • Optional background prefetch of the next batch (decompress/deserialize off
    the consumer thread), controlled by a new config (off by default).
  • Min/max statistics collection over Arrow vectors, kept consistent with the
    row-based ColumnStats path (NaN handling, collation-aware string
    comparison, null/decimal bounds).
• ArrowUtils.isSupportedByArrow -- recursive type-support check used by
  supportsColumnarInput.
• ObjectColumnStats -- now skips getSizeInBytes for columnar complex
  types (ColumnarArray/ColumnarMap/ColumnarRow), which are views into
  ColumnVectors and do not expose a size.
• New config spark.sql.execution.arrow.cache.prefetch.enabled (default
  false), Kryo registration for the new classes, and documentation
  (sql-arrow-cache-format.html, linked from the SQL docs menu).

Why are the changes needed?

The default cache format is row/column-encoded specifically for Spark. Using
Arrow as the cache format provides:

• Zero-copy columnar reads when the cached data is already in Arrow form (e.g.
  re-caching Arrow-cached data with column projection).
• Interoperability with the Arrow ecosystem and off-heap memory management via
  Arrow allocators.
• Min/max statistics for partition pruning, consistent with the default path.

In our benchmarks, the Arrow format is competitive with or faster than the
default format on columnar/primitive workloads, with the largest gains on the
zero-copy re-cache path. The default format can still be faster in some cases
(for example, at higher compression levels), so this is offered as an opt-in
alternative rather than a replacement. See the committed
sql/core/benchmarks/ArrowCacheBenchmark-jdk{17,21,25}-results.txt files,
generated by the ArrowCacheBenchmark suite via the GitHub Actions benchmark
workflow.

Does this PR introduce any user-facing change?

Yes, additively. A new opt-in cache serializer
(ArrowCachedBatchSerializer) and a new config
spark.sql.execution.arrow.cache.prefetch.enabled (default false) are added.
The default cache behavior is unchanged: spark.sql.cache.serializer still
defaults to DefaultCachedBatchSerializer.

How was this patch tested?

• New ArrowCachedBatchSerializerSuite covering primitive and complex/nested
  types, null handling, collation, NaN bounds, statistics correctness for both
  the row and columnar (Arrow-vector) paths, columnar input from Parquet,
  column projection, filter pushdown, and compression codecs (none/zstd/lz4),
  plus a check that the Arrow serializer is actually used.
• ArrowCachedBatchKryoRegistrationSuite verifying Kryo registration.
• Added ArrowCacheBenchmark for performance comparison against the default
  cache format. Result files for JDK 17/21/25 are generated in the consistent
  GitHub Actions environment via the benchmark workflow.

Locally: catalyst/compile + sql/Test/compile pass; the two suites above run
green (0 failures).

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Claude Opus 4.8)

[sunchao]

Summary

This PR introduces an opt-in Arrow-native implementation of Spark's in-memory Dataset cache. It is a substantial and promising addition: the implementation covers row and columnar inputs, row and columnar outputs, IPC serialization with compression, off-heap Arrow memory, cache statistics, zero-copy handling for Arrow-backed input, and optional asynchronous prefetching. The overall architecture follows CachedBatchSerializer cleanly and keeps the existing default serializer unchanged.

I reviewed the immutable head 0b7cceed26084ab35a47896b9915e7b52954fc57 with five independent passes and targeted runtime reproductions. The design direction looks reasonable, but I found six correctness issues in the new serializer: three P1 issues that can produce wrong results, invalid memory access, or non-terminating cache materialization, and three P2 issues affecting supported input shapes and iterator completeness.

Prior state and problem

Spark's existing in-memory cache uses a Spark-specific encoded representation. It integrates tightly with Spark execution and pruning, but it is not Arrow-native and cannot directly preserve Arrow vectors across cache boundaries. This PR addresses that gap for workloads where Arrow interoperability, off-heap vectors, compression, or Arrow-to-Arrow reuse are valuable.

The serializer is intentionally additive and opt-in through spark.sql.cache.serializer, so existing applications retain DefaultCachedBatchSerializer unless they select the new implementation. That is a sensible compatibility boundary for a cache format with different memory, serialization, and performance characteristics.

Design approach

The patch introduces ArrowCachedBatch as a SimpleMetricsCachedBatch containing an Arrow IPC record batch plus Spark-compatible per-column statistics. ArrowCachedBatchSerializer implements both InternalRow and ColumnarBatch write paths. Arrow-backed columnar input gets a fast path that reuses the underlying vectors; other inputs are materialized into a new VectorSchemaRoot through ArrowWriter.

For reads, the serializer can return ColumnarBatch objects backed by loaded Arrow vectors, or produce UnsafeRows through prebuilt typed readers. Complex types use a fallback through ColumnarBatch and UnsafeProjection. The patch also adds codec selection, Arrow capability checks, Kryo registration, documentation, benchmarks, and optional background prefetching.

Correctness / compatibility analysis

The default cache serializer and default prefetch setting remain unchanged, which limits compatibility risk for applications that do not enable this feature. IPC serialization, projection by expression ID, task-scoped allocators, and Spark's existing cache-statistics shape are appropriate integration choices.

However, the current implementation has correctness gaps at several ownership and contract boundaries. Columnar prefetch advances an iterator that owns and closes the vectors still exposed to the consumer. NaN bounds are not conservative under Spark SQL ordering and can cause incorrect partition pruning. The shared Arrow batch-size configuration does not retain its documented nonpositive-as-unlimited meaning, which can leave the write iterator producing empty batches forever. There are also unsupported but accepted Arrow/geospatial representations and an iterator termination bug around empty batches.

Key design decisions

• The feature remains opt-in through the existing static cache serializer configuration.
• Cached payloads use Arrow IPC record batches, with none, LZ4, and ZSTD codec choices.
• Input already backed by ArrowColumnVector uses a zero-copy extraction path.
• Columnar reads expose Arrow vectors directly; row reads use typed readers and direct UnsafeRowWriter output.
• Statistics retain the existing five-value per-column layout so SimpleMetricsCachedBatchSerializer pruning can be reused.
• Allocators and roots are task scoped, with cleanup registered through task-completion listeners.
• Prefetch is disabled by default and uses one executor per iterator when enabled.

These choices are individually reasonable, but the root/vector lifetime must remain valid until downstream consumption completes, and advertised type support must cover statistics and both output modes rather than schema conversion alone.

Implementation sketch

On the write side, each input partition is divided into cached batches, written or unloaded into Arrow buffers, optionally compressed, serialized through MessageSerializer, and paired with collected statistics. On the read side, each IPC payload is deserialized into a task allocator, loaded into a VectorSchemaRoot, projected to selected columns, and either returned as a ColumnarBatch or read into reusable UnsafeRow storage.

The zero-copy path bypasses row materialization for Arrow-backed input but still traverses vectors to calculate statistics. The row fast path creates one typed ArrowColumnReader per selected field and swaps its vector reference when loading a new batch. The prefetch wrappers attempt to overlap deserialization/decompression of the next batch with consumption of the current one.

Behavioral changes worth calling out

Users enabling the serializer get a different cache wire format, Arrow allocator usage, codec behavior, and batch-size handling. The new serializer also participates in Spark's cache partition pruning, so its statistics are correctness-critical rather than advisory. spark.sql.execution.arrow.cache.prefetch.enabled introduces concurrency and an additional batch of memory when enabled.

The blocking findings are:

• P1: Columnar prefetch can close the root backing the batch currently being consumed.
• P1: Float/Double statistics can prune batches containing NaN even when those rows satisfy the filter.
• P1: Nonpositive maxRecordsPerBatch can repeatedly emit empty batches without consuming input.
• P2: LargeVarCharVector input crashes the zero-copy statistics path.
• P2: Top-level Geometry and Geography are advertised as supported but fail during cache materialization.
• P2: A legal empty cached batch terminates row iteration and hides later populated batches.

Suggested improvements

Please resolve the six inline correctness findings before merging and add focused regressions for each boundary:

• Exercise delayed consumption of the current columnar batch while prefetching the next one.
• Compare pruning enabled and disabled for mixed finite/NaN and all-NaN Float/Double batches.
• Cover maxRecordsPerBatch values 0 and -1, including full materialization termination.
• Round-trip LargeVarCharVector through zero-copy write and both output paths.
• Materialize and read top-level fixed and ANY Geometry/Geography columns, or reject them consistently.
• Feed row conversion a sequence containing empty and nonempty cached batches and verify no rows are lost.

For validation, I built a temporary six-case positive-behavior suite against this exact head. All six tests failed at the predicted mechanisms: released Arrow buffers, zero-row output for the unlimited setting, a missing NaN result, the exact LargeVarCharVector cast failure, UNSUPPORTED_DATATYPE for top-level Geometry, and truncation after an empty cached batch.

[sunchao]

[P1] Please avoid advancing the root-owning iterator before the returned batch has been consumed. submitPrefetch() invokes underlying.next() on the background thread, and that method immediately closes previousRoot; however, the ColumnarBatch returned above still contains ArrowColumnVectors backed by that root. With at least two batches and spark.sql.execution.arrow.cache.prefetch.enabled=true, the consumer can therefore read released buffers. I reproduced this with two one-row batches: delaying the first read until prefetch starts makes getInt(0) fail with IndexOutOfBoundsException. Prefetch needs separate ownership/retention for each returned batch, or it must deserialize the next batch without advancing an iterator that closes the current root.

[viirya]

Fixed. You're right that the background prefetch advanced an iterator that closes the current root. I folded prefetch into ArrowCachedBatchToColumnarBatchIterator (mirroring the row-read iterator): the background thread now only deserializes the next batch into its own fresh root, and the previous root is closed exclusively on the consumer thread in next(), so a returned batch's vectors are never released while it may still be read. The standalone ArrowPrefetchColumnarBatchIterator wrapper is removed. Added a prefetch read test that I verified fails (IndexOutOfBoundsException) with the previous wrapper and passes now.

[sunchao]

[P1] Skipping NaN makes these bounds unsafe for cache pruning. Spark SQL orders NaN above every non-NaN value, so a batch containing [1.0, NaN] must satisfy v > 100.0; this code records both bounds as 1.0, causing InMemoryTableScanExec to discard the entire batch before evaluating the rows. I reproduced the query and it returned zero rows instead of NaN. The row-based cache statistics have the same inherited defect, but this new implementation explicitly propagates it into the Arrow path. Float/Double statistics need bounds that conservatively account for NaN, with mixed and all-NaN pruning tests.

[viirya]

Thanks for catching this and for the clear repro. You're right that the NaN bounds are not conservative under Spark SQL's ordering (where NaN sorts above everything), so a batch like [1.0, NaN] gets incorrectly pruned for v > 100.0. I confirmed the same wrong result with the default serializer: caching [1.0, NaN] and applying filter("v > 100.0") returns 0 rows instead of NaN.

I'd prefer not to fix this in the Arrow path alone here. This serializer's row path reuses the existing FloatColumnStats/DoubleColumnStats collectors directly (via createColumnStats), and the columnar calculateMinMax* path was deliberately written to match them. Diverging only the columnar path would make the two paths within this serializer inconsistent, and the underlying defect lives in the shared, pre-existing statistics collectors used by DefaultCachedBatchSerializer. Since this PR is an additive feature rather than a bug fix, I think the right scope is a separate ticket that makes NaN bounds conservative across both the row and columnar paths (default and Arrow together), with mixed and all-NaN pruning tests.

I'll keep the current behavior aligned with the default serializer here and file a follow-up.

[sunchao]

[P1] Please preserve the documented unlimited semantics for nonpositive spark.sql.execution.arrow.maxRecordsPerBatch. With 0 or a negative value this condition is false before consuming any row, so next() emits a zero-row batch while hasNext remains true because the source iterator was not advanced. Full cache materialization can consequently generate empty batches indefinitely rather than finishing. A one-row reproduction already returns numRows == 0; the existing Arrow conversion path uses (maxRecordsPerBatch <= 0 || rowCount < maxRecordsPerBatch), which is the guard needed here as well.

[viirya]

Fixed. Changed the loop guard to (maxRecordsPerBatch <= 0 || rowCount < maxRecordsPerBatch), matching ArrowConverters, so a nonpositive value means unlimited (one batch per partition) instead of emitting empty batches forever. Added a test for 0 and -1.

[sunchao]

[P2] This cast does not cover valid Arrow-backed Spark strings. ArrowColumnVector supports both VarCharVector and LargeVarCharVector, and the zero-copy input path accepts either without converting or validating the backing vector. A one-row LargeVarCharVector therefore reaches this line and fails with ClassCastException during statistics collection. Please handle both vector variants, and also ensure the row-reader path can consume the corresponding large string vector or deliberately route it through a compatible fallback.

[viirya]

Fixed, and it turned out to be worse than a cast failure. The zero-copy path reuses the input vectors but serializes/reloads them under a largeVarTypes=false schema, so a LargeVarCharVector (64-bit offsets) reinterpreted under that 32-bit-offset schema doesn't just throw in stats -- it silently corrupts data (I reproduced a roundtrip returning ["", "hello"] for input ["hello", "world"]). Rather than teach every path to carry the large-vector schema (which would mean persisting the Arrow schema per cached batch), I route any batch whose vectors are, or nest, a large var-width vector through the row-based slow path, which always produces standard var-width vectors. Added a test with a LargeVarCharVector input that now roundtrips correctly.

[sunchao]

[P2] Top-level Geometry and Geography columns are not currently supported end to end, despite being declared supported here. During ordinary cache().count(), createColumnStats routes these types to ObjectColumnStats, whose constructor calls ColumnType(dataType) and throws UNSUPPORTED_DATATYPE because ColumnType has no geospatial case. I reproduced this independently for GEOMETRY(0) and GEOGRAPHY(4326). Nested geospatial values can work, so the issue is narrower than all geospatial use, but top-level support should either be completed in statistics and row conversion or rejected by this capability check.

[viirya]

Fixed by completing top-level support rather than rejecting it. Geometry/Geography are stored as binary (WKB) internally and are already representable in Arrow (the nested case worked), so the only gap was the statistics hookup and the row read path: createColumnStats now returns BinaryColumnStats for them (size/count, null bounds), and needsFallback now routes them through the columnar-to-row fallback instead of the fast ArrowColumnReader. Top-level geometry/geography now roundtrip through both row and columnar reads. (DefaultCachedBatchSerializer has no column encoding for geospatial at all, so this is strictly more capable.) Also added a data-driven test that, for every type isSupportedByArrow claims, actually caches and reads it back -- this would have caught the original mismatch and guards against it recurring.

[sunchao]

[P2] hasNext must keep loading while the current batch has zero rows. A zero-row ColumnarBatch is legal input from a columnar data source and is preserved by convertColumnarBatchToCachedBatch; after loading such a batch, this returns false even when batchIter still contains populated batches. Normal iterator consumers then terminate permanently and silently lose all later rows. I reproduced this with an empty cached batch followed by a one-row batch: the row conversion returned no rows. Please loop past empty batches until a row is found or the input is exhausted.

[viirya]

Fixed. hasNext now loops past zero-row batches until a row is found or the input is exhausted, so an empty cached batch no longer terminates the row iterator early and drops subsequent batches. Added a test (empty batch followed by a one-row batch).

[zhengruifeng]

do we have test for duplicated column names?
I remember arrow doesn't support duplicated column names, so we have some special name handling in python<>jvm arrow exechange.

[viirya]

Good question -- there was no test, so I added one (duplicated column names roundtrip through the cache). It caches a DataFrame with three columns all named a, verifies both the row-based and vectorized read paths, and prunes a single one of the duplicated columns through the cache scan.

Duplicate names work without special handling here because nothing in the cache path is keyed by field name: vectors are accessed positionally (getFieldVectors / getVector(index)) and column pruning maps selected attributes to cache columns by exprId. Arrow schemas themselves permit duplicate field names. The renaming you remember from the Python<->JVM Arrow exchange is needed because pandas indexes columns by name on the Python side; cached batches never cross into Python, so that concern does not apply.

[dbtsai]

Zstd compression level is silently ignored. This branch constructs new ZstdCompressionCodec(compressionLevel) only to extract getCodecType() (a plain enum carrying no level), then rebuilds the codec via factory.createCodec(codecType) -- the single-argument overload, which produces a codec at the factory's default level. The captured compressionLevel (from spark.sql.execution.arrow.compression.level) never reaches the codec actually used by the VectorUnloader, so the documented config (range 1-22) has no effect for the Arrow cache; users tuning the compression level get the default regardless.

The codec instance that already carries the level can be used directly:

Suggested change

	case "zstd" =>
	val factory = CompressionCodec.Factory.INSTANCE
	val codecType = new ZstdCompressionCodec(compressionLevel).getCodecType()
	factory.createCodec(codecType)
	case "zstd" =>
	new ZstdCompressionCodec(compressionLevel)

[viirya]

Fixed by constructing the codec directly as you suggested -- createCompressionCodec now returns new ZstdCompressionCodec(compressionLevel) instead of rebuilding through the single-argument factory overload (the lz4 arm is simplified the same way). A comment notes why the factory must not be used here: the level only matters on the write side, since reads look up the codec by the type recorded in the IPC message.

Digging into why the committed benchmark results never caught this, I found a second, independent bug: ArrowCacheBenchmark was setting a nonexistent conf key (spark.sql.execution.arrow.compression.level instead of spark.sql.execution.arrow.compression.zstd.level), which spark.conf.set accepts silently. So the per-level benchmark rows were doubly broken -- even with the codec fixed, the level would never have reached it. The benchmark now references the SQLConf key constants so a typo can no longer bind to nothing. The doc page had the same wrong key plus a few other inaccuracies (codec default listed as zstd instead of none, level range stated as 1-22 though negative fast levels are supported, and a maxRecordsPerBatch key missing the execution segment); all fixed.

Added a regression test that compresses the same batch at zstd level -5 and 19 and asserts the higher level yields a strictly smaller payload; it fails against the previous codec construction. The per-level rows in the committed benchmark result files are stale (all three levels effectively measured the default level -- they are within noise of each other, which corroborates your finding); I will regenerate them with the benchmark GitHub Actions workflow.

[acruise]

whoops, left this hanging a few days

[acruise]

This denotes an exclusive setting, do we expect any scenarios where this cache serializer wouldn't work and a fallback might be necessary?

[viirya]

The cache serializer is selected once by the static spark.sql.cache.serializer conf, and once set it handles every cached relation -- Spark's cache framework has no per-relation fallback to a different serializer based on data type. So there's no automatic switch back to the default cache driver. The Arrow serializer aims to cover everything the default serializer does (and more -- e.g. geometry/geography, which the default has no encoding for). If a column uses a type Arrow genuinely cannot represent (e.g. ObjectType), materializing the cache throws UNSUPPORTED_DATATYPE rather than silently falling back. I've clarified both points in the docs.

[acruise]

Presumably when this returns false for any reason, we fallback to the default cache driver, that should be made clear in the docs if it isn't already

[acruise]

e.g. the doc says supports all Spark SQL data types but this implementation would seem to falsify that claim ;)

[acruise]

Well, I guess the claim may be true, and the fallthrough at the end might be defensive... In which case maybe we'd want to log a surprisingly unsupported type :)

[viirya]

Good catch -- "supports all Spark SQL data types" did overstate it, and the type list was also incomplete (it omitted Time, intervals, Geometry/Geography, Variant, Null, and UDTs, which are all supported). Fixed the doc to list the actually-supported set and to describe the unsupported-type behavior.

One clarification on the mechanism, re: your fallback question: isSupportedByArrow here only gates supportsColumnarInput, which the cache framework uses to choose the columnar-vs-row input path into this same serializer -- it isn't a fallback to the default cache driver (there's no such per-type fallback). A truly unsupported type isn't silently dropped either: toArrowSchema throws UNSUPPORTED_DATATYPE when the cache is materialized. The docs now state this explicitly.

[uros-b]

cc @MaxGekk to take a look ^^

[sunchao]

Additional findings on the current head.

[sunchao]

[P1] Please release each consumed input ColumnarBatch. This serializer receives columnar input by removing the normal ColumnarToRow consumer, which calls closeIfFreeable() after each batch. Both conversion branches finish synchronously, but this method returns without closing the input on either success or failure. For an Arrow-backed source that produces fresh off-heap vectors, every cached batch therefore remains live and executor memory grows until OOM. Please wrap the conversion in finally { batch.closeIfFreeable() }; this is safe for reusable writable/constant vectors because their closeIfFreeable() is a no-op.

[viirya]

Fixed. next() now wraps the conversion in tryWithSafeFinally { ... } { batch.closeIfFreeable() }, releasing the consumed input batch on both success and failure of either branch. closeIfFreeable() is a no-op for reusable writable/constant vectors, as you noted.

[sunchao]

[P1] Please retain ownership of the prefetched root during task cleanup instead of cancelling and discarding this future. With a short-circuiting consumer such as LIMIT, the next root may already have been produced when task completion runs; cancel(true) then returns without exposing or closing that root, and allocator.close() throws Memory was leaked by query. If the worker is still running, cancellation also does not join it, so allocation can race allocator shutdown and a root returned after cancellation is dropped. The row-reader listener has the same issue. Cleanup needs to obtain and close any completed result, coordinate worker termination, and only then close the allocator.

[viirya]

Fixed. Cleanup now goes through a shared helper that shuts the prefetch executor down and awaits termination (rather than cancel(true), which could race an in-flight allocation), then retrieves and closes any root the worker already produced before closing the allocator. Applied to both the columnar-reader and row-reader listeners.

[sunchao]

[P1] I rechecked the cleanup change on the current head, and an interrupted task can still race allocator shutdown. If the completion listener runs with the task thread interrupted, awaitTermination throws; the catch restores the interrupt and calls shutdownNow() without joining the worker. When the future is still in flight, the subsequent future.get() immediately throws InterruptedException, so the caller proceeds to allocator.close() while Arrow deserialization may still allocate or return an unclosed root. Please drain/join uninterruptibly, close any result, then restore the interrupt, with killed-task coverage for both row and columnar readers.

[viirya]

Fixed. drainAndClosePrefetch now clears the interrupt for the duration, loops awaitTermination until the worker actually terminates (re-clearing if interrupted again), then retrieves and closes any produced root, and only restores the interrupt at the very end. So a completion listener running on an already-interrupted (killed) task still joins the worker and closes the root before the allocator is closed. Added a test that calls it from an interrupted thread with a prefetch in flight and asserts the produced root is closed (no leak on allocator.close()) and the interrupt is restored. Applied to both the row and columnar readers via the shared helper.

[sunchao]

[P1] Please honor both Arrow batch limits on both input paths. The row path captures only arrowMaxRecordsPerBatch and ignores arrowMaxBytesPerBatch (whose default is 64 MiB), while the columnar path captures neither limit and serializes each upstream ColumnarBatch wholesale. Thus wide rows can form multi-gigabyte batches despite the byte guard, and columnar input ignores even maxRecordsPerBatch=1. This can exhaust executor memory or hit Arrow's 32-bit variable-width offset limit before serialization. The row path should stop when either threshold is reached, and the columnar path needs slicing/splitting.

[viirya]

Partially addressed, with a narrower scope after looking into it. The row path now honors maxBytesPerBatch in addition to maxRecordsPerBatch (stopping at whichever is hit first), matching ArrowConverters. For the columnar path I kept the one-batch-in/one-batch-out behavior: the upstream ColumnarBatch row count is already bounded by the source's batch-size config (e.g. spark.sql.parquet.columnarReaderBatchSize, default 4096), and the zero-copy path reuses already-formed Arrow vectors that cannot exceed Arrow's offset limits or balloon memory beyond what the upstream batch already holds. Let me know if you'd still prefer per-row byte enforcement on the columnar/zero-copy paths.

[sunchao]

[P1] This remains incomplete in two ways. The columnar path still captures neither limit and explicitly maps one input batch to one cached batch. A normal repro is to cache relation A with maxRecordsPerBatch=10000, change the session setting to 1, then cache a projection of A: the Arrow-columnar recache preserves the 10,000-row input batches. On the row path, every non-UnsafeRow is estimated as only numFields * 16, so a one-field GenericInternalRow containing a 100 MiB string contributes 16 bytes to the nominal 64 MiB guard. Please split/slice columnar input and use the actual Arrow writer size (or equivalent type-aware sizing) for generic rows.

[viirya]

Split into the two halves. The byte-estimate half is fixed: the row path now measures arrowWriter.sizeInBytes() (the actual bytes written to the Arrow vectors) instead of numFields * 16, so a large value in a GenericInternalRow is accounted for correctly. Added a test that a tiny maxBytesPerBatch forces multiple batches.\n\nOn splitting columnar input: I'd prefer not to add columnar slicing for the recache scenario. The upstream ColumnarBatch row count is already bounded by the producing source's batch-size config, and batch size is a performance knob rather than a correctness one -- a recached relation that keeps its original batch sizing produces identical results. The repro requires lowering maxRecordsPerBatch specifically between caching A and recaching a projection of A, which doesn't reduce memory pressure (the upstream batches are already materialized) and I couldn't find a realistic workload for it. Happy to revisit if you have one in mind.

[sunchao]

[P2] This capability check accepts any UDT whose sqlType is supported, but the cache statistics path does not unwrap UDTs. A valid UDT backed by VariantType passes here, then createColumnStats constructs ObjectColumnStats; ColumnType(udt.sqlType) has no VariantType case and throws UNSUPPORTED_DATATYPE during materialization. The default cache serializer does unwrap the same UDT and supports it, so this also contradicts the documented claim that Arrow covers every default-supported type. Please unwrap UDTs when selecting statistics collectors, or narrow this predicate and the documentation.

[viirya]

Fixed. createColumnStats now unwraps UDTs (case udt: UserDefinedType[_] => createColumnStats(udt.sqlType)), so a Variant- or Geometry-backed UDT gets the right collector instead of falling through to ObjectColumnStats and throwing UNSUPPORTED_DATATYPE. This keeps the capability check and the statistics path in agreement.

[sunchao]

[P2] Please do not drop size accounting for columnar complex values. Vectorized nested Parquet/ORC input reaches the Arrow serializer as a non-Arrow ColumnarBatch, takes the row-conversion path, and exposes ColumnarArray/ColumnarMap/ColumnarRow here. Every non-null complex value then contributes zero bytes. Because SimpleMetricsCachedBatch.sizeInBytes is the sum of these fields, a large complex-only cached relation reports exactly zero bytes and becomes eligible for automatic broadcast. Please derive the size from the produced Arrow vectors/root or otherwise account for the columnar payload.

[viirya]

Fixed at the source: the columnar slow path now derives statistics from the built Arrow root via collectStatistics(root) (reading vector.getBufferSize) instead of the row collectors, matching the zero-copy path. Columnar complex values therefore contribute their actual byte size rather than zero, so a complex-only relation no longer reports sizeInBytes=0. The ColumnStats columnar-complex guard is left in place defensively.

[sunchao]

[P2] This migration does not actually switch serializers within the same JVM once any cache has been materialized. InMemoryRelation stores the first CachedBatchSerializer in a process-wide singleton and never clears it when SparkSession.stop() stops the context; the new benchmark has to call the private InMemoryRelation.clearSerializer() hook explicitly for this reason. After using the default cache, following these steps still silently uses DefaultCachedBatchSerializer (and the reverse direction likewise stays Arrow). Please add a real lifecycle reset or document that changing formats requires a fresh JVM.

[viirya]

Fixed. The migration section now states that the serializer is resolved once and held process-wide, so switching cache formats requires a fresh JVM; the in-process "stop and reconfigure" steps were removed.

[sunchao]

[P2] This enablement example always throws on an existing session. spark.sql.cache.serializer is a static SQL configuration, and RuntimeConfig.set rejects static keys with CANNOT_MODIFY_CONFIG. Please remove this snippet and retain only the SparkSession.builder.config(...) example below.

[viirya]

Fixed. Removed the spark.conf.set snippet (it throws CANNOT_MODIFY_CONFIG on the static key) and kept only the SparkSession.builder.config(...) example.

[sunchao]

[P2] Please close the newly created root if VectorLoader.load fails. VectorLoader loads fields incrementally, so a later malformed/decompression/OOM failure can occur after earlier vectors have retained or allocated buffers. This method currently closes only recordBatch; the partially loaded root becomes unreachable, and task cleanup then fails while closing the allocator and can obscure the original error. The row deserializer has the same issue. Use close-on-failure while returning the root unchanged on success.

[viirya]

Fixed. Both deserializeToRoot and deserializeBatch now close the partially-loaded root if VectorLoader.load fails, then rethrow, so a mid-load failure no longer leaves an unreachable root that makes the later allocator.close() fail and mask the original error.

[sunchao]

[P2] This benchmark does not isolate filter pushdown: the timed case includes session creation, cache construction, a full noop scan that materializes the cache, and only then this filtered action. Both serializers already collect min/max bounds, so the result cannot support the documentation's attribution of the measured gain to Arrow pruning. The committed JDK 21 results reinforce this: default-uncompressed is slightly faster than Arrow-uncompressed, while the advertised difference is against default-compressed. Please materialize during untimed setup and time only the filtered action, or rename the benchmark and remove the pruning attribution.

[viirya]

Fixed by renaming and de-attributing rather than restructuring the timing: the case is now "Cache then filter" with a comment that it measures end-to-end cache build + filtered scan between the two formats, not pruning attributable to Arrow (both serializers collect min/max). The docs table and note were updated to match.

[sunchao]

[P2] The source and documentation were renamed, but the three committed result files were not regenerated. They still contain Cache with filter pushdown, Cache 5M rows + filter, and Arrow cache - filter (with stats), and their generating commits predate this fix. Since the documentation calls these files authoritative, they still preserve the pruning attribution this change was meant to remove. Please regenerate all three result files from the current head.

[viirya]

Fixed. I've triggered the benchmark workflow (JDK 17/21/25) to regenerate all three result files from the current head, so they will carry the renamed "Cache then filter" labels and drop the stale pruning-attribution naming.

[sunchao]

[P3] This test claims to verify that every supported type can be cached and read back, but it checks only the number of returned rows. Any corruption that preserves row count passes. Please compare the returned value with the input using type-aware equality; the geometry/geography roundtrip test should likewise validate WKB/SRID rather than only nullness.

[viirya]

Fixed. The roundtrip tests now compare the cached read against an uncached baseline with checkAnswer (type-aware), so corruption that preserves row count is caught; the geometry/geography test validates the actual value (WKB/SRID) rather than only nullness.

[sunchao]

Additional findings on the current head.

[sunchao]

[P2] Please do not advertise CalendarIntervalType as fully supported unless the cache representation is lossless over Spark's full value domain. Arrow's IntervalMonthDayNanoWriter multiplies CalendarInterval.microseconds by 1000 with Math.multiplyExact, but Spark permits the full Long range. I reproduced this with new CalendarInterval(0, 0, Long.MaxValue / 1000L + 1L): the default cache returns COUNT=1, while the Arrow cache aborts materialization with ArithmeticException: long overflow. Please use a lossless representation or reject and document this type rather than claiming parity with the default serializer.

[viirya]

Fixed with a clear diagnostic. Caching a CalendarInterval whose microseconds exceed +/-(Long.MaxValue / 1000) now throws an explanatory error (naming the type and the nanosecond-conversion limit) instead of an opaque ArithmeticException: long overflow. The check is installed only when the schema actually contains a CalendarInterval column, so there is no per-row cost for other schemas. Arrow's IntervalMonthDayNano is nanosecond-based and cannot losslessly hold the full Long microsecond domain, so I documented the value-range limit rather than changing the shared Arrow writer; the default serializer's lack of this restriction is noted in the docs.

[sunchao]

[P2] The diagnostic fix still covers only top-level intervals. hasCalendarInterval checks attr.dataType == CalendarIntervalType, while isSupportedByArrow accepts CalendarInterval recursively inside arrays, structs, maps, and UDTs, and their recursive Arrow writers still reach Math.multiplyExact(microseconds, 1000L). For example, an array<interval> containing Long.MaxValue / 1000 + 1 still escapes withIntervalOverflowTranslation and fails with raw ArithmeticException: long overflow, contrary to the guide and this reply. Please recurse through complex types and UserDefinedType.sqlType, with nested overflow coverage.

[ 🤖 posted by Codex on behalf of sunchao using the code-review-for-me skill 🤖 ]

[viirya]

Fixed. hasCalendarInterval now recurses through arrays, structs, maps, and UserDefinedType.sqlType (a dedicated traversal mirroring isSupportedByArrow, since DataType.existsRecursively does not descend into UDT sql types), so nested intervals get the same clear diagnostic. Added an array<interval> overflow case to the test.

[sunchao]

[P2] Please narrow the CalendarInterval overflow translation to the actual interval conversion. The new schema-wide catch wraps the entire batch loop whenever any interval field exists, and SparkArithmeticException extends ArithmeticException. I reproduced this with one partition containing two rows under ANSI mode: range(0, 2, 1, 1).selectExpr("1 / (1 - id) AS bad", "make_interval() AS i").cache().count(). Without the harmless interval column, the second row fails as DIVIDE_BY_ZERO; with it, the same failure is caught inside the batch loop and rethrown as DATETIME_OVERFLOW claiming a CalendarInterval overflow. Adding an unrelated interval column must not change the error condition or SQLSTATE. Please catch the raw Math.multiplyExact failure at IntervalMonthDayNanoWriter, or preserve structured/non-interval arithmetic exceptions unchanged.

[viirya]

Confirmed -- SparkArithmeticException extends ArithmeticException, and the batch-loop try covers rowIter.next() where upstream lazy evaluation runs, so an unrelated ANSI error can indeed be re-labeled. Adding a harmless interval column must not change the error condition.

Rather than narrowing the catch, we plan to make the translation unnecessary and delete it outright: following the same approach as SPARK-57975 / #57053 (opt-in lossless struct representation for nanosecond timestamps), a follow-up will store CalendarInterval losslessly as a struct of (months: int32, days: int32, microseconds: int64) -- the type's own field layout, mirroring the default cache's CALENDAR_INTERVAL ColumnType. With no * 1000 conversion on the cache path, the overflow (and withIntervalOverflowTranslation along with this mis-attribution bug) ceases to exist, and the cache regains the default serializer's full value domain for this type as well. The shared IntervalMonthDayNanoWriter will also get a structured DATETIME_OVERFLOW at the Math.multiplyExact site (the TimestampNTZNanosWriter pattern) for the interchange paths that keep the standard MonthDayNano encoding.

[viirya]

Follow-up filed: SPARK-58005 / #57088 implements the plan above. It extends SPARK-57975's opt-in lossless encoding to CalendarInterval -- a struct of (months, days, microseconds), the type's own layout, so the * 1000 conversion and its overflow do not exist on that path -- and, independently, moves the overflow translation into IntervalMonthDayNanoWriter itself: the structured DATETIME_OVERFLOW is now raised by a catch scoped to the single Math.multiplyExact expression (the TimestampNTZNanosWriter pattern), so it structurally cannot re-label an unrelated ArithmeticException from upstream evaluation. Once both land, this PR will opt the cache into the lossless encoding and delete withIntervalOverflowTranslation entirely, which removes the mis-attribution you found along with the wrapper; your ANSI DIVIDE_BY_ZERO repro then surfaces DIVIDE_BY_ZERO regardless of interval columns in the schema.

[viirya]

Done. With SPARK-57975 and SPARK-58005 merged, this PR now opts the cache into the lossless encodings (losslessInternalTypes = true at the schema-construction sites) and deletes withIntervalOverflowTranslation entirely -- on the cache path there is no longer a nanosecond conversion that can overflow, so the mis-attribution you found is gone along with the wrapper: your ANSI DIVIDE_BY_ZERO repro surfaces DIVIDE_BY_ZERO regardless of interval columns. The overflow diagnostic test is replaced by a full-domain round-trip test (microseconds = Long.MaxValue etc., top-level and nested), and the docs now state the cache matches the default serializer's full value domain for this type.

[sunchao]

[P2] Please add the standard Spark Jekyll front matter (layout, title, displayTitle, and the ASF license block). This is the only tracked docs/*.md page without a leading ---; without front matter it will not produce the sql-arrow-cache-format.html page referenced by the new menu and performance-guide links.

[viirya]

Fixed. Added the standard Jekyll front matter (layout, title, displayTitle, and the ASF license block) so the page produces sql-arrow-cache-format.html.

[sunchao]

[P2] Geometry/Geography statistics still assume an UnsafeRow representation. BinaryColumnStats calls row.getBinary, but the Catalyst physical value for these types is BinaryView; a valid GenericInternalRow (for example from a row-based DSv2 reader or direct serializer use) therefore throws ClassCastException, even though ArrowWriter correctly consumes the same value through getBinaryView. Please use a BinaryView-aware size collector and cover a direct generic-row input.

[viirya]

Fixed. Added a BinaryView-aware GeoColumnStats and routed Geometry/Geography to it in createColumnStats. It reads the value via getBinaryView (matching how ArrowWriter consumes it) instead of row.getBinary, so a GenericInternalRow storing a BinaryView no longer throws ClassCastException. Added a test that drives the collector with a BinaryView in a generic row.

[sunchao]

[P2] Please reconcile this capability whitelist with current master before merging. master now maps and reads/writes TimestampNTZNanosType and TimestampLTZNanosType through Arrow, but this predicate still rejects them. The serializer then takes the row path, falls through to ObjectColumnStats (whose ColumnType has no nanos-timestamp case), and fails materialization with UNSUPPORTED_DATATYPE; row output also lacks a typed reader or fallback. The synthetic merge is conflict-free but still has this behavioral incompatibility.

[viirya]

Good catch that master now maps these through Arrow. I looked at the cache paths, though, and the physical value for TimestampNTZNanosType/TimestampLTZNanosType is a TimestampNanosVal, not a plain Long, so the cache's stats collector and fast columnar reader can't treat them as long-backed without a dedicated, precision-aware path -- that's net-new support rather than a predicate tweak. Importantly there's no parity regression: the default cache serializer doesn't support these types either. I verified on current master that df.cache() of a TIMESTAMP_NTZ(9) column throws not support type: TimestampNTZNanosType(9) with the default serializer (its ColumnBuilder has no case for them). The Arrow serializer now rejects them with a clear checkSupportedSchema error at materialization. I'd rather add real support as a focused follow-up than land a half-correct fast path here -- let me know if you'd prefer it gated differently in the meantime.

[viirya]

Follow-up filed: SPARK-57735 / #56842 adds nanosecond-timestamp support to the default in-memory cache (DefaultCachedBatchSerializer), which is the prerequisite -- once that lands, the Arrow cache can route these types through the same statistics machinery rather than rejecting them here.

[sunchao]

[P2] The prerequisite cited here has now landed in this exact tree: DefaultCachedBatchSerializer has nanosecond timestamp builders in ColumnBuilder.scala, matching ColumnType cases and TimestampNanosColumnStats. This predicate still rejects both TimestampNTZNanosType and TimestampLTZNanosType, so the Arrow cache is now less capable than the default cache even though the guide says it covers every default-supported type. Since #56842 was the stated gate, please reconcile this before merging by completing Arrow-cache support or narrowing the compatibility claim.

[ 🤖 posted by Codex on behalf of sunchao using the code-review-for-me skill 🤖 ]

[viirya]

Done -- with the prerequisite merged (and now in this tree after rebase), this PR adds full Arrow-cache support for TimestampNTZNanosType/TimestampLTZNanosType: the capability predicate accepts them, statistics route to the same TimestampNanosColumnStats the default cache uses (min/max bounds, so cached nanos timestamps prune), the columnar-input stats path collects nanos bounds from the Arrow vectors, and reads work through both the vectorized path (ArrowColumnVector handles TimeStampNanoVector/TimeStampNanoTZVector natively) and the row path (via the columnar-to-row fallback). The write path was already in place via ArrowWriter's nanos writers, including their structured epoch-nanos overflow error. The type-coverage machinery now includes both types, so they are exercised by the cache-and-read-back tests under both reader modes.

[sunchao]

[P2] The follow-up support is still not parity with the default cache for valid nanos-timestamp values outside Arrow's signed-INT64 epoch-nanos window. Spark defines both TimestampNTZNanosType and TimestampLTZNanosType over years 0001-9999, and the default cache stores the (epochMicros, nanosWithinMicro) pair losslessly. This serializer instead routes non-Arrow input through ArrowWriter, which packs the value into one Long. I reproduced this with LocalDateTime.of(9999, 12, 31, 23, 59, 59, 999999999): the default serializer materializes the row, while the Arrow serializer fails with DATETIME_OVERFLOW because only roughly 1677-2262 is representable. There is no fallback after selecting this serializer, and the guide omits this restriction while claiming coverage of every type supported by the default cache. Please either use a lossless cache representation or document the reduced domain beside CalendarIntervalType and narrow the parity claim.

[viirya]

The prerequisite is now built: SPARK-57975 / #57053 (approved, pending CI) adds an opt-in lossless Arrow representation for the nanosecond timestamp types -- a struct of (epochMicros: int64, nanosWithinMicro: int16), i.e. TimestampNanosVal's own layout with no unit conversion -- covering the full 0001-9999 domain, including the 9999-12-31T23:59:59.999999999 value from your repro. The interchange mapping stays untouched: its consumers (pandas datetime64[ns] etc.) are themselves int64-bound, so the reduced domain there is inherent to the destination; the struct is for internal storage only.

Once that lands, this PR will opt in via losslessTimestampNanos = true at the cache's schema-construction sites. The struct is self-describing (tagged child metadata), so both read paths already work through ArrowColumnVector with no format flag; the stats collector will read the struct children directly (which also removes the epoch-nanos conversion from the stats path); and the zero-copy eligibility check will route native int64-nanos Arrow input through the row-conversion path, like the existing large-var-types handling. That restores full value-domain parity with the default cache for these types.

[viirya]

Done. The cache now stores nanosecond timestamps in SPARK-57975's lossless struct encoding, so the full 0001-9999 domain round-trips -- including the 9999-12-31T23:59:59.999999999 value from your repro -- restoring value-domain parity with the default serializer. Wiring details: the stats collector reads the struct components directly and compares with TimestampNanosVal ordering (which also removes the epoch-nanos conversion, and with it the precision-truncation asymmetry MaxGekk flagged); the zero-copy input eligibility check now routes interchange-shaped int64-nanos vectors (e.g. a Python UDF output feeding a cache) through the row-conversion path, like the existing large-var-types handling; both reader modes are covered by the existing type round-trip machinery. The docs' value-range note is updated accordingly.

[sunchao]

[P2] This describes the cache's memory model incorrectly. The durable ArrowCachedBatch payload is an Array[Byte], and the default Dataset.cache() storage level is deserialized MEMORY_AND_DISK, so those cached bytes live on JVM heap. Arrow allocators back only the transient encode/decode roots. Please distinguish long-lived heap cache storage from transient off-heap vectors; otherwise the executor.memoryOverhead and monitoring advice below can cause users to under-size heap and misdiagnose cache OOMs.

[viirya]

Fixed. Rewrote the Memory Management section: the durable cached payload is a heap Array[Byte] (the default Dataset.cache() level is the deserialized MEMORY_AND_DISK), and Arrow's off-heap allocators back only the transient encode/decode roots. The sizing guidance now points at executor heap for cache capacity and clarifies that spark.executor.memoryOverhead covers only the per-batch transient buffers, not the total cache size.

[sunchao]

[P2] Please do not recommend LZ4 as the faster/balanced choice with the dependencies used here. The benchmark source disables every LZ4 case because Arrow falls back to the Commons Compress implementation and measured roughly 50x slower than zstd. Following this troubleshooting advice can therefore make the reported slowdown much worse. Either provide and benchmark the fast LZ4 dependency, document that requirement, or recommend a codec supported by the committed measurements.

[viirya]

Fixed. The troubleshooting section no longer recommends lz4 as faster/balanced. It now suggests lowering the zstd level or using none for read-heavy workloads, with an explicit note that lz4 should be avoided unless the native LZ4 library is on the classpath, since otherwise Arrow falls back to the much slower pure-Java Commons Compress implementation.

[sunchao]

[P2] Please surface unsupported schemas with the documented user-facing condition rather than INTERNAL_ERROR. checkSupportedSchema explicitly handles unsupported types such as ObjectType, but SparkException.internalError sets condition INTERNAL_ERROR, while the guide promises UNSUPPORTED_DATATYPE. A focused invocation confirms getCondition == INTERNAL_ERROR, and the current test asserts only message text. Please use the existing structured unsupported-datatype error (or an equivalent structured error preserving column context) and assert the condition.

[ 🤖 posted by Codex on behalf of sunchao using the code-review-for-me skill 🤖 ]

[viirya]

Fixed. checkSupportedSchema now throws the structured UNSUPPORTED_DATATYPE error (via ExecutionErrors.unsupportedDataTypeError), which is also the condition toArrowSchema raises, so callers see one condition for the capability regardless of which layer detects it first. The test now asserts getCondition == "UNSUPPORTED_DATATYPE" instead of message text.

[sunchao]

[P2] This serializes only a schema-less IPC RecordBatch message, not the IPC stream claimed by ArrowCachedBatch, this serializer's class comment, and the guide. I verified the exact payload starts with MessageHeader.RecordBatch (3), and Arrow 19's standard ArrowStreamReader rejects it with Expected schema but header was 3. Spark's paired reader works only because it reconstructs the schema out of band and calls deserializeRecordBatch directly. Please either emit a complete stream (Schema, batch, and end marker) or describe this narrowly as an internal schema-less RecordBatch payload and remove the ecosystem data-sharing/interoperability claim.

[ 🤖 posted by Codex on behalf of sunchao using the code-review-for-me skill 🤖 ]

[viirya]

Fixed by describing it narrowly, as you suggested. The payload is intentionally a schema-less encapsulated RecordBatch message: the schema is constant for the whole cached relation and is reconstructed from the relation's attributes on read, so repeating Schema/EOS framing in every cached batch would only add bytes. The ArrowCachedBatch/serializer comments and the doc now say exactly that and no longer claim IPC-stream format or ecosystem interoperability.

[sunchao]

[P2] An interrupt delivered while this await is blocked is still lost. InterruptedException clears the status before entering this catch, so Thread.interrupted() observes false, and wasInterrupted records only the entry state. I reproduced this against the exact method by interrupting the caller after awaitTermination began; it returned with the interrupt flag clear. The current test pre-interrupts before entry, so it cannot cover this path. Please accumulate every caught interruption in a mutable flag, finish joining/closing, and restore that flag in finally.

[ 🤖 posted by Codex on behalf of sunchao using the code-review-for-me skill 🤖 ]

[viirya]

Fixed. Every caught InterruptedException (both in the awaitTermination loop and around future.get) now records into a mutable flag that finally restores, so an interrupt delivered while blocked is no longer lost -- you're right that the throw itself clears the status, so the old Thread.interrupted() in the catch observed nothing. Added a test that interrupts the draining thread while the worker is gated (so the interrupt necessarily lands during the drain) and asserts the flag is restored and the produced root closed.

[sunchao]

[P2] This is still not fixed, and Arrow 19 makes the native-library qualification inaccurate. The primary codec list still calls LZ4 fast, but Arrow 19's Lz4CompressionCodec unconditionally uses Commons Compress's FramedLZ4Compressor streams; it does not detect or switch to lz4-java, and this PR instantiates that codec directly. Spark already carries at.yawk.lz4, which does not change this path. Please remove the speed claim and the unless the native LZ4 library is on the classpath guidance (including the benchmark-source note), or replace and benchmark the codec implementation that actually provides a fast path.

[ 🤖 posted by Codex on behalf of sunchao using the code-review-for-me skill 🤖 ]

[viirya]

Fixed, thanks for checking the Arrow 19 source. All the LZ4 guidance (codec list, troubleshooting note, best practices, and the benchmark-source comment) now states that Arrow's Java LZ4 codec unconditionally uses the pure-Java Commons Compress framed LZ4 streams -- no native-backed path, no classpath switch -- and is much slower than zstd, so lz4 is simply not recommended.

[MaxGekk]

0 blocking, 6 non-blocking, 3 nits.
Large, high-quality feature PR already under deep review by @sunchao (3 rounds); the head commit addresses the substantive resource-lifecycle and error-condition items. My independent pass found no net-new blocking defect beyond @sunchao's open (correctly-scoped) stat threads. A focused deep-dive on the new temporal types (TIME(p), TIMESTAMP_LTZ/NTZ(q)) plus a local suite run surfaced a few non-blocking items below; deferring to @sunchao's in-flight review for the open stat threads.

Verification

Deep-dived the temporal types. Type mapping, precision-in-metadata round-trip, and the nanos write-side DATETIME_OVERFLOW guard are sound, and the row vs Arrow-vector stat paths agree for temporal columns on slot order and bound type. Traced Cast.castToTimestampNTZNanos -> truncateTimestampNanosToPrecision and makeTimestampNTZNanos: both truncate nanosWithinMicro to the column precision at write, so a p<9 column never stores a sub-precision value and the stat path's upper-bound floor-truncation is a no-op that matches the (non-truncating) read path -- not a wrong-results bug, only an invariant worth a guard test. Ran ArrowCachedBatchSerializerSuite locally: 69 succeeded, 0 failed.

Design / architecture (1)

• ArrowCachedBatchSerializer.scala:432: statistics are collected by two parallel implementations -- the row ColumnStats path and the Arrow-vector collectStatistics/calculateMinMax* family -- that must be hand-kept in sync; several of your open threads (NaN, geo, UDT-unwrap, columnar size) are individual divergences of this pair. A cross-path equivalence test would close the class. Temporal types verified in agreement.

Correctness (2)

• ArrowCachedBatchSerializer.scala:251: the CalendarInterval overflow translation throws INTERNAL_ERROR / SQLSTATE XX000 (observed in the suite run). Same anti-pattern you fixed for checkSupportedSchema (-> UNSUPPORTED_DATATYPE); DATETIME_OVERFLOW exists next door for the analogous nanos overflow. Refines your interval-capability thread on ArrowUtils.scala. See inline.
• ArrowCachedBatchSerializer.scala:1075: the nanos write-side overflow guard runs only on the row/ArrowWriter path; the zero-copy columnar path collects stats directly and never exercises it. Safe (Arrow input already holds valid int64 epoch-nanos) but worth confirming the asymmetry is intended. See inline.

Suggestions (3)

• ArrowCachedBatchSerializer.scala:693: the nanos stat path floor-truncates bounds to precision while the read path does not; safe today only because Cast truncates at write. A guard test pinning "stat upper bound never falls below any read value for a p<9 column" would stop a future non-truncating write path from silently reintroducing a pruning bug. See inline.
• Add a cross-path temporal-stats equivalence test (same TIME(p)/nanos data via row vs Arrow-columnar input, identical bound rows).
• Nested-container element-type gap: array/struct/map tests are thorough on structure but every container holds only Int/String. Untested code-handled cases: a nested temporal element (routed through the columnar-to-row fallback at :179, a different read path than the top-level typed reader), a nested CalendarInterval overflow (hasCalendarInterval recurses at :265 but the test is top-level-only), and a nested nanos DATETIME_OVERFLOW.

Nits (3)

• ArrowCachedBatchSerializer.scala:63: the class-header "Configuration options" list omits maxBytesPerBatch, cache.prefetch.enabled, and compression.zstd.level -- drifted behind the configs the code now reads. See inline.
• ArrowCachedBatchSerializer.scala:1096: the zero-copy "we don't own the vectors" comment reads in tension with the caller freeing the batch at :1066; the code is correct but a one-line clarification (input consumed synchronously before it's freed) would prevent a use-after-free misreading. See inline.
• docs/sql-arrow-cache-format.md:246: "has small per-batch overhead" -> "has a small per-batch overhead". See inline.

[MaxGekk]

The CalendarInterval overflow translation throws SparkException.internalError, which surfaces as INTERNAL_ERROR / SQLSTATE XX000 (I saw this directly in the ArrowCachedBatchSerializerSuite run: [INTERNAL_ERROR] Arrow cache cannot represent a CalendarInterval ... XX000). Caching a valid CalendarInterval value is a user-facing limitation, not an internal invariant violation -- this is the same anti-pattern you fixed for checkSupportedSchema (now UNSUPPORTED_DATATYPE), and a ready-made structured condition exists next door: QueryExecutionErrors.timestampNanosEpochNanosOverflowError uses DATETIME_OVERFLOW for the exact analogous nanos overflow. Suggest a structured user-facing condition here too. (This refines your ArrowUtils.scala interval-capability thread rather than standing separate -- if the overflow translation stays, its condition should be user-facing.)

[viirya]

Fixed. withIntervalOverflowTranslation now throws the structured DATETIME_OVERFLOW condition (added ExecutionErrors.datetimeOverflowError, mirroring QueryExecutionErrors.timestampNanosEpochNanosOverflowError's pattern) instead of SparkException.internalError. Caching a CalendarInterval that can't be losslessly converted is a user-facing limitation, same anti-pattern as the checkSupportedSchema fix. Test now asserts getCondition == "DATETIME_OVERFLOW".

[MaxGekk]

The nanos write-side overflow guard (DATETIME_OVERFLOW via Math.multiplyExact in TimestampNTZNanosWriter/LTZNanosWriter) only runs on the row/ArrowWriter path. This zero-copy path reuses the input Arrow vector and collects stats directly via TimeStampVector.get, so it never exercises that guard. It's safe today because an Arrow-backed source vector already holds valid int64 epoch-nanos, but it means the overflow error is reachable only on non-Arrow input. Worth confirming this asymmetry is intended (and, ideally, a comment noting why the zero-copy path needs no guard).

[viirya]

This is intended, not a gap: the zero-copy path only ever reuses vectors that came from an upstream Arrow source (ArrowColumnVector.getValueVector), and those vectors were already populated by an ArrowWriter (or equivalent) that enforced the guard on write. There's no new conversion happening on this path that could overflow. Added a comment explaining this and flagging that the assumption would need revisiting if a future caller could reach this path with unguarded vectors.

[MaxGekk]

The Arrow-vector stat path floor-truncates the min/max bounds to the column precision here (epochNanosToTimestampNanos(min/max, precision)), while the read path (ArrowColumnVector.decodeEpochNanos) returns the stored value without re-truncating. Floor-truncating the upper bound down would be unsound for pruning if a sub-precision value ever reached the vector. I verified this can't happen on the paths I could find -- CAST(... AS TIMESTAMP_NTZ(p)) truncates at write (truncateTimestampNanosToPrecision) and makeTimestampNTZNanos applies precision -- so it's not a live bug. But the safety rests entirely on write-time truncation while the read/stat paths disagree on precision handling. A guard test (assert the stat upper bound never falls below any read value for a p<9 column) would stop a future non-truncating write path from silently reintroducing a pruning bug.

[viirya]

Confirmed not a live bug for the reason you found -- write-time truncation covers it today. Added a comment spelling out the fragility explicitly, plus a guard test that drives a value through the real truncateTimestampNanosToPrecision write path and asserts the stat's upper bound matches the actual value exactly, so a future write path that stopped truncating would fail this test rather than silently under-reporting the bound.

[MaxGekk]

The class-header "Configuration options" list has drifted behind the code: it omits spark.sql.execution.arrow.maxBytesPerBatch (read at line 87, added for the byte-limit batch split), spark.sql.execution.arrow.cache.prefetch.enabled (line 148/214), and compression.zstd.level (line 90). Worth updating so the header matches the configs actually captured.

[viirya]

Fixed. Added spark.sql.execution.arrow.maxBytesPerBatch, spark.sql.execution.arrow.compression.zstd.level, and spark.sql.execution.arrow.cache.prefetch.enabled to the class header.

[MaxGekk]

This "we don't own the vectors ... owned by the input ColumnarBatch" note now reads in tension with the caller freeing that same batch via batch.closeIfFreeable() at line 1066. The code is correct -- getRecordBatch() serializes the vectors out synchronously before the caller's finally frees the input, and the root is intentionally left unclosed to avoid double-freeing the input's buffers -- but the two comments together no longer tell a coherent ownership story. A one-line clarification that the input is consumed before it's freed would prevent a use-after-free misreading.

[viirya]

Fixed. Clarified that this isn't a use-after-free: serializeBatch/collectStatistics fully materialize their outputs (a copied Array[Byte] and plain stat values) before this method returns, so nothing here still references the input ColumnarBatch's buffers once the caller frees them via closeIfFreeable().

[MaxGekk]

Missing article: "Arrow format has small per-batch overhead" -> "has a small per-batch overhead".

[viirya]

Fixed. "has small" -> "has a small".

[sunchao]

[P2] Please make the compact-decimal fast path endian-safe. Arrow Java stores Decimal128 in native byte order: on big-endian JVMs the first eight bytes are only the sign-extension word, while the unscaled Long is at startIndex + 8. This unconditional read therefore turns positive Decimal(p <= 18) values into zero and negative values into an unscaled -1 whenever the cache uses the row reader (for example, with vectorized cache reads disabled). A DECIMAL(10,2) value of 123.45 consequently round-trips as 0.00 on s390x. The previous DecimalVector.getObject path and Spark's default cache are endian-safe. Please use the vector accessor or choose the low-word offset based on native byte order, and cover the big-endian layout in a regression test.

[viirya]

Fixed. You're right on the layout: DecimalUtility.writeLongToArrowBuf branches on ByteOrder.nativeOrder() and places the unscaled low-order word in the first 8 bytes on little-endian but the last 8 on big-endian (I verified against the Arrow 19 bytecode). The fast-path reader now selects the word via a compactDecimalUnscaledOffset(nativeOrder) helper (startIndex + 0 on LE, + 8 on BE), keeping the zero-allocation path; ArrowBuf.getLong reads in native order, so choosing the right word is sufficient. The in-code comments claiming a little-endian layout were also wrong and are corrected. Coverage: a new test asserts both branches of the offset selection directly (so little-endian CI validates the big-endian branch) and cross-checks the selected offset against where a real DecimalVector stored the word on the running platform; the existing negative/zero compact-decimal round-trip test serves as the on-hardware regression for big-endian JVMs and its comment now documents that role.

[dongjoon-hyun]

Could you mention spark.sql.execution.arrow.maxBytesPerBatch together?

[viirya]

Added, together with spark.sql.execution.arrow.cache.prefetch.enabled which was also missing from this table.

[dongjoon-hyun]

This per-row columnType match reads the value an extra time just to check its runtime type (row.getArray(ordinal) here, then again inside actualSize). On UnsafeRow input each getArray/getMap/getStruct call allocates a fresh wrapper, so this adds per-row overhead to the default serializer's cache-write path for complex types as well. Also, silently skipping the size for columnar views records zero bytes, which underestimates the relation's sizeInBytes and can make it wrongly eligible for broadcast.

Suggest reading the value once and falling back to the type's defaultSize (ARRAY 28, MAP 68, STRUCT 20) as a conservative estimate for columnar views:

Suggested change

	// Check if this is a columnar complex type that doesn't support getSizeInBytes
	val isColumnarComplexType = columnType match {
	case _: ARRAY =>
	row.getArray(ordinal).isInstanceOf[ColumnarArray]
	case _: MAP =>
	row.getMap(ordinal).isInstanceOf[ColumnarMap]
	case struct: STRUCT =>
	row.getStruct(ordinal, struct.dataType.fields.length).isInstanceOf[ColumnarRow]
	case _ =>
	false
	}
	if (!isColumnarComplexType) {
	// Normal path: calculate size for unsafe types
	// (UnsafeArrayData/UnsafeMapData/UnsafeRow)
	val size = columnType.actualSize(row, ordinal)
	sizeInBytes += size
	}
	// else: Skip size calculation for columnar complex types
	// (ColumnarArray/ColumnarMap/ColumnarRow). These are views into ColumnVectors
	// and don't expose getSizeInBytes()
	count += 1
	// Read the value once: columnar complex values (ColumnarArray/ColumnarMap/ColumnarRow)
	// are views into ColumnVectors and do not expose getSizeInBytes, so fall back to the
	// type's default size estimate instead of recording zero bytes for them.
	val size = columnType match {
	case _: ARRAY => row.getArray(ordinal) match {
	case unsafe: UnsafeArrayData => 4 + unsafe.getSizeInBytes
	case _ => columnType.defaultSize
	}
	case _: MAP => row.getMap(ordinal) match {
	case unsafe: UnsafeMapData => 4 + unsafe.getSizeInBytes
	case _ => columnType.defaultSize
	}
	case struct: STRUCT =>
	row.getStruct(ordinal, struct.dataType.fields.length) match {
	case unsafe: UnsafeRow => 4 + unsafe.getSizeInBytes
	case _ => columnType.defaultSize
	}
	case _ => columnType.actualSize(row, ordinal)
	}
	sizeInBytes += size
	count += 1

As a side benefit, this also guards non-Unsafe values like GenericArrayData, where the previous code would have thrown ClassCastException from actualSize.

[dongjoon-hyun]

BTW, I didn't check the line width. So, please consider the logic only, @viirya .

[viirya]

Applied as suggested (reformatted for the 100-char limit). Reading the value once removes the extra per-row wrapper allocation on the default serializer's write path too, the defaultSize fallback replaces the zero-byte skip so columnar views no longer deflate the relation's sizeInBytes, and the Unsafe* match guards GenericArrayData from the ClassCastException in actualSize. Verified with ColumnStatsSuite, InMemoryColumnarQuerySuite, and PartitionBatchPruningSuite in addition to the Arrow suite.

[dongjoon-hyun]

This scans every row per column to count nulls. Arrow already tracks this in the validity buffer: getNullCount computes it with word-at-a-time bit counting (roughly O(n/64), O(1) for vectors with no validity buffer), instead of a per-row isNull call through the vector interface.

Since collectStatistics runs on the write path of both columnar branches — including the zero-copy re-cache path, where this per-row loop is the main remaining per-value work — this directly cuts the cache-build cost the benchmark highlights.

Suggested change

	val nullCount = (0 until rowCount).count(i => vector.isNull(i))
	val nullCount = vector.getNullCount

Semantics are unchanged for every shape this cache produces: validity-buffer-backed vectors count set bits exactly like the isNull loop, NullVector.getNullCount returns valueCount (all rows null, matching isNull = true), and the struct-backed lossless types (nanos timestamps, CalendarInterval) use the struct's own validity buffer, which is what isNull reads too.

[viirya]

Applied. I verified the semantic-equivalence claims: NullVector.getNullCount returns valueCount (all rows null, matching the per-row loop), validity-backed vectors count set bits word-at-a-time, and the struct-backed lossless types count the struct's own validity buffer -- the same buffer isNull reads. Since this was the last remaining per-value pass for non-orderable columns on the zero-copy re-cache path, it directly cuts cache-build cost there.

[viirya]

Hi, @dongjoon-hyun do you still have some more comments? Can you take another look? Thanks

[dongjoon-hyun]

+1, LGTM. Thank you, @viirya .

[viirya]

Merge Summary:

• merged into master 2347165
• merged into branch-4.x 7e43efe

Posted by merge_spark_pr.py

[viirya]

Thanks all!