andygrove commented on code in PR #6730:
URL: https://github.com/apache/datafusion-comet/pull/6730#discussion_r4212937968
##########
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:
Done in 90c3bbaee7. `reserveValues` sits next to `reserveData`, and all
three paths call it, with the unsafe array path calling it after its checks.
While I checked each call site with a mutation, dropping it from the dictionary
path failed nothing, because no test appended dictionary strings after nulls
the row path had left, or past the initial capacity. There's a test for that
now, which appends a dictionary-encoded field as a column after its struct took
the row path.
##########
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:
Thanks, fixed in 90c3bbaee7 with `filterNot(nestedTypes.contains)`, so
array<int> and array<string> each run once.
--
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]
