HyukjinKwon commented on code in PR #59933:
URL: https://github.com/apache/spark/pull/59933#discussion_r4417645316
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala:
##########
@@ -214,12 +227,119 @@ class ArrowBatchingDataWriter(delegate:
DataWriter[ColumnarBatch], options: Arro
}
}
+ override def writeBatch(batch: ColumnarBatch): Unit = {
+ val numBatchRows = batch.numRows()
+ arrowVectors(batch) match {
+ case Some(vectors) =>
+ flush()
+ // The vectors of the batch, which this root does not own: their
producer releases them.
+ val batchRoot =
+ new VectorSchemaRoot(vectors.map(_.getField).asJava, vectors.asJava,
numBatchRows)
+ var start = 0
+ while (start < numBatchRows) {
+ // A slice shares the buffers of the batch. Arrow Java copies a
vector of variable
+ // width, or a list, as if its offsets started at 0, which those of
a slice do.
+ val slice = batchRoot.slice(
+ start, math.min(numBatchRows - start, options.maxRecordsPerBatch))
+ try {
+ val length = batchLength(slice)
+ if (length == slice.getRowCount) {
+ writeCopy(slice)
+ } else {
+ val head = slice.slice(0, length)
+ try writeCopy(head) finally head.close()
+ }
+ start += length
+ } finally {
+ slice.close()
+ }
+ }
+ case None =>
+ var i = 0
+ while (i < numBatchRows) {
+ write(batch.getRow(i))
+ i += 1
+ }
+ }
+ }
+
+ /**
+ * Returns the Arrow vectors of the columns of a batch, if its columns can
be copied buffer by
+ * buffer: Arrow vectors with as many values as the batch has rows, whose
fields are physically
+ * congruent with those of `arrowSchema`, though they may differ in the
names, nullability and
+ * metadata of their fields, and in the time zones of timestamps, which the
copy keeps.
+ */
+ private def arrowVectors(batch: ColumnarBatch): Option[Seq[FieldVector]] = {
+ val fields = arrowSchema.getFields.asScala
+ // Without columns, the number of rows of a batch is only kept by the
conversion.
+ if (fields.isEmpty || batch.numCols() != fields.length) {
+ None
+ } else {
+ val vectors = fields.indices.map { i =>
+ batch.column(i) match {
+ case column: ArrowColumnVector => column.getValueVector match {
+ case vector: FieldVector if vector.getValueCount ==
batch.numRows() &&
+ ArrowUtils.isCompatibleWithDeclaredField(vector.getField,
fields(i)) =>
+ Some(vector)
+ case _ => None
+ }
+ case _ => None
+ }
+ }
+ if (vectors.forall(_.isDefined)) Some(vectors.flatten) else None
+ }
+ }
+
+ /**
+ * Returns the number of rows of the batch that starts with a slice of the
rows of a columnar
+ * batch: all of them, or as a batch of rows ends, up to the first row with
which it holds
+ * `maxBytesPerBatch` bytes, by the same measure of its size.
+ */
+ private def batchLength(slice: VectorSchemaRoot): Int = {
+ val numSliceRows = slice.getRowCount
+ // The size is not computed when the limit is the largest possible one, as
for rows.
+ if (options.maxBytesPerBatch >= Int.MaxValue) {
+ numSliceRows
+ } else {
+ def sizeInBytes(rows: Int): Long =
+ slice.getFieldVectors.asScala.map(_.getBufferSizeFor(rows).toLong).sum
+ // The size grows with the number of rows.
+ var lo = 1
+ var hi = numSliceRows
+ while (lo < hi) {
+ val mid = (lo + hi) >>> 1
+ if (sizeInBytes(mid) >= options.maxBytesPerBatch) hi = mid else lo =
mid + 1
+ }
+ lo
+ }
+ }
+
+ /** Writes a slice of the Arrow vectors of a columnar batch, copied buffer
by buffer. */
+ private def writeCopy(slice: VectorSchemaRoot): Unit = {
+ val length = slice.getRowCount
+ val copy = VectorSchemaRoot.create(slice.getSchema, allocator)
+ try {
+ copy.getFieldVectors.forEach(_.setInitialCapacity(length))
+ copy.allocateNew()
+ VectorSchemaRootAppender.append(false, copy, slice)
+ // Also sets the number of values of a null vector, which the copy
leaves at 0.
+ copy.setRowCount(length)
Review Comment:
Done in 59c4691: the copy now goes through `VectorLoader`, which sets the
row count, so that line and its comment are gone.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala:
##########
@@ -128,6 +138,181 @@ 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)
+ 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)
+
+ // 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, and their rows. The data writer keeps the data of the batches,
and reads it once
+ // `overwrite` overwrote the vectors of the batch.
+ def writeAndRead(
+ batch: ColumnarBatch,
+ overwrite: () => Unit,
+ schema: StructType = schema,
+ maxRecordsPerBatch: Int = 10,
+ maxBytesPerBatch: Long = Int.MaxValue): (Seq[Int], Seq[String],
Seq[InternalRow]) = {
+ val sizes = ArrayBuffer.empty[Int]
+ val timeZones = ArrayBuffer.empty[String]
+ val kept = ArrayBuffer.empty[TestArrowExport]
+ val dataWriter = new DataWriter[ColumnarBatch] {
+ override def write(batch: ColumnarBatch): Unit = {
+ sizes += batch.numRows()
+ timeZones +=
batch.column(schema.fieldIndex("c14")).asInstanceOf[ArrowColumnVector]
+
.getValueVector.getField.getType.asInstanceOf[ArrowType.Timestamp].getTimezone
+ 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, 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())
+ 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"),
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"), 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"), expected)))
+ }
+ withArrowBatch(largeVarTypes = true) { (batch, overwrite) =>
+ // Strings and binaries of other Arrow types than those of the schema
are converted.
+ assert(writeAndRead(batch, overwrite, maxRecordsPerBatch = 3) ==
+ ((Seq(3, 3, 1), Seq.fill(3)("America/Los_Angeles"), expected)))
+ }
+
+ // On-heap vectors are converted too. They do not support the null type.
+ val onHeapSchema = StructType(schema.filterNot(_.name == "c1"))
+ val onHeapRows = rows.map(row =>
InternalRow.fromSeq(row.toSeq(schema).patch(1, Nil, 1)))
+ val vectors: Seq[WritableColumnVector] =
+ OnHeapColumnVector.allocateColumns(onHeapRows.length, onHeapSchema).toSeq
+ try {
+ val converter = new RowToColumnConverter(onHeapSchema)
+ onHeapRows.foreach(converter.convert(_, vectors.toArray))
+ val batch = new ColumnarBatch(vectors.map(v => v: ColumnVector).toArray,
onHeapRows.length)
+ assert(writeAndRead(batch, () => (), onHeapSchema, maxRecordsPerBatch =
3) ==
+ ((Seq(3, 3, 1), Seq.fill(3)("America/Los_Angeles"),
+ onHeapRows.map(UnsafeProjection.create(onHeapSchema)(_).copy()))))
+ } finally {
+ vectors.foreach(_.close())
+ }
+ allocator.close()
Review Comment:
Done in 59c4691: the body of the test is in a `try`/`finally` that closes
the allocator, and `TestArrowExport.read` marks itself closed only after the
import.
--
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]
