andygrove commented on code in PR #6566:
URL: https://github.com/apache/datafusion-comet/pull/6566#discussion_r4178060500
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala:
##########
@@ -70,18 +74,30 @@ private[comet] class SparkColumnarArrowReader(
return false
}
- val startNs = System.nanoTime()
- val rowsRemaining = current.numRows() - rowsConsumedInCurrent
- val rowsToProduce =
- if (maxRecordsPerBatch <= 0) rowsRemaining
- else math.min(maxRecordsPerBatch, rowsRemaining)
-
- val writer = ArrowWriter.create(getVectorSchemaRoot, rowsToProduce)
- writer.writeColumns(current, rowsConsumedInCurrent, rowsToProduce)
- rowsConsumedInCurrent += rowsToProduce
-
+ // Without a limit, each Spark batch becomes one Arrow batch.
+ val batchSize =
+ if (maxRecordsPerBatch <= 0) current.numRows() - rowsConsumedInCurrent
+ else maxRecordsPerBatch
+ // Pulling the next source batch is the producer's time, not conversion
time.
+ var conversionNs = 0L
+ var startNs = System.nanoTime()
+ val writer = ArrowWriter.create(getVectorSchemaRoot, batchSize)
+ var rowsProduced = 0
+ var more = true
+ while (more) {
+ val rows = math.min(batchSize - rowsProduced, current.numRows() -
rowsConsumedInCurrent)
+ writer.writeColumns(current, rowsConsumedInCurrent, rows)
+ rowsConsumedInCurrent += rows
+ rowsProduced += rows
+ more = rowsProduced < batchSize && {
+ conversionNs += System.nanoTime() - startNs
+ val advanced = advanceToNonEmptyBatch()
Review Comment:
Fixed in `2b3776391`. `SparkColumnarArrowReader` now snapshots the
`InputFileBlockHolder` context for each source batch, only combines batches
with the same path, start, and length, and restores the emitted batch context
when it buffers a batch from the next file. I added a regression test with
reused vectors that covers both combining batches from one file and stopping at
a file boundary. The focused suites pass 54/54 on Spark 4.1 and Spark 3.4.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala:
##########
@@ -266,65 +691,149 @@ private[arrow] abstract class ArrowFieldWriter {
}
private[arrow] abstract class FixedWidthArrowFieldWriter extends
ArrowFieldWriter {
+ import ArrowFieldWriter._
override def valueVector: BaseFixedWidthVector
protected def setValueUnsafe(input: SpecializedGetters, ordinal: Int): Unit
- private def ensureCapacity(inputNumElements: Int): Unit = {
+ protected def ensureCapacity(inputNumElements: Int): Unit = {
while (valueVector.getValueCapacity < inputNumElements) {
valueVector.reAlloc()
}
}
- private def tryBulkCopyNoNull(input: ColumnVector, startRow: Int, numRows:
Int): Boolean = {
- // Spark's bulk getters allocate and fill a temporary array before the
copy into Arrow. Keep
- // slices below 32 on the scalar path to avoid the observed tiny-slice
regression.
- if (count != 0 || numRows < 32 || ByteOrder.nativeOrder() !=
ByteOrder.LITTLE_ENDIAN) {
- return false
- }
-
- val supportedInput = input match {
- case vector: OnHeapColumnVector => !vector.hasDictionary
- case vector: OffHeapColumnVector => !vector.hasDictionary
+ /**
+ * Copies the values of rows `[startRow, startRow + numRows)` of `input`
into the data buffer at
+ * index `count`, leaving validity alone. What lands under a null row is
unspecified. Returns
+ * false, having written nothing, for a vector type with no such copy.
+ *
+ * Spark's own vectors without a dictionary are copied in bulk, on-heap ones
from
+ * [[ArrowFieldWriter.MinOnHeapBulkCopyRows]] rows. Everything else,
including a
+ * dictionary-encoded vector, is read value by value, skipping null rows,
because a dictionary
+ * id under a null may be garbage.
+ */
+ protected def copyValues(input: ColumnVector, startRow: Int, numRows: Int):
Boolean = {
+ val target = valueVector.getDataBufferAddress + count.toLong *
valueVector.getTypeWidth
+ val bulk = input match {
+ case vector: OffHeapColumnVector => LittleEndian && !vector.hasDictionary
+ case vector: OnHeapColumnVector =>
+ LittleEndian && !vector.hasDictionary && numRows >=
MinOnHeapBulkCopyRows
Review Comment:
Fixed in `2b3776391`. Both the on-heap and off-heap fixed-width bulk paths
now exclude `isAllNull` vectors, so the fallback writes null validity without
reading undersized value storage. I added a regression test with 5,000 logical
rows over 4,096 allocated slots in both heap modes. The focused suites pass
54/54 on Spark 4.1 and Spark 3.4.
--
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]
