comphead commented on code in PR #6730:
URL: https://github.com/apache/datafusion-comet/pull/6730#discussion_r4212172334
##########
spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala:
##########
@@ -640,4 +685,144 @@ class CometArrowWriterSuite extends AnyFunSuite with
Matchers {
}
}
}
+
+ // Nested shapes whose unsafe forms take paths of their own: an array of
each primitive type,
+ // whose elements are either copied in one block or converted one at a time,
collections inside
+ // collections, and a struct wider than one word of null bits.
+ private val moreNestedTypes: Seq[DataType] =
primitiveTypes.map(ArrayType(_)) ++ Seq(
Review Comment:
`primitiveTypes.map(ArrayType(_))` includes `ArrayType(IntegerType)` and
`ArrayType(StringType)`, which `nestedTypes` already lists. Both tests below
concatenate the two lists, so the single-column `array<int>` and
`array<string>` cases run twice in each of the six parameterized tests and in
the 64-element test. The second run repeats the first exactly, because
`assertFilledRowsMatch` seeds its `Random` with the schema's hash. Could we add
a `.distinct` to the two `types` lists, or `filterNot(nestedTypes.contains)`
here?
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala:
##########
@@ -1262,6 +1436,107 @@ private[arrow] class StringWriter(val valueVector:
VarCharVector) extends ArrowF
super.reset()
dictionaryCache.clear()
}
+
+ // The value is checked before Arrow grows anything, so one that cannot be
written changes
+ // nothing.
+ override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal:
Int): Unit = {
+ if (row.isNullAt(ordinal)) {
+ setNull()
+ } else {
+ val offsetAndSize = row.getLong(ordinal)
+ val start = valueVector.getStartOffset(valueVector.getLastSet + 1)
+ val length = (checkedEnd(start.toLong, offsetAndSize) - start).toInt
+ // Grows the buffers, and gives null rows written since the last value
their empty offsets.
+ valueVector.setValueLengthSafe(count, length)
+ copyMemory(
+ row.getBaseObject,
+ row.getBaseOffset + (offsetAndSize >> 32),
+ valueVector.getDataBuffer.memoryAddress + start,
+ length.toLong)
+ BitVectorHelper.setBit(valueVector.getValidityBuffer, count.toLong)
+ }
+ count += 1
+ }
+
+ // Every element is checked before anything changes, as one value is, so the
buffers grow once
+ // and the bytes are copied without further checks.
+ override private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit
= {
+ val numElements = array.numElements()
+ if (numElements == 0) {
+ return
+ }
+ val start = valueVector.getStartOffset(valueVector.getLastSet + 1)
+ var end = start.toLong
+ var i = 0
+ while (i < numElements) {
+ if (!array.isNullAt(i)) {
+ end = checkedEnd(end, array.getLong(i))
+ }
+ i += 1
+ }
+ while (valueVector.getValueCapacity < count + numElements) {
+ valueVector.reallocValidityAndOffsetBuffers()
+ }
+ reserveData(valueVector, end)
+ if (valueVector.getLastSet < count - 1) {
+ valueVector.fillEmpties(count)
+ }
Review Comment:
`writeVariableWidth` and `writeDictionaryVariableWidth` open with the same
two steps as these lines: `fillEmpties` when `getLastSet < outStart - 1`, then
`reallocValidityAndOffsetBuffers` until `getValueCapacity` reaches `outStart +
numRows`. Could the three share one helper in `ArrowFieldWriter`, next to
`reserveData`? Then a change to how a variable-width append grows its buffers
has one place to go. The order relative to `reserveData` does not matter. In
Arrow 18.3 `fillEmpties` goes through `handleSafe(index, 0)`, which here can
grow only the validity and offset buffers.
--
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]
