dongjoon-hyun commented on code in PR #59933:
URL: https://github.com/apache/spark/pull/59933#discussion_r4417360994
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala:
##########
@@ -214,12 +227,119 @@ class ArrowBatchingDataWriter(delegate:
DataWriter[ColumnarBatch], options: Arro
}
}
+ override def writeBatch(batch: ColumnarBatch): Unit = {
+ val numBatchRows = batch.numRows()
+ arrowVectors(batch) match {
+ case Some(vectors) =>
+ flush()
+ // The vectors of the batch, which this root does not own: their
producer releases them.
+ val batchRoot =
+ new VectorSchemaRoot(vectors.map(_.getField).asJava, vectors.asJava,
numBatchRows)
+ var start = 0
+ while (start < numBatchRows) {
+ // A slice shares the buffers of the batch. Arrow Java copies a
vector of variable
+ // width, or a list, as if its offsets started at 0, which those of
a slice do.
+ val slice = batchRoot.slice(
+ start, math.min(numBatchRows - start, options.maxRecordsPerBatch))
+ try {
+ val length = batchLength(slice)
+ if (length == slice.getRowCount) {
+ writeCopy(slice)
+ } else {
+ val head = slice.slice(0, length)
+ try writeCopy(head) finally head.close()
+ }
+ start += length
+ } finally {
+ slice.close()
+ }
+ }
+ case None =>
+ var i = 0
+ while (i < numBatchRows) {
+ write(batch.getRow(i))
+ i += 1
+ }
+ }
+ }
+
+ /**
+ * Returns the Arrow vectors of the columns of a batch, if its columns can
be copied buffer by
+ * buffer: Arrow vectors with as many values as the batch has rows, whose
fields are physically
+ * congruent with those of `arrowSchema`, though they may differ in the
names, nullability and
+ * metadata of their fields, and in the time zones of timestamps, which the
copy keeps.
+ */
+ private def arrowVectors(batch: ColumnarBatch): Option[Seq[FieldVector]] = {
+ val fields = arrowSchema.getFields.asScala
+ // Without columns, the number of rows of a batch is only kept by the
conversion.
+ if (fields.isEmpty || batch.numCols() != fields.length) {
+ None
+ } else {
+ val vectors = fields.indices.map { i =>
+ batch.column(i) match {
+ case column: ArrowColumnVector => column.getValueVector match {
+ case vector: FieldVector if vector.getValueCount ==
batch.numRows() &&
+ ArrowUtils.isCompatibleWithDeclaredField(vector.getField,
fields(i)) =>
+ Some(vector)
+ case _ => None
+ }
+ case _ => None
+ }
+ }
+ if (vectors.forall(_.isDefined)) Some(vectors.flatten) else None
Review Comment:
**1. [Medium] One column that is not congruent sends the whole batch through
the conversion**
`arrowVectors` decides for the whole batch: when a single column is not an
Arrow vector congruent with the declared field, every column of the batch goes
through `batch.getRow(i)` and `ArrowWriter`, including the columns that could
be copied buffer by buffer. `spark_datasource.h` lets a library return
`large_utf8` and `utf8_view` strings, and `large_binary` and `binary_view`
binaries, in the data it reads, and `NativePartitionReader` keeps those types
in its vectors. So a copy between two native data sources of a library that
uses the view types never takes the copy path once the table has a string
column, and all of its other columns are converted value by value too.
Could the decision be made per column? For example, a string or binary
column of another Arrow type could be converted on its own into a vector of the
declared type, and then written together with the copied columns, so that the
other columns keep the copy.
##########
docs/sql-native-data-sources.md:
##########
@@ -37,7 +37,9 @@ package, and loads it when a query uses the data source.
Spark adapts it to
Data is exchanged only as columnar batches, through the
[Arrow C data
interface](https://arrow.apache.org/docs/format/CDataInterface.html), so Spark
reads
-the data that a library returns without copying it.
+the data that a library returns without copying it. When the data to write is
columnar too, such
+as the batches of a native scan, Spark writes them without converting them to
rows: it copies
+their Arrow buffers, since a library may keep the data that it writes.
Review Comment:
**7. [Low] The user doc does not say when the batches are still converted**
This reads as if every columnar write skipped the conversion. Two cases
still convert:
- With adaptive query execution, a query with a shuffle or a subquery is an
`AdaptiveSparkPlanExec`, which is not columnar, so the plan keeps
`ColumnarToRow`, as the PR description says.
- Batches whose columns are not Arrow vectors of the declared types, such as
the `large_utf8` or `utf8_view` strings that a library may return, are
converted value by value.
Could the sentence say that it applies when the batches of the scan are
written as they are, for example in a plain copy from one native data source to
another?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala:
##########
@@ -214,12 +227,119 @@ class ArrowBatchingDataWriter(delegate:
DataWriter[ColumnarBatch], options: Arro
}
}
+ override def writeBatch(batch: ColumnarBatch): Unit = {
+ val numBatchRows = batch.numRows()
+ arrowVectors(batch) match {
+ case Some(vectors) =>
+ flush()
+ // The vectors of the batch, which this root does not own: their
producer releases them.
+ val batchRoot =
+ new VectorSchemaRoot(vectors.map(_.getField).asJava, vectors.asJava,
numBatchRows)
+ var start = 0
+ while (start < numBatchRows) {
+ // A slice shares the buffers of the batch. Arrow Java copies a
vector of variable
+ // width, or a list, as if its offsets started at 0, which those of
a slice do.
+ val slice = batchRoot.slice(
+ start, math.min(numBatchRows - start, options.maxRecordsPerBatch))
Review Comment:
**2. [Low] Each batch ended by the byte limit slices the rest of the window
again**
When `maxBytesPerBatch` (64MB by default) ends a batch early, this loop
slices up to `maxRecordsPerBatch` rows, writes the first `length` of them, and
slices again from `start + length`. A slice is not free: for string, binary,
list and map columns whose first offset is not 0, `splitAndTransfer` rewrites
the offsets of every row of the slice (for lists always, and recursively into
their elements), and it copies the validity bits when `start` is not a multiple
of 8. So a batch of the query that holds k times `maxBytesPerBatch` is sliced
about k times, each time from the current row to the end of the window, which
costs about k/2 times the batch instead of once. A native scan whose batches
are much larger than the limit of the write pays this on every batch.
Could the binary search read the sizes from the vectors of `batchRoot`
directly? The size of the data of rows `[start, start + k)` of a string column
is the difference of two offsets, and the same holds recursively for lists, so
only the head of `length` rows would need to be sliced.
##########
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))
Review Comment:
**5. [Low] The copy starts with small buffers and grows them by doubling**
`setInitialCapacity(length)` sizes the data buffer of a string or binary
vector for 8 bytes per value, and the elements of a list for 5 per list.
`VectorAppender` then calls `reallocDataBuffer()` or `reAlloc()` until the data
fits, and each call allocates a buffer twice as large and copies the old one
into it. With strings of 200 bytes, that is about 5 doublings per column and
per slice, and the buffer that the library may keep can be up to twice the size
of its data.
The sizes are known from the slice. One way to copy each buffer once, at its
size, is to unload the slice with `VectorUnloader`, copy each `ArrowBuf` into a
buffer of `allocator`, and load the copy with `VectorLoader`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala:
##########
@@ -123,8 +135,10 @@ trait ColumnarStreamReader extends Serializable {
/**
* Writes data to a [[ColumnarDataSource]] for a batch write. Spark serializes
the writer to the
* executors, where `createWriter` is called for each task, and then commits
or aborts the write
- * on the driver. The batches passed to the data writers are backed by Arrow,
and are only valid
- * during the call.
+ * on the driver. The batches passed to the data writers are backed by new
Arrow vectors, with the
+ * Arrow types of the schema, except for the time zones of timestamps, though
not necessarily with
+ * its names and metadata. A batch is only valid during the call, but nothing
else writes to the
Review Comment:
**8. [Nit] The contract does not mention nullability, or the Spark types of
the columns**
The copy is created from `slice.getSchema`, so it has the fields of the
source. `isCompatibleWithDeclaredField` also lets these differ from the
declared ones in nullability, and in the `keysSorted` flag of a map, which is
part of its Arrow type. And `writeArrow` wraps the vectors in
`ArrowColumnVector`s, which infer their Spark type from the field metadata (the
precision of `TIME` and of nanosecond timestamps, the tags of variant and
geometry structs). So a `ColumnarDataSource` that reads `dataType()` or
`getField.isNullable` can see other values for a copied batch than for a
converted one.
Maybe: "except for the time zones of timestamps and the `keysSorted` flag of
maps, though not necessarily with its names, nullability and metadata, nor the
Spark types that `ArrowColumnVector` infers from them"?
##########
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:
**9. [Nit] `copy.setRowCount(length)` is a no-op**
`VectorSchemaRootAppender.append` ends with
`targetRoot.setRowCount(targetRoot.getRowCount() + delta.getRowCount())`, which
calls `setValueCount` on every vector, null vectors included (I checked this
against Arrow 19.0.0). The copy is created with a row count of 0, so its count
is already `length` here, and the comment that the copy leaves a null vector at
0 is not accurate. Unless I am missing a case, this line and its comment could
go.
##########
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()
Review Comment:
**6. [Low] No test writes rows that are pending before a copied batch**
This `flush()` keeps the order when a task gets a batch that is converted,
whose rows stay in `root`, and then a batch that is copied. But every test
gives a writer only one kind of batch, so removing it, or moving it after the
copy, would not fail a test. A small test that passes a batch of on-heap
vectors and then a batch of Arrow vectors with the same schema to one
`ArrowBatchingDataWriter`, and checks the order of the rows and the sizes of
the batches, would cover it.
A test whose query has a shuffle would also pin down what the PR description
says about adaptive query execution: such a write still converts the batches to
rows.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala:
##########
@@ -214,12 +227,119 @@ class ArrowBatchingDataWriter(delegate:
DataWriter[ColumnarBatch], options: Arro
}
}
+ override def writeBatch(batch: ColumnarBatch): Unit = {
+ val numBatchRows = batch.numRows()
+ arrowVectors(batch) match {
+ case Some(vectors) =>
+ flush()
+ // The vectors of the batch, which this root does not own: their
producer releases them.
+ val batchRoot =
+ new VectorSchemaRoot(vectors.map(_.getField).asJava, vectors.asJava,
numBatchRows)
+ var start = 0
+ while (start < numBatchRows) {
+ // A slice shares the buffers of the batch. Arrow Java copies a
vector of variable
+ // width, or a list, as if its offsets started at 0, which those of
a slice do.
+ val slice = batchRoot.slice(
+ start, math.min(numBatchRows - start, options.maxRecordsPerBatch))
+ try {
+ val length = batchLength(slice)
+ if (length == slice.getRowCount) {
+ writeCopy(slice)
+ } else {
+ val head = slice.slice(0, length)
+ try writeCopy(head) finally head.close()
Review Comment:
**3. [Low] An error from closing can hide the error of the native `write`**
`writeCopy(head)` calls the `write` of the library, and `head.close()` here,
`slice.close()` (L253-L255) and `copy.close()` (L328-L330) run in plain
`finally` blocks. If a close also throws after `write` failed, for example an
`IllegalStateException` from the reference counting of a buffer, its exception
replaces the `NATIVE_DATA_SOURCE_ERROR` that carries the `last_error` of the
library, which is the message a user needs. `runColumnar` in this PR uses
`Utils.tryWithSafeFinally` for the same reason; could these blocks use it too?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala:
##########
@@ -646,8 +669,13 @@ trait V2TableWriteExec
try {
sparkContext.runJob(
rdd,
- (context: TaskContext, iter: Iterator[InternalRow]) =>
- task.run(writerFactory, context, iter, useCommitCoordinator,
writeMetrics),
+ (context: TaskContext, iter: Iterator[Any]) => if (columnar) {
+ DataWritingSparkTask.runColumnar(writerFactory, context,
+ iter.asInstanceOf[Iterator[ColumnarBatch]], useCommitCoordinator,
writeMetrics)
Review Comment:
**4. [Low] The columnar path calls `DataWritingSparkTask.runColumnar`
whatever `writingTask` is**
When `columnar` is true, the batches go to
`DataWritingSparkTask.runColumnar`, and `task` is not used. That is correct
today because the only nodes that override `writingTask`, `ReplaceDataExec` and
`WriteDeltaExec`, are `RowLevelWriteExec`s, which override
`supportsColumnarInput` to false. But nothing ties the two together: a node
that overrides `writingTask` for a write that mixes in `SupportsColumnarInput`
would have its task skipped without an error.
Could `runColumnar` be a method of `WritingSparkTask` that fails by default,
which `DataWritingSparkTask` overrides, and be called here as
`task.runColumnar(...)`? Another option is to make `supportsColumnarInput`
require `writingTask eq DataWritingSparkTask`.
##########
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:
**10. [Nit] The test leaks its allocator, and an export, when it fails**
`allocator.close()` only runs when every assertion above passes, so a
failure leaves this child allocator of `ArrowUtils.rootAllocator` open for the
rest of the JVM. A `try`/`finally` around the body would close it.
Similarly, `TestArrowExport.read` (L748) sets `closed = true` before
`Data.importVectorSchemaRoot`, so if the import throws, `close()` does not
release the `ArrowArray` and the `ArrowSchema`. Setting it after the import,
which moves the data, would keep them released.
--
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]
