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()) } }