HyukjinKwon commented on code in PR #59933:
URL: https://github.com/apache/spark/pull/59933#discussion_r4417643695
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala:
##########
@@ -214,12 +227,119 @@ class ArrowBatchingDataWriter(delegate:
DataWriter[ColumnarBatch], options: Arro
}
}
+ override def writeBatch(batch: ColumnarBatch): Unit = {
+ val numBatchRows = batch.numRows()
+ arrowVectors(batch) match {
+ case Some(vectors) =>
+ flush()
+ // The vectors of the batch, which this root does not own: their
producer releases them.
+ val batchRoot =
+ new VectorSchemaRoot(vectors.map(_.getField).asJava, vectors.asJava,
numBatchRows)
+ var start = 0
+ while (start < numBatchRows) {
+ // A slice shares the buffers of the batch. Arrow Java copies a
vector of variable
+ // width, or a list, as if its offsets started at 0, which those of
a slice do.
+ val slice = batchRoot.slice(
+ start, math.min(numBatchRows - start, options.maxRecordsPerBatch))
+ try {
+ val length = batchLength(slice)
+ if (length == slice.getRowCount) {
+ writeCopy(slice)
+ } else {
+ val head = slice.slice(0, length)
+ try writeCopy(head) finally head.close()
+ }
+ start += length
+ } finally {
+ slice.close()
+ }
+ }
+ case None =>
+ var i = 0
+ while (i < numBatchRows) {
+ write(batch.getRow(i))
+ i += 1
+ }
+ }
+ }
+
+ /**
+ * Returns the Arrow vectors of the columns of a batch, if its columns can
be copied buffer by
+ * buffer: Arrow vectors with as many values as the batch has rows, whose
fields are physically
+ * congruent with those of `arrowSchema`, though they may differ in the
names, nullability and
+ * metadata of their fields, and in the time zones of timestamps, which the
copy keeps.
+ */
+ private def arrowVectors(batch: ColumnarBatch): Option[Seq[FieldVector]] = {
+ val fields = arrowSchema.getFields.asScala
+ // Without columns, the number of rows of a batch is only kept by the
conversion.
+ if (fields.isEmpty || batch.numCols() != fields.length) {
+ None
+ } else {
+ val vectors = 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
+ }
+ }
+ if (vectors.forall(_.isDefined)) Some(vectors.flatten) else None
+ }
+ }
+
+ /**
+ * Returns the number of rows of the batch that starts with a slice of the
rows of a columnar
+ * batch: all of them, or as a batch of rows ends, up to the first row with
which it holds
+ * `maxBytesPerBatch` bytes, by the same measure of its size.
+ */
+ private def batchLength(slice: VectorSchemaRoot): Int = {
+ val numSliceRows = slice.getRowCount
+ // The size is not computed when the limit is the largest possible one, as
for rows.
+ if (options.maxBytesPerBatch >= Int.MaxValue) {
+ numSliceRows
+ } else {
+ def sizeInBytes(rows: Int): Long =
+ slice.getFieldVectors.asScala.map(_.getBufferSizeFor(rows).toLong).sum
+ // The size grows with the number of rows.
+ var lo = 1
+ var hi = numSliceRows
+ while (lo < hi) {
+ val mid = (lo + hi) >>> 1
+ if (sizeInBytes(mid) >= options.maxBytesPerBatch) hi = mid else lo =
mid + 1
+ }
+ lo
+ }
+ }
+
+ /** Writes a slice of the Arrow vectors of a columnar batch, copied buffer
by buffer. */
+ private def writeCopy(slice: VectorSchemaRoot): Unit = {
+ val length = slice.getRowCount
+ val copy = VectorSchemaRoot.create(slice.getSchema, allocator)
+ try {
+ copy.getFieldVectors.forEach(_.setInitialCapacity(length))
Review Comment:
Done in 59c4691: the slice is unloaded with `VectorUnloader`, each buffer is
copied once into a buffer of its size, and the copies are loaded with
`VectorLoader`, instead of `VectorAppender`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala:
##########
@@ -214,12 +227,119 @@ class ArrowBatchingDataWriter(delegate:
DataWriter[ColumnarBatch], options: Arro
}
}
+ override def writeBatch(batch: ColumnarBatch): Unit = {
+ val numBatchRows = batch.numRows()
+ arrowVectors(batch) match {
+ case Some(vectors) =>
+ flush()
Review Comment:
Done in 59c4691: "write the pending rows of converted batches before a
copied batch" passes an on-heap batch and then an Arrow batch to one
`ArrowBatchingDataWriter`, and checks the order of the rows and the sizes of
the batches. "write the columnar batches of a query" also checks that a write
whose query has a shuffle is an adaptive plan with a `ColumnarToRow`.
##########
docs/sql-native-data-sources.md:
##########
@@ -37,7 +37,9 @@ package, and loads it when a query uses the data source.
Spark adapts it to
Data is exchanged only as columnar batches, through the
[Arrow C data
interface](https://arrow.apache.org/docs/format/CDataInterface.html), so Spark
reads
-the data that a library returns without copying it.
+the data that a library returns without copying it. When the data to write is
columnar too, such
+as the batches of a native scan, Spark writes them without converting them to
rows: it copies
+their Arrow buffers, since a library may keep the data that it writes.
Review Comment:
Done in 59c4691: the sentence now says that it applies when a write gets the
batches of a scan as they are, as in a plain copy from one native data source
to another, that columns of other Arrow types, such as `utf8_view` strings, are
converted value by value, and that with adaptive query execution the data of a
query with a shuffle or a subquery is still converted to rows.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala:
##########
@@ -123,8 +135,10 @@ trait ColumnarStreamReader extends Serializable {
/**
* Writes data to a [[ColumnarDataSource]] for a batch write. Spark serializes
the writer to the
* executors, where `createWriter` is called for each task, and then commits
or aborts the write
- * on the driver. The batches passed to the data writers are backed by Arrow,
and are only valid
- * during the call.
+ * on the driver. The batches passed to the data writers are backed by new
Arrow vectors, with the
+ * Arrow types of the schema, except for the time zones of timestamps, though
not necessarily with
+ * its names and metadata. A batch is only valid during the call, but nothing
else writes to the
Review Comment:
Done in 59c4691: the contract now also mentions the `keysSorted` flag of
maps, nullability, and the Spark types that `ArrowColumnVector` infers.
--
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]
