viirya commented on code in PR #6730:
URL: https://github.com/apache/datafusion-comet/pull/6730#discussion_r4210628124
##########
spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala:
##########
@@ -640,4 +656,146 @@ class CometArrowWriterSuite extends AnyFunSuite with
Matchers {
}
}
}
+
+ // Nested shapes whose unsafe forms take paths of their own: elements
converted one at a time,
+ // collections inside collections, and a struct wider than one word of null
bits.
+ private val moreNestedTypes: Seq[DataType] = Seq(
Review Comment:
`unsafeValueWidth` in `ArrowWriters.scala` lists twelve vector types whose
values are copied out of an unsafe array in one block. That relies on Spark's
element size matching the Arrow type width for each of them. In these tests
only int and long elements reach that copy (`ArrayType(IntegerType)`,
`ArrayType(ArrayType(LongType))`, and the int and long map keys). Byte, short,
float, double, date, timestamp, timestamp_ntz, the two interval types and time
never go through it as array elements.
Could we add `primitiveTypes.map(ArrayType(_))` here? `fill` already handles
every one of those element types, so this should be a one-line change that
covers each entry in the list, including the 1- and 2-byte widths.
##########
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:
`ArrayWriter`, `StructWriter` and `MapWriter` each have the same
`writeUnsafeRowField` and `writeArrayElements`. The only difference is which
reused view `writeUnsafe` points at the value. Would it work to move those two
methods into a small shared base class with an abstract `writeUnsafe(base,
baseOffset, offsetAndSize)` that each writer implements? Then the null and
`count` handling would live in one place rather than three. The stale-view
mutations in the description show that a change here currently has to be made
in each writer separately.
--
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]
