viirya commented on code in PR #6730:
URL: https://github.com/apache/datafusion-comet/pull/6730#discussion_r4215125697
##########
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:
The retention is fixed in 0623dbac51 by allocating a view per value again,
and the new reachability test fails on the previous head. Resolved from my side.
--
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]
