perf: reuse the nested unsafe views in ArrowWriter's row path, detaching them per batch

pom.xml

@@ -1020,6 +1020,7 @@ under the License.
                     <arg>-Wconf:cat=unused-params&amp;site=org\.apache\.comet\.serde\.CometAggregateExpressionSerde\.supportsSparkPartialToNativeFinal\..*:s</arg>
                     <arg>-Wconf:cat=unused-params&amp;site=org\.apache\.comet\.serde\.CometAggregateExpressionSerde\.supportsNativePartialToSparkFinal\..*:s</arg>
                     <arg>-Wconf:cat=unused-params&amp;site=org\.apache\.spark\.sql\.comet\.CometScanWithPlanData\.withDynamicPruningFilters\..*:s</arg>
+                    <arg>-Wconf:cat=unused-params&amp;site=org\.apache\.spark\.sql\.comet\.execution\.arrow\.ArrowFieldWriter\.collectViews\..*:s</arg>
                     <!-- Extractors and public UDF/aggregate hooks retain caller-required arguments. -->
                     <arg>-Wconf:cat=unused-params&amp;site=org\.apache\.comet\.CometExecIterator\.InputFailure\.unapply\.nativeFailure:s</arg>
                     <arg>-Wconf:cat=unused-params&amp;site=org\.apache\.comet\.udf\.CometUDF\.evaluate\.(partitionIndex|planId):s</arg>

spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala

@@ -22,6 +22,7 @@ package org.apache.spark.sql.comet.execution.arrow
 import java.math.{BigDecimal => JavaBigDecimal, BigInteger}
 import java.nio.ByteOrder
 
+import scala.collection.mutable.ArrayBuffer
 import scala.jdk.CollectionConverters._
 
 import org.apache.arrow.memory.{ArrowBuf, BufferAllocator}
@@ -149,6 +150,11 @@ class ArrowWriter(val root: VectorSchemaRoot, fields: Array[ArrowFieldWriter]) {
 
   private var count: Int = 0
 
+  // The views of unsafe arrays, maps and structs that the writers at any depth keep from one value
+  // to the next. Each holds the last row it read until the batch is finished or the writer reset,
+  // when all are pointed away from their rows: a detach per row would cost what the reuse saves.
+  private val views = UnsafeViews.of(fields)
+
   def write(row: InternalRow): Unit = {
     var i = 0
     row match {
@@ -204,12 +210,14 @@ class ArrowWriter(val root: VectorSchemaRoot, fields: Array[ArrowFieldWriter]) {
   def finish(): Unit = {
     root.setRowCount(count)
     fields.foreach(_.finish())
+    views.detach()
   }
 
   def reset(): Unit = {
     root.setRowCount(0)
     count = 0
     fields.foreach(_.reset())
+    views.detach()
   }
 }
 
@@ -636,6 +644,58 @@ private[arrow] object ArrowFieldWriter {
     }
 }
 
+/**
+ * The views of unsafe arrays, maps and structs that the writers under an [[ArrowWriter]] keep
+ * from one value to the next, rather than allocating a view per value. Between values a view
+ * holds the last row it read, and the memory behind it: at most one row per nested writer.
+ * [[detach]] points each at an empty value, which [[ArrowWriter]] does once a batch is finished
+ * or reset.
+ */
+private[arrow] final class UnsafeViews(
+    arrays: Array[UnsafeArrayData],
+    maps: Array[UnsafeMapData],
+    rows: Array[UnsafeRow]) {
+  import UnsafeViews._
+
+  def detach(): Unit = {
+    var i = 0
+    while (i < arrays.length) {
+      arrays(i).pointTo(EmptyArray, Platform.LONG_ARRAY_OFFSET.toLong, 8)
+      i += 1
+    }
+    i = 0
+    while (i < maps.length) {
+      maps(i).pointTo(EmptyMap, Platform.LONG_ARRAY_OFFSET.toLong, 24)
+      i += 1
+    }
+    i = 0
+    while (i < rows.length) {
+      rows(i).pointTo(null, 0L, 0)
+      i += 1
+    }
+  }
+}
+
+private[arrow] object UnsafeViews {
+  // An array of no elements and a map of two, for the views to point at between rows. A null base
+  // would not do for either, as `pointTo` reads a header there.
+  private val EmptyArray = Array[Long](0L)
+  private val EmptyMap = Array[Long](8L, 0L, 0L)
+
+  /** Collects the views of the writers under `fields`. */
+  final class Builder {
+    val arrays: ArrayBuffer[UnsafeArrayData] = ArrayBuffer.empty
+    val maps: ArrayBuffer[UnsafeMapData] = ArrayBuffer.empty
+    val rows: ArrayBuffer[UnsafeRow] = ArrayBuffer.empty
+  }
+
+  def of(fields: Array[ArrowFieldWriter]): UnsafeViews = {
+    val builder = new Builder
+    fields.foreach(_.collectViews(builder))
+    new UnsafeViews(builder.arrays.toArray, builder.maps.toArray, builder.rows.toArray)
+  }
+}
+
 /**
  * For a dictionary-encoded input, one more than the output index of the first row written for
  * each dictionary id, or zero. An output row holds an id's bytes only while the input, and with
@@ -756,6 +816,12 @@ private[arrow] abstract class ArrowFieldWriter {
   /** Called before the rows of each input batch, to drop anything cached about the last one. */
   private[arrow] def startInputBatch(): Unit = {}
 
+  /**
+   * Adds to `views` the view of Spark's unsafe data this writer keeps from one value to the next,
+   * if any, and those of the writers under it, for [[ArrowWriter]] to detach after each batch.
+   */
+  private[arrow] def collectViews(views: UnsafeViews.Builder): Unit = {}
+
   /**
    * Whether [[maskNulls]] can null out this field's values under null parents. A list or map
    * would need its offsets rewritten as well.
@@ -1625,6 +1691,9 @@ private[arrow] class TimeNanoWriter(val valueVector: TimeNanoVector)
 private[arrow] class ArrayWriter(val valueVector: ListVector, val elementWriter: ArrowFieldWriter)
     extends ArrowFieldWriter {
 
+  // Pointed at each unsafe array in turn rather than allocating a view per value; see UnsafeViews.
+  private val elements = new UnsafeArrayData
+
   override def setNull(): Unit = {}
 
   override def setValue(input: SpecializedGetters, ordinal: Int): Unit = {
@@ -1643,13 +1712,11 @@ private[arrow] class ArrayWriter(val valueVector: ListVector, val elementWriter:
     }
   }
 
-  // Each unsafe array gets a view of its own from Spark's getter. A view reused across values
-  // would keep the last row it read reachable after that row's values were copied.
   override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal: Int): Unit = {
     if (row.isNullAt(ordinal)) {
       setNull()
     } else {
-      writeElements(row.getArray(ordinal))
+      writeUnsafe(row.getBaseObject, row.getBaseOffset, row.getLong(ordinal))
     }
     count += 1
   }
@@ -1661,19 +1728,30 @@ private[arrow] class ArrayWriter(val valueVector: ListVector, val elementWriter:
       if (array.isNullAt(i)) {
         setNull()
       } else {
-        writeElements(array.getArray(i))
+        writeUnsafe(array.getBaseObject, array.getBaseOffset, array.getLong(i))
       }
       count += 1
       i += 1
     }
   }
 
+  /** Sets value `count` to the unsafe array that `offsetAndSize` locates from `baseOffset`. */
+  private def writeUnsafe(base: AnyRef, baseOffset: Long, offsetAndSize: Long): Unit = {
+    elements.pointTo(base, baseOffset + (offsetAndSize >> 32), offsetAndSize.toInt)
+    writeElements(elements)
+  }
+
   private def writeElements(array: UnsafeArrayData): Unit = {
     valueVector.startNewValue(count)
     elementWriter.writeArrayElements(array)
     valueVector.endValue(count, array.numElements())
   }
 
+  override private[arrow] def collectViews(views: UnsafeViews.Builder): Unit = {
+    views.arrays += elements
+    elementWriter.collectViews(views)
+  }
+
   // Spark's vectors keep every array's elements in one child vector. The offsets are written in
   // one pass and the elements one run of adjacent rows at a time.
   override def writeColumnSlice(input: ColumnVector, startRow: Int, numRows: Int): Unit = {
@@ -1715,6 +1793,9 @@ private[arrow] class StructWriter(
     children: Array[ArrowFieldWriter])
     extends ArrowFieldWriter {
 
+  // Pointed at each unsafe struct in turn, as in ArrayWriter.
+  private val fields = new UnsafeRow(children.length)
+
   override def setNull(): Unit = {
     var i = 0
     while (i < children.length) {
@@ -1739,12 +1820,11 @@ private[arrow] class StructWriter(
     }
   }
 
-  // Each unsafe struct gets a view of its own, as in ArrayWriter.
   override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal: Int): Unit = {
     if (row.isNullAt(ordinal)) {
       setNull()
     } else {
-      writeFields(row.getStruct(ordinal, children.length))
+      writeUnsafe(row.getBaseObject, row.getBaseOffset, row.getLong(ordinal))
     }
     count += 1
   }
@@ -1756,13 +1836,19 @@ private[arrow] class StructWriter(
       if (array.isNullAt(i)) {
         setNull()
       } else {
-        writeFields(array.getStruct(i, children.length))
+        writeUnsafe(array.getBaseObject, array.getBaseOffset, array.getLong(i))
       }
       count += 1
       i += 1
     }
   }
 
+  /** Sets value `count` to the unsafe struct that `offsetAndSize` locates from `baseOffset`. */
+  private def writeUnsafe(base: AnyRef, baseOffset: Long, offsetAndSize: Long): Unit = {
+    fields.pointTo(base, baseOffset + (offsetAndSize >> 32), offsetAndSize.toInt)
+    writeFields(fields)
+  }
+
   private def writeFields(struct: UnsafeRow): Unit = {
     valueVector.setIndexDefined(count)
     var i = 0
@@ -1772,6 +1858,11 @@ private[arrow] class StructWriter(
     }
   }
 
+  override private[arrow] def collectViews(views: UnsafeViews.Builder): Unit = {
+    views.rows += fields
+    children.foreach(_.collectViews(views))
+  }
+
   private val childrenSupportNullMask = children.forall(_.supportsNullMask)
 
   // Writes each field as a column. Under a null struct the fields must come out null, as setNull
@@ -1835,6 +1926,9 @@ private[arrow] class MapWriter(
     val valueWriter: ArrowFieldWriter)
     extends ArrowFieldWriter {
 
+  // Pointed at each unsafe map in turn, as in ArrayWriter.
+  private val entries = new UnsafeMapData
+
   override def setNull(): Unit = {}
 
   override def setValue(input: SpecializedGetters, ordinal: Int): Unit = {
@@ -1857,12 +1951,11 @@ private[arrow] class MapWriter(
     }
   }
 
-  // Each unsafe map gets a view of its own, as in ArrayWriter.
   override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal: Int): Unit = {
     if (row.isNullAt(ordinal)) {
       setNull()
     } else {
-      writeEntries(row.getMap(ordinal))
+      writeUnsafe(row.getBaseObject, row.getBaseOffset, row.getLong(ordinal))
     }
     count += 1
   }
@@ -1874,13 +1967,25 @@ private[arrow] class MapWriter(
       if (array.isNullAt(i)) {
         setNull()
       } else {
-        writeEntries(array.getMap(i))
+        writeUnsafe(array.getBaseObject, array.getBaseOffset, array.getLong(i))
       }
       count += 1
       i += 1
     }
   }
 
+  /** Sets value `count` to the unsafe map that `offsetAndSize` locates from `baseOffset`. */
+  private def writeUnsafe(base: AnyRef, baseOffset: Long, offsetAndSize: Long): Unit = {
+    entries.pointTo(base, baseOffset + (offsetAndSize >> 32), offsetAndSize.toInt)
+    writeEntries(entries)
+  }
+
+  override private[arrow] def collectViews(views: UnsafeViews.Builder): Unit = {
+    views.maps += entries
+    keyWriter.collectViews(views)
+    valueWriter.collectViews(views)
+  }
+
   private def writeEntries(map: UnsafeMapData): Unit = {
     val numElements = map.numElements()
     valueVector.startNewValue(count)

spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala

@@ -886,38 +886,55 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers {
     false
   }
 
-  test("the row path keeps no reference to an unsafe row once it is written") {
-    // Arrays, maps and structs at the top level, inside one another and as array elements. A view
-    // reused across values would keep the last row it read, and the row's memory, reachable.
+  test("the row path keeps no reference to an unsafe row once its batch is finished") {
+    // Arrays, maps and structs at the top level, inside one another and as array elements. The
+    // writers point a reused view at each value, so between rows a view holds the last row it
+    // read: at most one row per nested writer. Finishing or resetting the writer must point every
+    // view away, or those rows, and the memory behind them, would stay reachable. With nulls, the
+    // views of different writers last read different rows.
     val schema = new StructType()
       .add("a", ArrayType(ArrayType(IntegerType)))
       .add("s", new StructType().add("x", StringType).add("y", ArrayType(LongType)))
       .add("m", MapType(StringType, new StructType().add("z", BinaryType)))
       .add("as", ArrayType(new StructType().add("w", IntegerType)))
       .add("am", ArrayType(MapType(IntegerType, StringType)))
     val n = 20
-    val rnd = new Random(7)
-    val vectors = schema.fields.map(f => newVector(n, f.dataType, offHeap = false))
-    val allocator = new RootAllocator(Long.MaxValue)
-    val root = VectorSchemaRoot.create(Utils.toArrowSchema(schema, "UTC"), allocator)
-    try {
-      schema.fields.zip(vectors).foreach { case (field, v) =>
-        fill(v, field.dataType, n, rnd, nullFraction = 0.0, reversed = false)
-      }
-      val batch = new ColumnarBatch(vectors.toArray[ColumnVector], n)
-      val project = UnsafeProjection.create(schema)
-      val writer = ArrowWriter.create(root, n)
-      (0 until n).foreach { i =>
-        val row = project(batch.getRow(i)).copy()
-        writer.write(row)
-        withClue(s"row $i: ") {
-          reachable(writer, row.getBaseObject) shouldBe false
+    Seq(0.0, 0.5).foreach { nullFraction =>
+      val rnd = new Random(7)
+      val vectors = schema.fields.map(f => newVector(n, f.dataType, offHeap = false))
+      val allocator = new RootAllocator(Long.MaxValue)
+      val root = VectorSchemaRoot.create(Utils.toArrowSchema(schema, "UTC"), allocator)
+      try {
+        schema.fields.zip(vectors).foreach { case (field, v) =>
+          fill(v, field.dataType, n, rnd, nullFraction, reversed = false)
+        }
+        val batch = new ColumnarBatch(vectors.toArray[ColumnVector], n)
+        val project = UnsafeProjection.create(schema)
+        val rows = (0 until n).map(i => project(batch.getRow(i)).copy())
+        Seq("finish", "reset").foreach { end =>
+          val writer = ArrowWriter.create(root, n)
+          rows.foreach { row =>
+            writer.write(row)
+            // A generic row's unsafe values take the writers' generic path into the same views.
+            writer.write(new GenericInternalRow(Array.tabulate[Any](schema.length)(f =>
+              if (row.isNullAt(f)) null else row.get(f, schema(f).dataType))))
+          }
+          val held = rows.count(row => reachable(writer, row.getBaseObject))
+          withClue(s"nulls=$nullFraction, rows held before $end: ") {
+            held should be <= 5 // one per nested writer at most
+          }
+          if (end == "finish") writer.finish() else writer.reset()
+          rows.zipWithIndex.foreach { case (row, i) =>
+            withClue(s"nulls=$nullFraction row $i after $end: ") {
+              reachable(writer, row.getBaseObject) shouldBe false
+            }
+          }
         }
+      } finally {
+        root.close()
+        allocator.close()
+        vectors.foreach(_.close())
       }
-    } finally {
-      root.close()
-      allocator.close()
-      vectors.foreach(_.close())
     }
   }