sunchao commented on code in PR #6566:
URL: https://github.com/apache/datafusion-comet/pull/6566#discussion_r4174917856
##########
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:
[P2] Preserve file context before advancing the source. When a native filter
consumes converted Parquet batches and a Spark projection evaluates
`input_file_name()` or `input_file_block_*()`, this call advances `FileScanRDD`
and changes its `InputFileBlockHolder` before the current rows reach that
projection. Combining batches from two files therefore assigns the second
file's metadata to the first file's rows. The reproduced query returned
incorrect metadata for 8,192 rows, while Spark and the base implementation
returned the correct values. Retain source-batch boundaries for these plans, or
materialize the metadata before coalescing.
Evidence: Full Spark 4.1.3/JDK 17 reproduction: write two Parquet files with
8,192 rows each, read both in one partition, set
`spark.sql.parquet.columnarReaderBatchSize=8192`,
`spark.comet.batchSize=16384`, `spark.comet.scan.enabled=false`, and
`spark.comet.convert.parquet.enabled=true`. Apply `where("id >= 0")` then
select `id` and the three `input_file_*()` expressions. Verified plan contains
`Project -> CometColumnarToRow -> CometFilter -> CometSparkColumnarToColumnar`.
Head: 8,192 mismatches among 16,384 rows, including the wrong filename and
block length. Base reader/writer substitution: zero mismatches. Harness:
`/tmp/review6566-bx8ympbd/FullPlanProbe.scala`; results: `full-head.log` and
`full-base.log` in that directory.
##########
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:
[P2] Exclude Spark's missing/all-null vectors from bulk data reads. Spark
marks an absent nested Parquet field as `isAllNull` and does not grow its
backing storage with the surrounding collection. An array containing 5,000
structs can therefore have a missing integer child backed by only 4,096 slots.
This new bulk branch calls `getInts(0, 5000)` despite every value being null,
throwing `ArrayIndexOutOfBoundsException`. Materializing that valid evolved
schema through `ArrowCachedBatchSerializer` now fails instead of caching nulls
as Spark and the base writer do. Check `isAllNull` before selecting either
Spark-vector bulk path and write only validity for that case.
Evidence: Reproduced through the full Spark cache path with on-heap vectors
and a 4,096-row Parquet batch size. Write one row containing
`transform(sequence(1,5000), x -> named_struct('a', x)) AS items`, then read
with schema `items array<struct<a:int, added:int>>` and call `.cache().count()`
using `ArrowCachedBatchSerializer`. Head fails with `arraycopy: last source
index 5000 out of bounds for int[4096]`, through `getInts -> bulkCopy ->
StructWriter -> ArrayWriter -> ArrowCachedBatchSerializer`. Base reader/writer
substitution returns `1`; the direct writer comparison also confirms 5,000 null
`added` values. Harnesses and logs are under `/tmp/review6566-bx8ympbd/`
(`FullPlanProbe.scala`, `NestedEvolutionProbe.scala`, `full-head.log`,
`full-base.log`).
--
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]
