[SPARK-58381][SQL] Gate the Arrow cache zero-copy write path on physical congruence with the cache schema

sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala

@@ -1004,22 +1004,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) {
@@ -1040,9 +1042,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])
@@ -1110,35 +1112,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.
  */

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, SparkUnsupportedOperationException}
+import org.apache.spark.{SparkConf, 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}
@@ -730,6 +732,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