gengliangwang commented on code in PR #59933:
URL: https://github.com/apache/spark/pull/59933#discussion_r4420671900
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala:
##########
@@ -214,15 +237,223 @@ class ArrowBatchingDataWriter(delegate:
DataWriter[ColumnarBatch], options: Arro
}
}
- private def flush(): Unit = {
- if (root != null) {
- arrowWriter.finish()
- val columns = root.getFieldVectors.asScala.map(new ArrowColumnVector(_):
ColumnVector)
+ override def writeBatch(batch: ColumnarBatch): Unit = {
+ val numBatchRows = batch.numRows()
+ val copied = copiedVectors(batch)
+ if (copied.forall(_.isEmpty)) {
+ // Without columns to copy, the values are converted as rows are.
+ var i = 0
+ while (i < numBatchRows) {
+ write(batch.getRow(i))
+ i += 1
+ }
+ } else if (numBatchRows > 0) {
+ // A batch without rows leaves the pending batch of rows as it is.
+ flush()
+ // A window of up to `maxRecordsPerBatch` rows of the batch, whose
columns that are not copied
+ // are converted at once. The next window starts where the last batch
written ended, so that
+ // the batches end as those of rows do.
+ var windowStart = 0
+ while (windowStart < numBatchRows) {
+ val first = windowStart
+ val windowEnd = math.min(numBatchRows, first +
options.maxRecordsPerBatch)
+ val converted = ArrayBuffer.empty[FieldVector]
+ windowStart = Utils.tryWithSafeFinally {
+ val columns = copied.indices.map { i =>
+ copied(i) match {
+ case Some(vector) => WindowColumn(vector, 0, copied = true)
+ case None =>
+ val vector = convertColumn(batch, i, first, windowEnd)
+ converted += vector
+ WindowColumn(vector, first, copied = false)
Review Comment:
**[P1] Bound column conversion by the byte limit**
When a batch mixes copied columns with a converted `utf8_view` or
`large_utf8` column, this materializes the entire record-count window before
`batchLength` applies the byte limit. For example, 1,000 strings of 3 MiB each
require a roughly 3 GiB canonical string vector despite the default 64 MiB
batch limit, potentially causing OOM or Arrow's 2 GiB buffer-limit exception.
The previous row path flushed bounded batches successfully. Convert
incrementally or bound the conversion window by bytes before allocating it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala:
##########
@@ -214,15 +237,223 @@ class ArrowBatchingDataWriter(delegate:
DataWriter[ColumnarBatch], options: Arro
}
}
- private def flush(): Unit = {
- if (root != null) {
- arrowWriter.finish()
- val columns = root.getFieldVectors.asScala.map(new ArrowColumnVector(_):
ColumnVector)
+ override def writeBatch(batch: ColumnarBatch): Unit = {
+ val numBatchRows = batch.numRows()
+ val copied = copiedVectors(batch)
+ if (copied.forall(_.isEmpty)) {
+ // Without columns to copy, the values are converted as rows are.
+ var i = 0
+ while (i < numBatchRows) {
+ write(batch.getRow(i))
+ i += 1
+ }
+ } else if (numBatchRows > 0) {
+ // A batch without rows leaves the pending batch of rows as it is.
+ flush()
+ // A window of up to `maxRecordsPerBatch` rows of the batch, whose
columns that are not copied
+ // are converted at once. The next window starts where the last batch
written ended, so that
+ // the batches end as those of rows do.
+ var windowStart = 0
+ while (windowStart < numBatchRows) {
+ val first = windowStart
+ val windowEnd = math.min(numBatchRows, first +
options.maxRecordsPerBatch)
+ val converted = ArrayBuffer.empty[FieldVector]
+ windowStart = Utils.tryWithSafeFinally {
+ val columns = copied.indices.map { i =>
+ copied(i) match {
+ case Some(vector) => WindowColumn(vector, 0, copied = true)
+ case None =>
+ val vector = convertColumn(batch, i, first, windowEnd)
+ converted += vector
+ WindowColumn(vector, first, copied = false)
+ }
+ }
+ var start = first
+ var next = -1
+ while (next < 0) {
+ val (length, full) = batchLength(columns, start, windowEnd)
+ if (start > first && start + length == windowEnd && windowEnd <
numBatchRows && !full) {
+ // The batch ends with the window, before it holds
`maxBytesPerBatch` bytes, so it may
+ // hold rows after the window: it starts the next window, whose
conversion converts
+ // its rows again.
+ next = start
+ } else {
+ writeRows(columns, start, length)
+ start += length
+ if (start == windowEnd) next = start
+ }
+ }
+ next
+ } {
+ AutoCloseables.close(converted.asJava)
+ }
+ }
+ }
+ }
+
+ /**
+ * Returns, for each column of a batch, its Arrow vector if it can be copied
buffer by buffer: an
+ * Arrow vector with as many values as the batch has rows, whose field is
physically congruent
+ * with the one of `arrowSchema`, though it may differ in the names,
nullability and metadata of
+ * its fields, and in the time zones of timestamps, which the copy keeps.
The other columns are
+ * converted value by value.
+ */
+ private def copiedVectors(batch: ColumnarBatch): Seq[Option[FieldVector]] = {
+ val fields = arrowSchema.getFields.asScala
+ if (batch.numCols() != fields.length) {
+ Seq.fill(fields.length)(None)
+ } else {
+ fields.indices.map { i =>
+ batch.column(i) match {
+ case column: ArrowColumnVector => column.getValueVector match {
+ case vector: FieldVector if vector.getValueCount ==
batch.numRows() &&
+ ArrowUtils.isCompatibleWithDeclaredField(vector.getField,
fields(i)) =>
+ Some(vector)
Review Comment:
**[P2] Exclude unsliceable Arrow vectors from the copy path**
When a scan returns a fixed-width Arrow vector larger than 2 GiB, these
checks accept it even though `writeRows` cannot safely slice it. Arrow's
`splitAndTransfer` computes fixed-width byte offsets in 32-bit arithmetic, so
later slices overflow and abort a valid write that previously converted into
small row batches. Add a recursive slice-safety check, analogous to
`BaseSliceArrowOutputProcessor.canSlice`, and fall back to value conversion for
affected columns.
--
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]
