HyukjinKwon commented on code in PR #59933:
URL: https://github.com/apache/spark/pull/59933#discussion_r4417641978
##########
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
Review Comment:
Done in 59c4691: each column is copied or converted on its own. The
congruent Arrow columns are copied, and the others, such as `utf8_view`
strings, are converted into vectors of the declared type for a window of at
most `maxRecordsPerBatch` rows, which are sliced into the batches with the
copied columns.
##########
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))
Review Comment:
Done in 59c4691: `ArrowBatchingDataWriter.rangeSize` computes the size of a
range of rows from the offsets of the vectors, at every level, as
`getBufferSizeFor` measures a slice of it, and only the head that is written is
sliced. A test checks it against the size of a slice for every type.
##########
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()
Review Comment:
Done in 59c4691: the closes after the write of the library, and in `flush`,
use `Utils.tryWithSafeFinally`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala:
##########
@@ -646,8 +669,13 @@ trait V2TableWriteExec
try {
sparkContext.runJob(
rdd,
- (context: TaskContext, iter: Iterator[InternalRow]) =>
- task.run(writerFactory, context, iter, useCommitCoordinator,
writeMetrics),
+ (context: TaskContext, iter: Iterator[Any]) => if (columnar) {
+ DataWritingSparkTask.runColumnar(writerFactory, context,
+ iter.asInstanceOf[Iterator[ColumnarBatch]], useCommitCoordinator,
writeMetrics)
Review Comment:
Done in 59c4691: `runColumnar` is a method of `WritingSparkTask` that fails
with an internal error unless `DataWritingSparkTask` overrides it, and
`writeWithV2` calls `task.runColumnar`.
--
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]
