HyukjinKwon commented on code in PR #59933:
URL: https://github.com/apache/spark/pull/59933#discussion_r4423374058
##########
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:
Done in 518df52: the columns that are not copied are now converted row by
row, from the start of each batch up to the row at which it ends, instead of a
window of `maxRecordsPerBatch` rows at once. As for rows, the size of the batch
is computed after each row: `rangeSize` for the copied columns and the field
writers for the converted ones. So a batch that mixes copied columns with
converted `utf8_view` or `large_utf8` ones ends at `maxBytesPerBatch` as the
row path does, and the conversion never holds more than the rows of that batch.
Each row is converted once, and the converted vectors are written as they are,
without a slice. A batch whose columns are all copied still finds its end with
the binary search over `rangeSize`.
The new test "convert the columns of a batch that are not copied up to the
end of each batch" writes 20 rows of 1 MiB `large_utf8` names with copied ids,
with limits of 10 rows and 2 MiB. It gets 10 batches of 2 rows and 20
conversions, and the peak memory of the writer stays below 5 MiB. With the
window, the same test peaks at 24 MiB.
##########
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:
Done in 518df52: `ArrowBatchingDataWriter.canCopy`, which `copiedVectors`
now uses, also requires `BaseSliceArrowOutputProcessor.canSlice`. It checks the
vector and its children: a fixed-width vector of 2 GiB or more, a view vector
with 2^27 rows or more, or a list vector with 2^29 - 1 rows or more. Such a
column is converted value by value instead, into batches bounded as above. I
widened `canSlice` to `private[execution]` instead of copying it.
The new test "copy only the Arrow columns that Arrow can slice" checks that
`canCopy` rejects vectors that report more values than they hold, as the test
of `canSlice` does: a fixed-width vector of 2 GiB, and a list with such a
child. It also checks that `canCopy` still accepts an ordinary vector. The
guide now lists these columns among those that Spark converts.
--
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]
