From 28bec46d893a1acf1c413cea7dab8ec9b1743487 Mon Sep 17 00:00:00 2001 From: Alex Cruise Date: Mon, 2 Nov 2026 09:47:38 -0700 Subject: [PATCH] perf: reuse the nested unsafe views in ArrowWriter's row path, detaching them per batch The row path read each array, map and struct of an unsafe row through Spark's getters, which allocate a view per value. ArrayWriter, StructWriter and MapWriter now keep one UnsafeArrayData, UnsafeRow or UnsafeMapData each and point it at every value, as an earlier revision of #6730 did. A view holds the last row it read, and that revision was dropped because nothing ever let go of it. Detaching after every row turned out to cost what the reuse saves, one pointTo per nested writer per row, so ArrowWriter collects the views at construction (UnsafeViews, by type so each detach is a direct call) and points them at an empty layout in finish and reset. Between rows a writer holds at most one row per nested writer, for one batch. Every production caller builds a writer per batch as a local and finishes it, so a batch that throws drops its views with its frame. Throwaway harness, one column of 8192 unsafe rows per shape, all shapes warmed together, best of 30, median of three JVMs, JDK 21, assertions off, ns per row before -> after: array 28.5 -> 26.0, map 50.3 -> 45.6, array> 55.4 -> 51.5, array> 44.0 -> 42.3, map 29.5 -> 26.2, struct 21.9 -> 19.5, a row of eight mixed fields 119.6 -> 113.4. Flat columns are unchanged. The reachability test now allows one held row per nested writer before finish or reset and none after, for unsafe rows and generic rows holding unsafe values, with and without nulls. The strict-warnings profile gets a filter for the unused parameter of the default collectViews hook, as it has for the other extension hooks that leave a parameter for overrides. Part of #6765. Co-Authored-By: Claude Fable 5.1 --- pom.xml | 1 + .../comet/execution/arrow/ArrowWriters.scala | 125 ++++++++++++++++-- .../arrow/CometArrowWriterSuite.scala | 63 +++++---- 3 files changed, 156 insertions(+), 33 deletions(-) diff --git a/pom.xml b/pom.xml index a215e2a818f..93707c1fa55 100644 --- a/pom.xml +++ b/pom.xml @@ -1020,6 +1020,7 @@ under the License. -Wconf:cat=unused-params&site=org\.apache\.comet\.serde\.CometAggregateExpressionSerde\.supportsSparkPartialToNativeFinal\..*:s -Wconf:cat=unused-params&site=org\.apache\.comet\.serde\.CometAggregateExpressionSerde\.supportsNativePartialToSparkFinal\..*:s -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.CometScanWithPlanData\.withDynamicPruningFilters\..*:s + -Wconf:cat=unused-params&site=org\.apache\.spark\.sql\.comet\.execution\.arrow\.ArrowFieldWriter\.collectViews\..*:s -Wconf:cat=unused-params&site=org\.apache\.comet\.CometExecIterator\.InputFailure\.unapply\.nativeFailure:s -Wconf:cat=unused-params&site=org\.apache\.comet\.udf\.CometUDF\.evaluate\.(partitionIndex|planId):s diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala index 53b4ae7888d..b4c84c0ce24 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala +++ b/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) diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala index 400d9aad797..798d4ea7bc5 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala @@ -886,9 +886,12 @@ 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))) @@ -896,28 +899,42 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { .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()) } }