dongjoon-hyun commented on code in PR #59933:
URL: https://github.com/apache/spark/pull/59933#discussion_r4417897708
##########
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:
**1. [Medium] Data kept past the close of a data writer is still charged to
the allocator that `close()` closes**
The new contract says that a data writer may take ownership of the buffers
of a batch to keep its data, and `spark_datasource.h` lets a library move an
input array to take ownership, with no bound on how long it keeps it. But every
buffer that the data writer gets, whether a copy, a converted vector or a batch
of rows, is allocated from this child allocator, which `close()` (L451-L458)
closes right after `delegate.close()`.
So a library that keeps the arrays past `close_data_writer`, for example an
in-memory sink that appends them to a table that outlives the task, leaves
buffers outstanding when the allocator closes:
- `allocator.close()` throws `IllegalStateException` ("Memory was leaked by
query"), which `Utils.tryLogNonFatalError` logs as an error, for every task.
- When the library later releases the array, each release of these buffers
starts with `allocator.assertOpen()` (`BufferLedger.decrement`, and
`BaseAllocator.dissociateLedger` for the last one; I checked this in the Arrow
19.0.0 classes). With Java assertions enabled, as in the Spark tests, that
throws `IllegalStateException` inside the release callback.
The test data writers release what they keep in `commit` or `close`, so the
suites do not hit this. Could the contract and the header say that the data
must be released by the time the data writer is closed? Or, if a library should
be able to keep it longer, could these buffers come from an allocator that
outlives the task, such as `ArrowUtils.rootAllocator`, so that closing the data
writer does not depend on them?
##########
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:
**2. [Low] The converted vectors leak when a later column fails to convert**
`converted` is built by `copied.indices.map` before the `tryWithSafeFinally`
that closes it. So if `convertColumn` throws for a column, the vectors already
converted for the columns before it are never closed: `convertColumn` closes
only its own vector. It can throw on the data, for example
`ArrowColumnVector.getDecimal` when a value exceeds the declared precision, or
with an `OutOfMemoryException` from the allocation. `close()` then logs a leak,
and the memory stays charged to `ArrowUtils.rootAllocator`.
Building the columns inside the `tryWithSafeFinally`, and closing what was
converted so far, would avoid it, for example:
```scala
val converted = ArrayBuffer.empty[FieldVector]
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, windowStart, windowEnd)
converted += vector
WindowColumn(vector, windowStart, copied = false)
}
}
// ... the loop over the batches of the window, as now
} {
converted.foreach(_.close())
}
```
This also removes the `internalError` case of L257-L258 (see item 8).
##########
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:
**3. [Low] `copyRoot` leaks the buffers already copied when an allocation
fails**
The `map` that copies the buffers runs before the `try` whose `finally`
closes `buffers`. If `allocator.buffer(size)` throws for one buffer, for
example with an `OutOfMemoryException` while copying a wide batch, `buffers` is
never assigned and the copies made before it stay allocated in the child
allocator, which `close()` then reports as leaked.
Collecting the copies into a buffer inside the `try` would release them on
any failure:
```scala
val buffers = ArrayBuffer.empty[ArrowBuf]
try {
unloaded.getBuffers.asScala.foreach { buffer =>
val size = buffer.readableBytes()
val copied = allocator.buffer(size)
buffers += copied
copied.setBytes(0, buffer, buffer.readerIndex(), size)
copied.writerIndex(size)
}
// ... the ArrowRecordBatch and the VectorLoader, as now
} finally {
buffers.foreach(_.close())
}
```
##########
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:
**4. [Low] The ends of the windows also end the batches cut by the byte
limit**
The windows are fixed at multiples of `maxRecordsPerBatch` from the start of
the columnar batch. When `maxBytesPerBatch` ends a batch inside a window, the
next batch can only go to the end of that window, not to `start +
maxRecordsPerBatch` as a batch of rows can. For example, with
`maxRecordsPerBatch` at 3 and rows that reach `maxBytesPerBatch` every 2 rows:
- rows give batches of 2, 2, 2, ...
- a copied columnar batch gives 2, 1, 2, 1, ..., since the windows [0, 3)
and [3, 6) are written as [0, 2), [2, 3), [3, 5) and [5, 6).
That is half as many batches again, and half of them are half the size. The
class doc (L204-L205) says that the batches end as those of rows do, except at
the end of the columnar batch. Could a window start where the previous batch
ended, or could the doc say that the windows end the batches too?
##########
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:
**10. [Nit] The slice allocates some buffers that the copy then copies
again**
This is about the change for
[5](https://github.com/apache/spark/pull/59933#discussion_r4417361029).
`batchRoot.slice(start, length)` shares most buffers, but it allocates and
writes new offsets for every list, and for every string or binary column whose
first offset is not 0, and new validity bits when `start` is not a multiple of
8. `copyRoot` then copies these new buffers again, and it also copies the
validity buffers that `VectorLoader` replaces when a column has no nulls or
only nulls. So the comment at L352-L353, that the buffers are copied once each,
holds only for the buffers that the slice shares.
The cost is small next to the data, so updating the comment may be enough.
Otherwise, the buffers that the slice allocated could be moved into the copy
instead of copied.
##########
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()
Review Comment:
**7. [Nit] An empty batch whose columns could be copied ends the pending
batch of rows**
`flush()` runs even when `numBatchRows` is 0, so a batch without rows, whose
columns are congruent Arrow vectors, writes the pending rows as a smaller batch
although nothing is copied after them. For example, on-heap batches of 2 rows
with an empty Arrow batch between them give batches of 2 rows instead of up to
`maxRecordsPerBatch`. Returning early from `writeBatch` when `numBatchRows` is
0 would avoid it.
##########
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)
Review Comment:
**5. [Low] A copied column skips the checks that a converted value goes
through**
A copied column forwards the buffers as they are, while the same values on
the conversion path go through `ArrowColumnVector` and `ArrowWriter`, which
check their domain: `getDecimal` throws when a value exceeds the declared
precision, `DecimalWriter` writes a null when `changePrecision` fails, and the
nanosecond timestamps check `DATETIME_OVERFLOW`. Spark does not check these
values in the arrays that a library returns.
So for the same data, whether a write fails depends on the plan. If a native
reader returns a `DECIMAL(10, 2)` array that holds 10^12, `INSERT INTO sink
SELECT * FROM source` copies it into the sink, while the same query with a
`.repartition(2)` (an adaptive plan, so rows) fails with a precision error. The
doc of `ArrowUtils.isCompatibleWithDeclaredField` also says that the row-based
path is where these checks are reinstated.
If this is intended, could the contract say that the values of copied
columns are not checked? Otherwise, the decimal columns could be checked, or
converted.
##########
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:
**11. [Nit] After a failed import, `close()` throws and hides the error of
the import**
This is about the change for
[10](https://github.com/apache/spark/pull/59933#discussion_r4417361062).
`Data.importVectorSchemaRoot` releases and closes the `ArrowSchema` in a
`finally` (`Data.importField`), and then moves and closes the `ArrowArray`
before it imports its buffers. So when the import throws, one or both structs
are already closed, but `closed` stays false, and `close()` calls `release()`
on a closed struct, which throws "ArrowSchema is already closed" (or the same
for the `ArrowArray`). Since `close()` runs from the `close()` of the data
writer in a `finally`, that exception replaces the error of the import in the
test report.
Importing the schema and the array in two steps, `Data.importSchema` and
then `Data.importIntoVectorSchemaRoot`, and marking each struct as closed once
the step that consumes it is called, would let `close()` release only the
struct that is still open, and keep the original error.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala:
##########
@@ -128,6 +140,244 @@ class ColumnarDataSourceSuite extends QueryTest with
SharedSparkSession {
}
}
+ test("write the columnar batches of a query") {
+ // Runs a write, and returns the plan of the query that it writes.
+ def writtenPlan(write: => Unit): SparkPlan = {
+ QueryTest.withQueryExecutionsCaptured(spark)(write).flatMap { qe =>
+ qe.executedPlan.collectFirst { case writeExec: V2TableWriteExec =>
writeExec.query }
+ }.head
+ }
+ // Writes the ids [0, 10) and their names, read in Arrow batches of
`arrowBatchSize` rows, or in
+ // a batch of on-heap vectors for each of the two partitions of the scan.
+ def copy(table: String, arrowBatchSize: Int = 0, keep: Boolean = false):
SparkPlan = {
+ writtenPlan {
+ spark.read.format(format).option("end", "10").option("arrowBatchSize",
arrowBatchSize)
+ .load()
+ .write.format(format).option("table", table).option("keep",
keep).mode("append").save()
+ }
+ }
+ val expected = (0L until 10L).map(id => (id, s"name$id"))
+
+ withSQLConf(SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "3") {
+ // The data writers get the batches of the scan, without a transition to
rows. On-heap
+ // vectors are converted to Arrow value by value, into batches that end
as those of rows do.
+ assert(copy("columnar_write").isInstanceOf[BatchScanExec])
+ assert(TestColumnarSink.rows("columnar_write") == expected)
+ assert(TestColumnarSink.batchSizes("columnar_write") == Seq(3, 2, 3, 2))
+
+ // Arrow batches are copied buffer by buffer, each in batches of at most
3 rows.
+ assert(copy("arrow_write", arrowBatchSize =
2).isInstanceOf[BatchScanExec])
+ assert(TestColumnarSink.rows("arrow_write") == expected)
+ assert(TestColumnarSink.batchSizes("arrow_write") == Seq(2, 2, 1, 2, 2,
1))
+ copy("large_arrow_write", arrowBatchSize = 5)
+ assert(TestColumnarSink.rows("large_arrow_write") == expected)
+ assert(TestColumnarSink.batchSizes("large_arrow_write") == Seq(3, 2, 3,
2))
+
+ // A data writer may keep the data of the batches that it gets, though
the scan overwrites
+ // its Arrow vectors with each batch.
+ copy("kept_arrow_write", arrowBatchSize = 2, keep = true)
+ assert(TestColumnarSink.rows("kept_arrow_write") == expected)
+
+ // A query that does not produce columnar batches is written as rows.
+ val plan = writtenPlan {
+ spark.read.format(format).option("end", "10").load()
+ .selectExpr("id", "upper(name) AS name")
+ .write.format(format).option("table",
"row_write").mode("append").save()
+ }
+ assert(plan.exists(_.isInstanceOf[ColumnarToRowExec]), plan)
+
+ // With adaptive query execution, a query with a shuffle is an adaptive
plan, which is not
+ // columnar, so its batches are converted to rows.
+ val shuffled = writtenPlan {
+ spark.read.format(format).option("end", "10").option("arrowBatchSize",
2).load()
+ .repartition(2)
+ .write.format(format).option("table",
"shuffled_write").mode("append").save()
+ }
+ assert(shuffled.isInstanceOf[AdaptiveSparkPlanExec], shuffled)
+ assert(find(shuffled)(_.isInstanceOf[ColumnarToRowExec]).isDefined,
shuffled)
+ assert(TestColumnarSink.rows("shuffled_write").sorted == expected)
+ assert(TestColumnarSink.rows("row_write") == expected.map { case (id,
name) =>
+ (id, name.toUpperCase(Locale.ROOT))
+ })
+ }
+ }
+
+ test("copy the Arrow columns of every type, or convert the columns of a
batch") {
+ val nullable = Seq(
+ "id % 2 = 0", "CAST(id AS TINYINT)", "CAST(id AS SMALLINT)", "CAST(id AS
INT)", "id",
+ "CAST(id AS FLOAT)", "CAST(id AS DOUBLE)", "CAST(id AS DECIMAL(10, 2))",
+ "CAST(id AS DECIMAL(38, 10))", "CAST(id AS STRING)", "CAST(CAST(id AS
STRING) AS BINARY)",
+ "DATE_ADD(DATE'2020-01-01', CAST(id AS INT))", "TIMESTAMP_SECONDS(id)",
+ "CAST(TIMESTAMP_SECONDS(id) AS TIMESTAMP_NTZ)",
"MAKE_YM_INTERVAL(CAST(id AS INT), 1)",
+ "MAKE_DT_INTERVAL(CAST(id AS INT))", "MAKE_INTERVAL(0, 1, 0, CAST(id AS
INT))",
+ "PARSE_JSON(CAST(id AS STRING))", "ARRAY(id, NULL)", "ARRAY(CAST(id AS
STRING))",
+ "MAP(CAST(id AS STRING), id)", "MAP(id, ARRAY(CAST(id AS STRING)))",
+ "NAMED_STRUCT('a', id, 'b', CAST(id AS STRING))",
"ARRAY(NAMED_STRUCT('x', id))")
+ // c0 is not nullable, c1 has the null type, and the other columns have
nulls.
+ val df = spark.range(7).selectExpr(
+ Seq("id AS c0", "NULL AS c1") ++ nullable.zipWithIndex.map { case (expr,
i) =>
+ s"IF(id % 3 = 1, NULL, $expr) AS c${i + 2}"
+ }: _*)
+ val schema = df.schema
+ val rows = df.queryExecution.toRdd.map(_.copy()).collect().toSeq
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("test", 0,
Long.MaxValue)
+ try {
+ // Writes a batch with the session time zone America/Los_Angeles, and
returns the sizes of
+ // the batches that the data writer got, the time zone of their
timestamp column c14, which
+ // a copy keeps, the Arrow types of their string column c11, and their
rows. The data writer
+ // keeps the data of the batches, and reads it once `overwrite`
overwrote the batch.
+ def writeAndRead(
+ batch: ColumnarBatch,
+ overwrite: () => Unit,
+ schema: StructType = schema,
+ maxRecordsPerBatch: Int = 10,
+ maxBytesPerBatch: Long = Int.MaxValue)
+ : (Seq[Int], Seq[String], Seq[ArrowType], Seq[InternalRow]) = {
+ val sizes = ArrayBuffer.empty[Int]
+ val timeZones = ArrayBuffer.empty[String]
+ val stringTypes = ArrayBuffer.empty[ArrowType]
+ val kept = ArrayBuffer.empty[TestArrowExport]
+ def field(batch: ColumnarBatch, name: String) =
batch.column(schema.fieldIndex(name))
+ .asInstanceOf[ArrowColumnVector].getValueVector.getField
+ val dataWriter = new DataWriter[ColumnarBatch] {
+ override def write(batch: ColumnarBatch): Unit = {
+ sizes += batch.numRows()
+ timeZones += field(batch,
"c14").getType.asInstanceOf[ArrowType.Timestamp].getTimezone
+ stringTypes += field(batch, "c11").getType
+ kept += TestArrowExport(batch)
+ }
+ override def commit(): WriterCommitMessage = null
+ override def abort(): Unit = {}
+ override def close(): Unit = kept.foreach(_.close())
+ }
+ val writer = new ArrowBatchingDataWriter(dataWriter,
ArrowBatchingOptions(
+ schema, "America/Los_Angeles", maxRecordsPerBatch, maxBytesPerBatch))
+ try {
+ writer.writeBatch(batch)
+ writer.commit()
+ overwrite()
+ val toUnsafe = UnsafeProjection.create(schema)
+ val written = kept.toSeq.flatMap { exported =>
+
exported.read(_.rowIterator().asScala.map(toUnsafe(_).copy()).toSeq)
+ }
+ kept.clear()
+ (sizes.toSeq, timeZones.toSeq, stringTypes.toSeq, written)
+ } finally {
+ writer.close()
+ }
+ }
+
+ // Runs `f` with a batch of Arrow vectors with the time zone UTC, and a
function that
+ // overwrites the vectors with the rows in reverse order.
+ def withArrowBatch(largeVarTypes: Boolean)(f: (ColumnarBatch, () =>
Unit) => Unit): Unit = {
+ val root = VectorSchemaRoot.create(
+ ArrowUtils.toArrowSchema(
+ schema, "UTC", errorOnDuplicatedFieldNames = true, largeVarTypes =
largeVarTypes),
+ allocator)
+ try {
+ val arrowWriter = ArrowWriter.create(root)
+ rows.foreach(arrowWriter.write)
+ arrowWriter.finish()
+ val batch = new ColumnarBatch(
+ root.getFieldVectors.asScala.map(new ArrowColumnVector(_):
ColumnVector).toArray,
+ rows.length)
+ f(batch, () => {
+ arrowWriter.reset()
+ rows.reverse.foreach(arrowWriter.write)
+ arrowWriter.finish()
+ })
+ } finally {
+ root.close()
+ }
+ }
+
+ val expected = rows.map(UnsafeProjection.create(schema)(_).copy())
+ val utf8 = ArrowType.Utf8.INSTANCE
+ Seq(false, true).foreach { largeVarTypes =>
+ withArrowBatch(largeVarTypes) { (batch, _) =>
+ // The size of a range of rows, which the batches of a copy end at,
is the size of a slice
+ // of these rows.
+ (0 until batch.numCols()).foreach { i =>
+ val vector =
batch.column(i).asInstanceOf[ArrowColumnVector].getValueVector
+ Seq((0, 7), (1, 4), (3, 7), (6, 7), (2, 2)).foreach { case (start,
end) =>
+ val pair = vector.getTransferPair(vector.getAllocator)
+ pair.splitAndTransfer(start, end - start)
+ val sliceSize =
+ try pair.getTo.getBufferSizeFor(end - start) finally
pair.getTo.close()
+ assert(ArrowBatchingDataWriter.rangeSize(vector, start, end) ==
sliceSize,
+ s"${vector.getField} [$start, $end)")
+ }
+ }
+ }
+ }
+ withArrowBatch(largeVarTypes = false) { (batch, overwrite) =>
+ // The columns are Arrow vectors with the types of the schema, so they
are copied.
+ assert(writeAndRead(batch, overwrite) == ((Seq(7), Seq("UTC"),
Seq(utf8), expected)))
+ }
+ withArrowBatch(largeVarTypes = false) { (batch, overwrite) =>
+ // In slices of at most 3 rows, which start in the middle of the bytes
of validity.
+ assert(writeAndRead(batch, overwrite, maxRecordsPerBatch = 3) ==
+ ((Seq(3, 3, 1), Seq.fill(3)("UTC"), Seq.fill(3)(utf8), expected)))
+ }
+ withArrowBatch(largeVarTypes = false) { (batch, overwrite) =>
+ // A copied batch also ends after the first row with which it holds
the maximum size.
+ assert(writeAndRead(batch, overwrite, maxBytesPerBatch = 1) ==
+ ((Seq.fill(7)(1), Seq.fill(7)("UTC"), Seq.fill(7)(utf8), expected)))
Review Comment:
**6. [Low] The byte limit of a copied batch is only tested at 1 byte**
With `maxBytesPerBatch = 1`, every search of `batchLength` returns 1, since
`lo` starts at 1. So the tests do not check where a batch ends when the limit
is reached in the middle of a window, nor that it ends where a batch of rows
would. For example, comparing with `>` instead of `>=` at L345 of
`ColumnarWrite.scala` still passes every test.
A case whose limit is the size of a few rows, which `rangeSize` can compute
for the columns of the batch, compared with the batch sizes that the same rows
get on the row path (an on-heap batch, as at the end of this test), would cover
it. With a `maxRecordsPerBatch` that is not a multiple of that number of rows,
it would also show the difference of item 4.
##########
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
Review Comment:
**9. [Nit] The search runs even when the whole window fits**
With the default `maxBytesPerBatch` of 64MB, most windows are smaller than
the limit, yet each batch runs the whole binary search, about 14 probes over
all the columns for the 10000 rows of the default `maxRecordsPerBatch`, before
it returns `end - start`. Checking the whole range first would take one probe
in that case:
```scala
if (sizeInBytes(end - start) < options.maxBytesPerBatch) {
end - start
} else {
// the search, as now
}
```
##########
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 {
Review Comment:
**8. [Nit] The branch without copied columns, and the column that is neither
copied nor converted, cannot happen**
`writeBatch` calls `writeWindow` only when at least one column is copied, so
`copiedColumns` is never empty here, and this branch with `Iterator.empty`
never runs. Similarly, `copied` and `converted` are complementary, so the
`internalError` case of L257-L258 is unreachable. Building the columns in one
pass, as in the example of item 2, would remove the second, and the first could
become a plain `val copy = ...`.
##########
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:
**12. [Nit] An empty line ends the doc of `batchLength`**
```suggestion
*/
```
--
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]
