HyukjinKwon commented on code in PR #59933:
URL: https://github.com/apache/spark/pull/59933#discussion_r4418555342
##########
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())
Review Comment:
Done in f2ca665: the comment now says which buffers the slice allocates and
that each buffer of the slice is copied.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala:
##########
@@ -404,31 +765,95 @@ object TestColumnarSink {
case class TestCommitMessage(rows: Seq[(Long, String)], batchSizes: Seq[Int])
extends WriterCommitMessage
-class TestDataWriter(fail: Boolean) extends DataWriter[ColumnarBatch] {
+/**
+ * @param keep whether to keep the data of each batch until the commit, as a
native data writer
+ * can, by taking ownership of the Arrow buffers of the batch
+ */
+class TestDataWriter(fail: Boolean, keep: Boolean = false) extends
DataWriter[ColumnarBatch] {
private val rows = ArrayBuffer.empty[(Long, String)]
private val batchSizes = ArrayBuffer.empty[Int]
+ private val kept = ArrayBuffer.empty[TestArrowExport]
override def write(batch: ColumnarBatch): Unit = {
if (fail) {
throw new IllegalStateException("injected failure")
}
assert((0 until
batch.numCols()).forall(batch.column(_).isInstanceOf[ArrowColumnVector]))
batchSizes += batch.numRows()
+ if (keep) {
+ kept += TestArrowExport(batch)
+ } else {
+ read(batch)
+ }
+ }
+
+ private def read(batch: ColumnarBatch): Unit = {
(0 until batch.numRows()).foreach { i =>
rows += ((batch.column(0).getLong(i),
batch.column(1).getUTF8String(i).toString))
}
}
- override def commit(): WriterCommitMessage = TestCommitMessage(rows.toSeq,
batchSizes.toSeq)
+ override def commit(): WriterCommitMessage = {
+ kept.foreach(_.read(read))
+ kept.clear()
+ TestCommitMessage(rows.toSeq, batchSizes.toSeq)
+ }
override def abort(): Unit = {}
- override def close(): Unit = {}
+ override def close(): Unit = {
+ kept.foreach(_.close())
+ kept.clear()
+ }
+}
+
+/** A batch exported to the Arrow C data interface, which keeps its data until
it is read. */
+class TestArrowExport(array: ArrowArray, schema: ArrowSchema) {
+ private var closed = false
+
+ /** Imports the batch, and passes it to `f`. */
+ def read[T](f: ColumnarBatch => T): T = {
+ val root = Data.importVectorSchemaRoot(ArrowUtils.rootAllocator, array,
schema, null)
+ // The import moved the data to the root, and closed the structs.
+ closed = true
Review Comment:
Done in f2ca665: `TestArrowExport.read` imports with `Data.importSchema` and
then `Data.importIntoVectorSchemaRoot`, and marks each struct as closed once
its import is called, so `close()` only releases the struct that is still open.
##########
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.
+ *
+ */
Review Comment:
Done in f2ca665.
--
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]
