HyukjinKwon commented on code in PR #59933:
URL: https://github.com/apache/spark/pull/59933#discussion_r4418552457
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala:
##########
@@ -185,17 +193,26 @@ object ArrowBatchingOptions {
}
/**
- * Converts the rows of a write task to Arrow-backed [[ColumnarBatch]]es for
the [[DataWriter]]
- * of a [[ColumnarDataSource]]. A new Arrow vector schema root is created for
each batch, so a
- * native writer that took ownership of the data of a batch can keep it after
the call.
+ * Converts the data of a write task to Arrow-backed [[ColumnarBatch]]es for
the [[DataWriter]]
+ * of a [[ColumnarDataSource]]: its rows, or the columnar batches of a
columnar query. Each batch
+ * gets new Arrow vectors, so a native writer that took ownership of the data
of a batch can keep
+ * it after the call, while the producer of a columnar batch may reuse its
vectors once the call
+ * returns.
+ *
+ * Each column of a columnar batch that is an Arrow vector in the layout of
the Arrow schema of the
+ * write is copied buffer by buffer, and the other columns are converted value
by value. The
+ * batches end as those of rows do, though a batch with copied columns also
ends with the columnar
+ * batch. A columnar batch without columns to copy is converted as rows are.
*/
class ArrowBatchingDataWriter(delegate: DataWriter[ColumnarBatch], options:
ArrowBatchingOptions)
- extends DataWriter[InternalRow] {
+ extends ColumnarInputDataWriter {
+ import ArrowBatchingDataWriter.WindowColumn
private val arrowSchema = ArrowUtils.toArrowSchema(
options.schema, options.timeZoneId, errorOnDuplicatedFieldNames = true,
largeVarTypes = false)
private val allocator =
ArrowUtils.rootAllocator.newChildAllocator("columnar data source writer",
0, Long.MaxValue)
Review Comment:
Done in f2ca665: the contract of `ColumnarWriter`, the doc of
`ArrowBatchingDataWriter`, `write` in the header and the guide now say that a
data writer may keep the data of a batch until it is closed, and must release
it by then, since its buffers belong to the data writer.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala:
##########
@@ -214,15 +231,201 @@ 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 {
+ flush()
+ // A window of the rows of the batch, whose columns that are not copied
are converted at once.
+ var windowStart = 0
+ while (windowStart < numBatchRows) {
+ val windowEnd = math.min(numBatchRows, windowStart +
options.maxRecordsPerBatch)
+ val converted = copied.indices.map { i =>
+ if (copied(i).isEmpty) Some(convertColumn(batch, i, windowStart,
windowEnd)) else None
+ }
Review Comment:
Done in f2ca665: the columns are built in one pass inside the
`tryWithSafeFinally`, which closes the vectors converted so far. A new test
checks that a decimal that fails to convert after another converted column
leaves nothing allocated.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala:
##########
@@ -214,15 +231,201 @@ 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 {
+ flush()
+ // A window of the rows of the batch, whose columns that are not copied
are converted at once.
+ var windowStart = 0
+ while (windowStart < numBatchRows) {
+ val windowEnd = math.min(numBatchRows, windowStart +
options.maxRecordsPerBatch)
+ val converted = copied.indices.map { i =>
+ if (copied(i).isEmpty) Some(convertColumn(batch, i, windowStart,
windowEnd)) else None
+ }
+ Utils.tryWithSafeFinally {
+ val columns = copied.zip(converted).map {
+ case (Some(vector), _) => WindowColumn(vector, 0, copied = true)
+ case (_, Some(vector)) => WindowColumn(vector, windowStart, copied
= false)
+ case _ =>
+ throw SparkException.internalError("A column is neither copied
nor converted.")
+ }
+ var start = windowStart
+ while (start < windowEnd) {
+ val length = batchLength(columns, start, windowEnd)
+ writeWindow(columns, start, length)
+ start += length
+ }
+ } {
+ converted.flatten.foreach(_.close())
+ }
+ windowStart = windowEnd
+ }
+ }
+ }
+
+ /**
+ * 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)
+ case _ => None
+ }
+ case _ => None
+ }
+ }
+ }
+ }
+
+ /** Converts the values of the rows [start, end) of a column into a vector
of `arrowSchema`. */
+ private def convertColumn(
+ batch: ColumnarBatch,
+ ordinal: Int,
+ start: Int,
+ end: Int): FieldVector = {
+ val vector = arrowSchema.getFields.get(ordinal).createVector(allocator)
+ try {
+ vector.setInitialCapacity(end - start)
+ vector.allocateNew()
+ val fieldWriter = ArrowWriter.createFieldWriter(vector)
+ var i = start
+ while (i < end) {
+ fieldWriter.write(batch.getRow(i), ordinal)
+ i += 1
+ }
+ fieldWriter.finish()
+ vector
+ } catch {
+ case e: Throwable =>
+ vector.close()
+ throw e
+ }
+ }
+
+ /**
+ * Returns the number of rows of the batch that starts at the row `start` of
the columns: up to
+ * `end`, or as a batch of rows ends, up to the first row with which it
holds `maxBytesPerBatch`
+ * bytes, by the measure of `getBufferSizeFor`, which `rangeSize` computes
without slicing.
+ *
+ */
+ private def batchLength(columns: Seq[WindowColumn], start: Int, end: Int):
Int = {
+ // The size is not computed when the limit is the largest possible one, as
for rows.
+ if (options.maxBytesPerBatch >= Int.MaxValue) {
+ end - start
+ } else {
+ def sizeInBytes(rows: Int): Long = columns.map { column =>
+ ArrowBatchingDataWriter.rangeSize(
+ column.vector, start - column.first, start - column.first + rows)
+ }.sum
+ // The size grows with the number of rows.
+ var lo = 1
+ var hi = end - start
+ while (lo < hi) {
+ val mid = (lo + hi) >>> 1
+ if (sizeInBytes(mid) >= options.maxBytesPerBatch) hi = mid else lo =
mid + 1
+ }
+ lo
+ }
+ }
+
+ /**
+ * Writes the rows [start, start + length) of the columns. A copied column
is sliced, which shares
+ * the buffers of the batch, and its buffers copied once each at their size.
Arrow Java reads the
+ * offsets of a vector of variable width, or of a list, as if they started
at 0, which those of a
+ * slice do. A converted column, whose vector nothing else writes to, is
only sliced.
+ */
+ private def writeWindow(columns: Seq[WindowColumn], start: Int, length:
Int): Unit = {
+ val copiedColumns = columns.filter(_.copied).map(_.vector)
+ // The vectors of the batch, which this root does not own: their producer
releases them.
+ val batchRoot = new VectorSchemaRoot(copiedColumns.map(_.getField).asJava,
+ copiedColumns.asJava,
copiedColumns.headOption.map(_.getValueCount).getOrElse(0))
+ val written = ArrayBuffer.empty[AutoCloseable]
+ Utils.tryWithSafeFinally {
+ val copies = if (copiedColumns.isEmpty) {
+ Iterator.empty
+ } else {
+ val slice = batchRoot.slice(start, length)
+ val copy = Utils.tryWithSafeFinally(copyRoot(slice))(slice.close())
+ written += copy
+ copy.getFieldVectors.asScala.iterator
+ }
+ val vectors = columns.map { column =>
+ if (column.copied) {
+ copies.next()
+ } else {
+ val pair = column.vector.getTransferPair(column.vector.getAllocator)
+ pair.splitAndTransfer(start - column.first, length)
+ val slice = pair.getTo.asInstanceOf[FieldVector]
+ written += slice
+ slice
+ }
+ }
+ delegate.write(new ColumnarBatch(
+ vectors.map(new ArrowColumnVector(_): ColumnVector).toArray, length))
+ } {
+ written.foreach(_.close())
+ }
+ }
+
+ /** Copies the buffers of the vectors of a root into new vectors, once each,
at their size. */
+ private def copyRoot(source: VectorSchemaRoot): VectorSchemaRoot = {
+ val copy = VectorSchemaRoot.create(source.getSchema, allocator)
+ try {
+ val unloaded = new VectorUnloader(source).getRecordBatch
try {
- delegate.write(new ColumnarBatch(columns.toArray, root.getRowCount))
+ val buffers = unloaded.getBuffers.asScala.map { buffer =>
+ val size = buffer.readableBytes()
+ val copied = allocator.buffer(size)
+ copied.setBytes(0, buffer, buffer.readerIndex(), size)
+ copied.writerIndex(size)
+ copied
+ }
Review Comment:
Done in f2ca665: the copies are collected into an `ArrayBuffer` inside the
`try`, as suggested.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala:
##########
@@ -214,15 +231,201 @@ 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 {
+ flush()
+ // A window of the rows of the batch, whose columns that are not copied
are converted at once.
+ var windowStart = 0
+ while (windowStart < numBatchRows) {
+ val windowEnd = math.min(numBatchRows, windowStart +
options.maxRecordsPerBatch)
Review Comment:
Done in f2ca665: a window now starts where the last batch written ended.
When a batch that the byte limit did not end reaches the end of the window, it
starts the next window instead, whose conversion converts those rows again.
--
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]
