spark-test commented on code in PR #59933:
URL: https://github.com/apache/spark/pull/59933#discussion_r4468509210
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala:
##########
@@ -873,6 +942,34 @@ object DataWritingSparkTask extends
WritingSparkTask[DataWriter[InternalRow]] {
writer: DataWriter[InternalRow], iter: java.util.Iterator[InternalRow]):
Unit = {
writer.writeAll(iter)
}
+
+ override def runColumnar(
+ writerFactory: DataWriterFactory,
+ context: TaskContext,
+ iter: Iterator[ColumnarBatch],
+ useCommitCoordinator: Boolean,
+ customMetrics: Map[String, SQLMetric]): DataWritingSparkTaskResult = {
+ val dataWriter = writerFactory.createWriter(context.partitionId(),
context.taskAttemptId())
+ writeAndCommit(dataWriter, context, useCommitCoordinator, customMetrics) {
+ val columnarWriter = dataWriter match {
+ case writer: ColumnarInputDataWriter => writer
+ case writer => throw SparkException.internalError(
+ s"Expected a data writer of columnar batches, but got
${writer.getClass.getName}.")
+ }
+ var numRows = 0L
+ iter.foreach { batch =>
+ // Releases the batch once it is written, as ColumnarToRowExec does
once it returned its
+ // rows: a producer of new column vectors for each batch leaves them
to the consumer.
+ Utils.tryWithSafeFinally {
+ columnarWriter.writeBatch(batch)
+ numRows += batch.numRows()
+ }(batch.closeIfFreeable())
Review Comment:
**Nit (P3):** Nothing in the tests can observe this release. Every scan that
feeds a columnar write reuses its vectors, so `closeIfFreeable()` is a no-op
for them: `ReusedArrowColumnVector` in the native reader and in
`TestArrowBatchReader`, and `OnHeapColumnVector` (a `WritableColumnVector`) in
`TestBatchReader`. The plain `ArrowColumnVector`s of `ColumnarDataSourceSuite`
go directly to `ArrowBatchingDataWriter.writeBatch`, which does not release
them. So deleting this call, skipping it when `writeBatch` throws, or moving it
before `writeBatch` keeps `ColumnarDataSourceSuite` and `NativeDataSourceSuite`
green. For a producer whose batches own their vectors, such as a connector that
returns plain `ArrowColumnVector`s, that regression would leak the vectors of
every batch, at least until the task ends, or free a batch before it is copied.
**Verification:**
- **Behavior:** A columnar write whose scan allocates new freeable vectors
for every batch releases each batch exactly once, after the data writer has
taken its data, and also when the data writer fails on that batch, so that no
memory of the scan stays allocated once the task ends.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsColumnarInput.scala:
##########
@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.internal.connector
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.connector.write.DataWriter
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+/**
+ * An internal mixin for a `Write`, or for the `StreamingWrite` of a
micro-batch write, whose data
+ * writers are [[ColumnarInputDataWriter]]s. When the query of an append or an
overwrite, or of a
+ * micro-batch, produces columnar batches, Spark passes them to the data
writers instead of
Review Comment:
**Nit (P3):** This sentence states the behavior without its conditions, and
so does the scaladoc of `ColumnarWrite` (`ColumnarWrite.scala:61-62`). Spark
passes the batches only when `V2TableWriteExec.writesColumnar` holds, that is
when the query supports columnar output but not rows, and only when the session
has no columnar rules (`Columnar.scala:597`). The scan of a cached DataFrame, a
query that adaptive query execution wraps, and any query in a session whose
extensions have columnar rules are written as rows, as
`ColumnarDataSourceSuite` asserts. A reader of these two docs would instead
expect every columnar query to reach the data writers as batches. The scaladoc
of `writesColumnar` and the user guide already state the conditions.
**Verification:**
- **Inspection:** The documentation of the internal mixin for columnar input
and of the columnar write says that the data writers get the batches of a query
only when the query supports columnar output and not rows and the session has
no columnar rules, names or refers to the queries that are written as rows,
such as the scan of a cached DataFrame, a query that adaptive query execution
wraps and a query in a session whose extensions have columnar rules, and agrees
with the documentation of the write node and with the user guide.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala:
##########
@@ -873,6 +942,34 @@ object DataWritingSparkTask extends
WritingSparkTask[DataWriter[InternalRow]] {
writer: DataWriter[InternalRow], iter: java.util.Iterator[InternalRow]):
Unit = {
writer.writeAll(iter)
}
+
+ override def runColumnar(
+ writerFactory: DataWriterFactory,
+ context: TaskContext,
+ iter: Iterator[ColumnarBatch],
+ useCommitCoordinator: Boolean,
+ customMetrics: Map[String, SQLMetric]): DataWritingSparkTaskResult = {
+ val dataWriter = writerFactory.createWriter(context.partitionId(),
context.taskAttemptId())
+ writeAndCommit(dataWriter, context, useCommitCoordinator, customMetrics) {
+ val columnarWriter = dataWriter match {
+ case writer: ColumnarInputDataWriter => writer
+ case writer => throw SparkException.internalError(
+ s"Expected a data writer of columnar batches, but got
${writer.getClass.getName}.")
+ }
+ var numRows = 0L
+ iter.foreach { batch =>
+ // Releases the batch once it is written, as ColumnarToRowExec does
once it returned its
+ // rows: a producer of new column vectors for each batch leaves them
to the consumer.
+ Utils.tryWithSafeFinally {
+ columnarWriter.writeBatch(batch)
+ numRows += batch.numRows()
Review Comment:
**Nit (P3):** This count is new code for columnar writes, and no test checks
it. `writeWithV2` adds it to the `numOutputRows` metric of the write node, and
a micro-batch reports it as the sink's `numOutputRows` through
`commitProgress`, but no test in `ColumnarDataSourceSuite` or
`NativeDataSourceSuite` asserts either one for a columnar write. Counting
batches instead of rows, or never adding to the count, keeps both suites green,
while the SQL metric and the streaming progress of every columnar write would
be wrong.
**Verification:**
- **Behavior:** A batch write and a micro-batch write that take the columnar
batches of a scan, in several batches per task, report exactly the number of
rows that they wrote: as the number of output rows in the SQL metrics of the
batch write, and as the output rows of the sink in the progress of the
micro-batch.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala:
##########
@@ -380,6 +970,57 @@ class TestBatchReader(lo: Long, hi: Long, schema:
StructType)
override def close(): Unit = if (consumed) batch.close()
}
+/**
+ * Returns the ids [lo, hi) and their names in Arrow batches of at most
`batchSize` rows. It writes
+ * each batch to the same Arrow vectors, over the previous one.
+ */
+/** @param largeVarTypes whether the strings are `large_utf8`, which the
writers convert */
Review Comment:
**Nit (P3):** This `@param` line opens a second doc comment, so Scaladoc
attaches only it to `TestArrowBatchReader`, and the description above it (lines
973-976) is left as a detached doc comment.
**Verification:**
- **Inspection:** The Arrow batch test reader is preceded by a single doc
comment that holds its description followed by the tag of its parameter for
large strings, with no detached doc comment before it.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -75,12 +75,17 @@ static boolean isSupported() {
static final long MICRO_BATCH_READ = 2;
static final long BATCH_WRITE = 4;
static final long STREAMING_WRITE = 8;
- static final long TRUNCATE = 16;
+ static final long BATCH_TRUNCATE = 16;
+ static final long STREAMING_TRUNCATE = 32;
Review Comment:
**Nit (P3):** `BATCH_TRUNCATE` and `STREAMING_TRUNCATE` must equal
`SPARK_DS_BATCH_TRUNCATE` and `SPARK_DS_STREAMING_TRUNCATE` in
`spark_datasource.h`, but no test compares them with the header. The test "the
bindings match the layout of spark_datasource.h" compares only the sizes and
offsets of `NativeLibrary.layout()`, and the test "operations that a data
source does not support" builds its capabilities option from these same Java
constants, which the test library returns unchanged; the only complete-mode
writes pass such options. So swapping 16 and 32, or giving `STREAMING_TRUNCATE`
another value, keeps `NativeDataSourceSuite` green, while a library that
declares only `SPARK_DS_BATCH_TRUNCATE` could have its batch overwrites
rejected and its complete-mode writes accepted, or the reverse.
**Verification:**
- **Compatibility:** The capability values that Spark checks, including the
batch and the streaming truncate bits, and the write modes equal the macros of
the spark_datasource.h that libraries compile against, so that swapping the two
truncate bits or changing either value fails a check that needs only a C++
compiler, not Java 22.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala:
##########
@@ -214,15 +239,240 @@ 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 if (numBatchRows > 0) {
+ // A batch without rows leaves the pending batch of rows as it is.
+ flush()
+ // Each batch starts where the last one ended, so that the batches end
as those of rows do.
+ var start = 0
+ while (start < numBatchRows) {
+ start += writeRows(
+ batch, copied, start, start + math.min(numBatchRows - start,
options.maxRecordsPerBatch))
+ }
+ }
+ }
+
+ /**
+ * Returns, for each column of a batch, its Arrow vector if it can be copied
buffer by buffer, see
+ * `ArrowBatchingDataWriter.canCopy`. 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 ArrowBatchingDataWriter.canCopy(vector, fields(i),
batch.numRows()) =>
+ Some(vector)
+ case _ => None
+ }
+ case _ => None
+ }
+ }
+ }
+ }
+
+ /**
+ * Writes the batch of the rows of a columnar batch that starts at the row
`start`, and returns
+ * its number of rows: 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 for the copied columns without slicing them.
+ *
+ * @param copied the vector of each column of the batch that is copied, see
`copiedVectors`
+ */
+ private def writeRows(
+ batch: ColumnarBatch,
+ copied: Seq[Option[FieldVector]],
+ start: Int,
+ end: Int): Int = {
+ // The columns that are not copied, and the vectors that they are
converted into.
+ val ordinals = copied.indices.filter(copied(_).isEmpty)
+ val converted = ArrayBuffer.empty[FieldVector]
+ Utils.tryWithSafeFinally {
+ val writers = ordinals.map { ordinal =>
+ val vector = arrowSchema.getFields.get(ordinal).createVector(allocator)
+ converted += vector
+ vector.setInitialCapacity(end - start)
+ vector.allocateNew()
+ ArrowWriter.createFieldWriter(vector)
+ }
+ val length = if (writers.isEmpty) {
+ batchLength(copied.flatten, start, end)
+ } else {
+ convertRows(batch, copied.flatten, ordinals.zip(writers), start, end)
+ }
+ writers.foreach(_.finish())
+ writeSlices(copied, converted.toSeq, start, length)
+ length
+ } {
+ AutoCloseables.close(converted.asJava)
+ }
+ }
+
+ /**
+ * Returns the number of rows of the batch of copied columns that starts at
the row `start`, see
+ * `writeRows`.
+ */
+ private def batchLength(copied: Seq[FieldVector], start: Int, end: Int): Int
= {
+ def sizeInBytes(rows: Int): Long =
+ copied.map(ArrowBatchingDataWriter.rangeSize(_, start, start + rows)).sum
+ // The size is not computed when the limit is the largest possible one, as
for rows.
+ if (options.maxBytesPerBatch >= Int.MaxValue ||
+ sizeInBytes(end - start) < options.maxBytesPerBatch) {
+ end - start
+ } else {
+ // 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
+ }
+ }
+
+ /**
+ * Converts the columns of a columnar batch that are not copied with their
writers, by ordinal,
+ * row by row from the row `start` up to the end of the batch, and returns
its number of rows, see
+ * `writeRows`. As for rows, the size of the batch is computed after each
row, so that the
+ * conversion stops at the row with which the batch holds `maxBytesPerBatch`
bytes.
+ */
+ private def convertRows(
+ batch: ColumnarBatch,
+ copied: Seq[FieldVector],
+ writers: Seq[(Int, ArrowFieldWriter)],
+ start: Int,
+ end: Int): Int = {
+ val copiedVectors = copied.toArray
+ val ordinals = writers.map(_._1).toArray
+ val fieldWriters = writers.map(_._2).toArray
+ // The size of the batch once it holds the rows [start, rowEnd).
+ def sizeInBytes(rowEnd: Int): Long = {
+ var size = 0L
+ var i = 0
+ while (i < copiedVectors.length) {
+ size += ArrowBatchingDataWriter.rangeSize(copiedVectors(i), start,
rowEnd)
+ i += 1
+ }
+ i = 0
+ while (i < fieldWriters.length) {
+ size += fieldWriters(i).getSizeInBytes()
+ i += 1
+ }
+ size
+ }
+ // The size is not computed when the limit is the largest possible one, as
for rows.
+ val sized = options.maxBytesPerBatch < Int.MaxValue
+ var next = start
+ var full = false
+ while (next < end && !full) {
+ val row = batch.getRow(next)
+ var i = 0
+ while (i < fieldWriters.length) {
+ fieldWriters(i).write(row, ordinals(i))
+ i += 1
+ }
+ next += 1
+ full = sized && sizeInBytes(next) >= options.maxBytesPerBatch
+ }
+ next - start
+ }
+
+ /**
+ * Writes the rows [start, start + length) of a columnar batch, with the
vectors that its columns
+ * that are not copied were converted into from these rows, in their order.
The copied columns are
+ * sliced, which shares most buffers of the batch, though it allocates new
offsets for lists, and
+ * for strings and binaries unless the offset of the row `start` is 0, and
new validity buffers,
+ * and value buffers for booleans, unless `start` is a multiple of 8. Then
each buffer of the
+ * slice is copied at its 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. The converted
vectors, which
+ * nothing else writes to, are written as they are.
+ */
+ private def writeSlices(
+ copied: Seq[Option[FieldVector]],
+ converted: Seq[FieldVector],
+ start: Int,
+ length: Int): Unit = {
+ // The slices of the copied columns share the buffers of the batch, which
their producer
+ // releases, so they are closed once they are copied.
+ val slices = ArrayBuffer.empty[FieldVector]
+ val copy = Utils.tryWithSafeFinally {
+ copied.flatten.foreach { vector =>
+ // The slice is closed also when the buffers that it allocates fail to
allocate.
+ val pair = vector.getTransferPair(vector.getAllocator)
+ slices += pair.getTo.asInstanceOf[FieldVector]
+ pair.splitAndTransfer(start, length)
Review Comment:
**Nit (P3):** In the tests, only vectors allocated in the JVM reach this
call with a `start` other than 0. `NativeDataSourceSuite` never lowers
`maxRecordsPerBatch` or `maxBytesPerBatch`, and the scans of its test library
return batches of at most 1000 rows (the default `batch_size`), below the
default `maxRecordsPerBatch` of 10000, so vectors imported from a library
through the C data interface are only ever copied whole, from row 0. At a later
start, `splitAndTransfer` allocates new offset or validity buffers and
transfers slices of the imported buffers, which is what every native copy does
when its scan returns batches larger than the write's limits, for example
65536-row batches. A regression there, such as releasing the imported buffers
when a slice is closed or misreading their offsets, would keep both suites
green while such a copy fails or writes corrupt data.
**Verification:**
- **Behavior:** Copying columns imported through the Arrow C data interface,
in batches larger than the record limit and under byte limits of a few rows so
that slices start at non-zero and unaligned rows, writes the same rows and
batch sizes as the row path for fixed-width, boolean, variable-width and nested
columns, leaves the imported batch readable until its producer releases it, and
leaves no memory allocated, in a test that needs no native data source library
and so runs in every build.
--
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]
