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