You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by vi...@apache.org on 2026/07/28 23:15:56 UTC
(spark) branch branch-4.x updated: [SPARK-58381][SQL] Gate the Arrow cache zero-copy write path on physical congruence with the cache schema
This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new dbea368bbc6b [SPARK-58381][SQL] Gate the Arrow cache zero-copy write path on physical congruence with the cache schema
dbea368bbc6b is described below
commit dbea368bbc6b9400fcec70a07c9a8584a799ad6b
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Tue Jul 28 16:15:36 2026 -0700
[SPARK-58381][SQL] Gate the Arrow cache zero-copy write path on physical congruence with the cache schema
### What changes were proposed in this pull request?
Replace the blocklist that decides zero-copy eligibility on the Arrow cache write path (`ColumnarBatchToArrowCachedBatchIterator`) with a per-column physical-congruence check against the cache schema, using `ArrowUtils.isCompatibleWithDeclaredField` (introduced by SPARK-58258 for the columnar Python UDF pass-through). Input vectors whose field trees are not physically congruent with the corresponding cache schema field now take the row-based conversion path, which rewrites the values [...]
### Why are the changes needed?
The zero-copy path serializes the input vectors' buffers verbatim under the cache's own schema: `serializeBatch` writes only the record batch, and the read path reconstructs the schema from `cacheSchema` and loads the buffers positionally into it. The blocklist covered only large var-width vectors and interchange-shaped nanosecond timestamp / `CalendarInterval` vectors, so every other physical divergence was forwarded verbatim and silently reinterpreted under the canonical layout when [...]
The corruption is concrete: an Arrow map vector whose entry-struct children sit in `[value, key]` order is tolerated by Arrow and read correctly by `ArrowColumnVector` (key/value are addressed by name), but its buffers are positional, so a cached `{123: 7}` comes back as `{7: 123}` -- silently, and for the lifetime of the cached relation. The new test demonstrates this: it fails on master with `Array((7, 123)) did not equal Array((123, 7))` and passes with the gate.
In-tree columnar producers happened to be covered by the blocklist (the cache scan itself produces canonical vectors; Python UDF output's large var-width and interchange nanos/interval shapes were the blocked entries), so this is reachable today only from external Arrow-backed DSv2 columnar sources -- `supportsColumnarInput` accepts any of them. The congruence check closes the whole class rather than the enumerated instances, and keeps the write gate symmetric with the read-side gate [...]
A false reject is safe by construction: the row-based fallback re-encodes values under the cache schema (with its value-domain guards), so eligible batches lose only the zero-copy optimization, never correctness.
### Does this PR introduce _any_ user-facing change?
No. The Arrow cache serializer (SPARK-57268) is unreleased.
### How was this patch tested?
New test in `ArrowCachedBatchSerializerSuite` that builds a spec-tolerated but non-canonical map vector (entry children `[value, key]`), verifies the source reads correctly, round-trips it through the serializer's columnar write and read paths, and asserts the values survive. Without the fix it fails with the keys and values swapped. Full `ArrowCachedBatchSerializerSuite` passes (72 tests), confirming canonical shapes remain zero-copy eligible.
### Was this patch authored or co-authored using generative AI tooling?
Yes, this pull request and its description were written by Claude Code.
Closes #57577 from viirya/arrow-cache-write-zerocopy-gate.
Authored-by: Liang-Chi Hsieh <vi...@gmail.com>
Signed-off-by: Liang-Chi Hsieh <vi...@gmail.com>
(cherry picked from commit ece8eacb8f44742395f39e67a0714eca2efc4a3e)
Signed-off-by: Liang-Chi Hsieh <vi...@gmail.com>
---
.../columnar/ArrowCachedBatchSerializer.scala | 67 +++++++---------------
.../columnar/ArrowCachedBatchSerializerSuite.scala | 61 +++++++++++++++++++-
2 files changed, 80 insertions(+), 48 deletions(-)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala
index 58b385071e77..f90f1a9dcd22 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala
@@ -1017,22 +1017,24 @@ private class ColumnarBatchToArrowCachedBatchIterator(
Utils.tryWithSafeFinally {
val rowCount = batch.numRows()
- // Check if batch is already Arrow-based for zero-copy path. The zero-copy path reuses the
- // input vectors but serializes them under the cache's own schema, and the read path
- // reconstructs that same schema, so the input vectors' physical shape must match it:
- // - the cache schema is built with largeVarTypes=false, so large var-width vectors
- // (64-bit offsets) would be silently corrupted when read back under 32-bit offsets;
- // - the cache schema is built with losslessInternalTypes=true, so nanosecond timestamps
- // and CalendarInterval are lossless structs, and interchange-shaped input vectors
- // (TimeStampNano(TZ)Vector, IntervalMonthDayNanoVector, e.g. from a Python UDF output)
- // would not match the reconstructed struct schema.
- // Fall back to the row-based conversion (which rewrites through ArrowWriter under the
- // cache schema) whenever any input vector is, or nests, such a mismatched shape.
+ // Check if batch is already Arrow-based for zero-copy path. The zero-copy path serializes
+ // the input vectors' buffers verbatim under the cache's own schema (serializeBatch writes
+ // only the record batch; the read path reconstructs the schema from cacheSchema and loads
+ // the buffers positionally into it), so each input vector's field tree must be physically
+ // congruent with the corresponding cache schema field. Any divergence -- a var-width or
+ // list offset width disagreeing with the canonical one, view or dictionary encodings, an
+ // interchange-shaped nanosecond timestamp or CalendarInterval vector where the cache
+ // schema has the lossless structs (losslessInternalTypes=true), a tagged struct carrying
+ // extra children, map entry children in the wrong order -- would be silently reinterpreted
+ // under the canonical layout when the cached batch is read back. Incongruent input takes
+ // the row-based conversion instead, which rewrites the values through ArrowWriter under
+ // the cache schema.
+ val declaredFields = arrowSchema.getFields
val vectors = (0 until batch.numCols()).map(batch.column)
- val zeroCopyEligible = vectors.forall {
- case acv: ArrowColumnVector =>
- !ColumnarBatchToArrowCachedBatchIterator.containsCacheSchemaMismatch(
- acv.getValueVector)
+ val zeroCopyEligible = vectors.zipWithIndex.forall {
+ case (acv: ArrowColumnVector, i) =>
+ ArrowUtils.isCompatibleWithDeclaredField(
+ acv.getValueVector.getField, declaredFields.get(i))
case _ => false
}
if (zeroCopyEligible) {
@@ -1053,9 +1055,9 @@ private class ColumnarBatchToArrowCachedBatchIterator(
schema: Seq[Attribute],
vectors: Seq[ColumnVector]): ArrowCachedBatch = {
// Zero-copy path: extract Arrow vectors directly from ArrowColumnVector. Vectors reaching
- // this path have already passed containsCacheSchemaMismatch, so nanosecond timestamp and
- // CalendarInterval columns are in the lossless struct shape matching the cache schema; no
- // value conversion happens here, so no overflow is possible.
+ // this path are physically congruent with the cache schema (isCompatibleWithDeclaredField),
+ // so nanosecond timestamp and CalendarInterval columns are in the lossless struct shape
+ // matching it; no value conversion happens here, so no overflow is possible.
val arrowVectors = vectors.map(
_.asInstanceOf[ArrowColumnVector].getValueVector.asInstanceOf[
org.apache.arrow.vector.FieldVector])
@@ -1123,35 +1125,6 @@ private class ColumnarBatchToArrowCachedBatchIterator(
}
}
-private object ColumnarBatchToArrowCachedBatchIterator {
- import org.apache.arrow.vector.{FieldVector, LargeVarBinaryVector, LargeVarCharVector}
-
- /**
- * Whether the vector is, or nests, a large var-width vector (64-bit offsets). These are not
- * eligible for the zero-copy path because that path serializes and reloads under a schema built
- * with largeVarTypes=false; reinterpreting 64-bit offset buffers as 32-bit would corrupt data.
- */
- /**
- * Whether the vector tree contains any shape the cache schema cannot serialize as-is: large
- * var-width vectors (the cache schema uses 32-bit offsets) or interchange-shaped nanosecond
- * timestamp / CalendarInterval vectors (the cache schema uses the lossless struct
- * representations from losslessInternalTypes=true). Such input must take the row-conversion
- * path, which rewrites values through ArrowWriter under the cache schema. The lossless struct
- * vectors themselves (e.g. from re-caching a projection of a cached relation) match the cache
- * schema and stay zero-copy eligible.
- */
- def containsCacheSchemaMismatch(
- vector: org.apache.arrow.vector.ValueVector): Boolean = vector match {
- case _: LargeVarCharVector | _: LargeVarBinaryVector => true
- case _: org.apache.arrow.vector.TimeStampNanoVector |
- _: org.apache.arrow.vector.TimeStampNanoTZVector |
- _: org.apache.arrow.vector.IntervalMonthDayNanoVector => true
- case fv: FieldVector =>
- fv.getChildrenFromFields.asScala.exists(containsCacheSchemaMismatch)
- case _ => false
- }
-}
-
/**
* Iterator that converts ArrowCachedBatch to ColumnarBatch.
*/
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala
index f4119a306ce1..621fc4d48130 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala
@@ -25,8 +25,10 @@ import org.apache.arrow.vector.{
Float4Vector, Float8Vector, IntVector, LargeVarCharVector, SmallIntVector,
TimeNanoVector, TimeStampMicroTZVector, TimeStampMicroVector, TinyIntVector,
VarBinaryVector, VarCharVector, VectorSchemaRoot, VectorUnloader}
+import org.apache.arrow.vector.complex.{MapVector, StructVector}
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
-import org.apache.spark.{SparkConf, SparkException, SparkUnsupportedOperationException}
+import org.apache.spark.{SparkConf, SparkException, SparkUnsupportedOperationException, TaskContext}
import org.apache.spark.sql.{QueryTest, Row}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{AttributeReference, GenericInternalRow}
@@ -731,6 +733,63 @@ class ArrowCachedBatchSerializerSuite extends QueryTest with SharedSparkSession
}
}
+ test("zero-copy cache write requires physical congruence with the cache schema") {
+ // An Arrow map vector whose entry-struct children sit in [value, key] order: Arrow tolerates
+ // the layout and ArrowColumnVector reads it correctly (key/value are addressed by name), but
+ // record-batch buffers are laid out positionally, so serializing them verbatim under the
+ // cache's canonical [key, value] schema would silently swap every key with its value when the
+ // cached batch is read back. Such input must take the row-based conversion path, which
+ // rewrites the values through ArrowWriter under the cache schema.
+ val serializer = new ArrowCachedBatchSerializer
+ val attrs = Seq(AttributeReference("m", MapType(IntegerType, IntegerType, false))())
+ val input = spark.sparkContext.parallelize(Seq(0), 1).mapPartitions { _ =>
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator(
+ "swapped-map-input", 0, Long.MaxValue)
+ Option(TaskContext.get()).foreach(
+ _.addTaskCompletionListener[Unit](_ => allocator.close()))
+ val intType = new ArrowType.Int(32, true)
+ val entriesField = new Field(
+ "entries",
+ FieldType.notNullable(ArrowType.Struct.INSTANCE),
+ java.util.Arrays.asList(
+ new Field("value", FieldType.notNullable(intType), null),
+ new Field("key", FieldType.notNullable(intType), null)))
+ val mapField = new Field(
+ "m",
+ FieldType.nullable(new ArrowType.Map(false)),
+ java.util.Collections.singletonList(entriesField))
+ val mapVector = mapField.createVector(allocator).asInstanceOf[MapVector]
+ mapVector.allocateNew()
+ val entries = mapVector.getDataVector.asInstanceOf[StructVector]
+ entries.getChild("key").asInstanceOf[IntVector].setSafe(0, 123)
+ entries.getChild("value").asInstanceOf[IntVector].setSafe(0, 7)
+ entries.setIndexDefined(0)
+ mapVector.startNewValue(0)
+ mapVector.endValue(0, 1)
+ mapVector.setValueCount(1)
+ val column = new ArrowColumnVector(mapVector)
+ // The source itself reads correctly; only the cache round trip is under test.
+ val sourceMap = column.getMap(0)
+ assert(sourceMap.keyArray.getInt(0) == 123 && sourceMap.valueArray.getInt(0) == 7)
+ Iterator(new ColumnarBatch(Array[ColumnVector](column), 1))
+ }
+ val conf = spark.sessionState.conf
+ val cached = serializer.convertColumnarBatchToCachedBatch(
+ input, attrs, StorageLevel.MEMORY_ONLY, conf)
+ val roundTripped = serializer
+ .convertCachedBatchToColumnarBatch(cached, attrs, attrs, conf)
+ .mapPartitions { batches =>
+ batches.flatMap { batch =>
+ (0 until batch.numRows()).map { i =>
+ val m = batch.column(0).getMap(i)
+ (m.keyArray.getInt(0), m.valueArray.getInt(0))
+ }
+ }
+ }
+ .collect()
+ assert(roundTripped === Array((123, 7)))
+ }
+
test("columnar input with empty arrays and maps from parquet") {
withTempPath { dir =>
val path = dir.getAbsolutePath
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org