viirya commented on code in PR #6730:
URL: https://github.com/apache/datafusion-comet/pull/6730#discussion_r4211555262
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala:
##########
@@ -1384,16 +1627,59 @@ 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.
+ private val elements = new UnsafeArrayData
+
override def setNull(): Unit = {}
override def setValue(input: SpecializedGetters, ordinal: Int): Unit = {
- val array = input.getArray(ordinal)
+ input.getArray(ordinal) match {
+ case unsafe: UnsafeArrayData =>
+ writeElements(unsafe)
+ case array =>
+ val numElements = array.numElements()
+ valueVector.startNewValue(count)
+ var i = 0
+ while (i < numElements) {
+ elementWriter.write(array, i)
+ i += 1
+ }
+ valueVector.endValue(count, numElements)
+ }
+ }
+
+ override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal:
Int): Unit = {
+ if (row.isNullAt(ordinal)) {
+ setNull()
+ } else {
+ writeUnsafe(row.getBaseObject, row.getBaseOffset, row.getLong(ordinal))
+ }
+ count += 1
+ }
+
+ override private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit
= {
+ val numElements = array.numElements()
var i = 0
- valueVector.startNewValue(count)
- while (i < array.numElements()) {
- elementWriter.write(array, i)
+ while (i < numElements) {
+ if (array.isNullAt(i)) {
+ setNull()
+ } else {
+ 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)
Review Comment:
+1, and I should have caught this in my first pass. Old code allocated a
fresh view per value, so nothing outlived the copy. Here each of `elements`,
`fields` and `entries` keeps the last backing object it was pointed at, and
`entries` does the same through its key and value arrays. It is at most one
object per nested writer, and only until the batch's writer is dropped. But
when those objects are per-row copies, or on-heap pages that the memory manager
has already freed, that is heap Spark doesn't account for.
One thing to watch in the fix: `pointTo(null, 0L, 0)` is fine for
`UnsafeRow`, but `UnsafeArrayData.pointTo` and `UnsafeMapData.pointTo` read
their header with `Platform.getLong(baseObject, baseOffset)`. With a null base
and a zero offset, that reads address 0 and would crash the JVM. Pointing them
at a shared empty layout avoids that. That would be a `long[] {0}` for an array
of no elements, and `long[] {8, 0, 0}` with a size of 24 for a map, which is a
key array of 8 bytes holding 0 elements and then an empty value array.
Detaching once per row, after `ArrowWriter.write` finishes the row, may be
cheaper than after every nested value, since an `array<struct>` would otherwise
detach once per element. Whichever measures better is fine with me.
For the test, could we check that each view no longer references the input
after a write rather than relying on GC with weak references? Something like a
package-private accessor for the views' base objects would keep it
deterministic.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala:
##########
@@ -1384,16 +1627,59 @@ 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.
+ private val elements = new UnsafeArrayData
+
override def setNull(): Unit = {}
override def setValue(input: SpecializedGetters, ordinal: Int): Unit = {
- val array = input.getArray(ordinal)
+ input.getArray(ordinal) match {
+ case unsafe: UnsafeArrayData =>
+ writeElements(unsafe)
+ case array =>
+ val numElements = array.numElements()
+ valueVector.startNewValue(count)
+ var i = 0
+ while (i < numElements) {
+ elementWriter.write(array, i)
+ i += 1
+ }
+ valueVector.endValue(count, numElements)
+ }
+ }
+
+ override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal:
Int): Unit = {
Review Comment:
Thanks for measuring it. A shared call site going megamorphic is a good
reason to keep the three copies, and the stale-view mutations show each copy is
covered. I'm fine leaving this as is.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
