mbutrovich commented on code in PR #5046:
URL: https://github.com/apache/datafusion-comet/pull/5046#discussion_r3731419307
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala:
##########
Review Comment:
`RowArrowReader` (RowArrowReader.scala:62) drives this method through
`writer.write(row)` for every row of the fallback, non-columnar path. For a
`FixedWidthArrowFieldWriter`, `setValue` here is still the original `setSafe`
call, since `setValueUnsafe` is only reached from `writeCol`/`writeColNoNull`.
So every field of every row on this path still runs Arrow's `handleSafe`
capacity check, the exact per-write cost the PR description says this change
removes.
`RowArrowReader` now builds its writer with `create(root,
maxRecordsPerBatch.toInt)`, so `count < capacity` holds for the whole batch by
construction, the same invariant `writeCol`/`writeColNoNull` rely on to call
`setValueUnsafe` directly. Routing `write` through the unsafe setter for
fixed-width fields looks safe under that same invariant.
This path backs `CometSparkRowToColumnar`, used whenever the child doesn't
support columnar output. That's likely to be a common case, and it's worth
having a benchmark number for it the same way `CometArrowWriterBenchmark`
covers `SparkColumnarArrowReader`, since right now there is no measurement for
how much of the intended win this path is actually getting.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala:
##########
@@ -224,89 +228,133 @@ private[arrow] abstract class ArrowFieldWriter {
}
}
-private[arrow] class BooleanWriter(val valueVector: BitVector) extends
ArrowFieldWriter {
+private[arrow] abstract class FixedWidthArrowFieldWriter extends
ArrowFieldWriter {
+
+ override def valueVector: BaseFixedWidthVector
+
+ protected def setValueUnsafe(input: SpecializedGetters, ordinal: Int): Unit
override def setNull(): Unit = {
valueVector.setNull(count)
}
- override def setValue(input: SpecializedGetters, ordinal: Int): Unit = {
- valueVector.setSafe(count, if (input.getBoolean(ordinal)) 1 else 0)
+ protected def setNullUnsafe(): Unit = {
+ BitVectorHelper.unsetBit(valueVector.getValidityBuffer, count)
}
-}
-private[arrow] class ByteWriter(val valueVector: TinyIntVector) extends
ArrowFieldWriter {
+ override def writeCol(input: ColumnarArray): Unit = {
+ val inputNumElements = input.numElements()
+ while (valueVector.getValueCapacity < inputNumElements) {
+ valueVector.reAlloc()
+ }
+ while (count < inputNumElements) {
+ if (input.isNullAt(count)) {
+ setNullUnsafe()
+ } else {
+ setValueUnsafe(input, count)
+ }
+ count += 1
+ }
+ }
- override def setNull(): Unit = {
- valueVector.setNull(count)
+ override def writeColNoNull(input: ColumnarArray): Unit = {
Review Comment:
For the no-null case, is a bulk copy of the source column into the Arrow
data buffer worth exploring instead of one `setValueUnsafe` call per row?
`BaseFixedWidthVector.getDataBuffer()` / `getDataBufferAddress()` plus a raw
memory copy is already the pattern `copyFrom` uses internally in arrow-java
when the source is another Arrow vector. Spark's on-heap and off-heap
`ColumnVector` implementations back fixed-width columns with contiguous
primitive arrays, so the same kind of bulk transfer should be reachable here
for the no-null case, cutting the per-row virtual dispatch entirely rather than
just skipping the capacity check on it.
This is a bigger change than this PR's scope, so if it's not something you
want to take on here, could you file a follow-up issue for it? I'd rather have
that tracked than have it drop once this merges.
--
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]
