HyukjinKwon commented on code in PR #59933:
URL: https://github.com/apache/spark/pull/59933#discussion_r4418553754
##########
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:
Done in f2ca665: the contract of `ColumnarWriter`, the doc of
`ArrowBatchingDataWriter`, `write` in the header and the guide now say that
copied values are written as the scan returned them, without the checks of a
conversion, such as the precision of a decimal.
##########
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:
Done in f2ca665: a new test compares the batch sizes of a copied batch with
those of the same rows on the row path, with limits around the size of 2 rows
and windows of 5 rows, and checks the batches of 2 rows across the windows.
##########
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:
Done in f2ca665: a columnar batch without rows no longer flushes the pending
rows, with a test.
##########
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:
Done in f2ca665: the `internalError` case is gone with the one-pass build,
and the copy of the copied columns is a plain `val` in `writeRows` (renamed
from `writeWindow`).
##########
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:
Done in f2ca665: `batchLength` checks the whole range first.
--
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]
