dongjoon-hyun commented on code in PR #59933:
URL: https://github.com/apache/spark/pull/59933#discussion_r4418472608
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/Columnar.scala:
##########
@@ -589,6 +590,9 @@ case class ApplyColumnarRulesAndInsertTransitions(
case write: DataWritingCommandExec
if write.cmd.isInstanceOf[V1WriteCommand] &&
conf.plannedWriteEnabled =>
write.child.supportsColumnar
+ // A V2 write whose data writers take columnar batches writes those of
a columnar query.
+ case write: V2TableWriteExec if write.supportsColumnarInput =>
+ write.query.supportsColumnar
Review Comment:
**1. [Medium] A cached DataFrame is now read as batches, which are kept
until the end of the task**
This case makes the query columnar whenever it `supportsColumnar`, also when
it supports rows too. `InMemoryTableScanExec` supports both, and SPARK-37369
made Spark read it as rows when its parent takes rows. Its batches are never
Arrow vectors, so nothing is copied, and the columnar path only adds work and
memory.
For example, with numeric columns only, which
`DefaultCachedBatchSerializer.supportsColumnarOutput` accepts, and a cached
plan without a shuffle, so that adaptive query execution does not apply:
```scala
val df = spark.read.parquet(path).select("id", "value").cache()
df.count()
df.write.format("native_sink").mode("append").save()
```
The query of the write is the `InMemoryTableScanExec`, so `writeWithV2`
calls `executeColumnar()`:
- Before this PR, the scan ran as rows, through
`convertCachedBatchToInternalRow`, without column vectors.
- Now `convertCachedBatchToColumnarBatch` decompresses each cached batch
into new `OnHeapColumnVector`s (`OffHeapColumnVector`s with
`spark.sql.columnVector.offheap.enabled`), and registers a task completion
listener that closes them. `closeIfFreeable` in `runColumnar` is a no-op for
these vectors, so every decompressed batch of the partition stays reachable
until the task ends.
- `writeBatch` then converts these batches value by value, as it does for
any batch that is not Arrow.
Could this case keep rows for a query that also supports them
(`write.query.supportsColumnar && !write.query.supportsRowBased`), with the
same condition in `writeWithV2`? A test that writes a cached DataFrame would
cover it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala:
##########
@@ -214,15 +232,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:
**2. [Low] Small columnar batches are no longer merged into batches of
`maxRecordsPerBatch` rows**
A columnar batch with a copied column first writes the pending rows, and the
batches cut from it end with it, as the doc at L205 says. Before this PR,
`ColumnarToRowExec` and `ArrowWriter` merged the rows of small columnar batches
into batches of up to `maxRecordsPerBatch` rows. So the size of the batches
that a library gets now follows the batch size of the source:
- A native source that returns batches of 500 rows, with the default
`maxRecordsPerBatch` of 10000, now gives 20 times as many calls of `write`,
each with its own export through the C data interface.
- A micro-batch whose source returns batches of a few rows gives a call of
`write` for every few rows.
A sink that turns each batch into a unit, such as a row group, a page or an
IPC record batch, then writes much smaller units. The header only promises an
upper bound, so this stays within it, but the PR description says that the
batches end at `maxRecordsPerBatch` rows or `maxBytesPerBatch` bytes "as the
header promises", and the user doc does not mention smaller batches either.
Could a small copied batch be appended to the pending batch, or be converted as
rows are, until the batch reaches the limits? If the current behavior is
intended, could the user doc say that a library may get smaller batches when
Spark copies them?
##########
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)
Review Comment:
**4. [Nit] The test depends on the defaults of adaptive query execution**
This assertion holds only with `spark.sql.adaptive.enabled=true`, and the
assertions that the plain copies are a `BatchScanExec` (L164, L169) only with
`spark.sql.adaptive.forceApply=false`, but the `withSQLConf` at L161 sets only
`maxRecordsPerBatch`. With adaptive query execution disabled, `shuffled` is a
`ShuffleExchangeExec`, and with `forceApply`, the copies are adaptive plans.
Setting both in the `withSQLConf` would make the test independent of the
defaults:
```scala
withSQLConf(
SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "3",
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.ADAPTIVE_EXECUTION_FORCE_APPLY.key -> "false") {
```
##########
docs/sql-native-data-sources.md:
##########
@@ -37,7 +37,12 @@ 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 a write gets the
batches of a scan as
+they are, as in a plain copy from one native data source to another, Spark
writes them without
+converting them to rows: it copies the buffers of their columns, since a
library may keep the data
+that it writes, and converts value by value the columns of other Arrow types
than those of the
+schema, such as `utf8_view` strings. With adaptive query execution, the data
of a query with a
+shuffle or a subquery is still converted to rows.
Review Comment:
**3. [Nit] A query with a shuffle is converted to rows with or without
adaptive query execution**
This is about the change for
[7](https://github.com/apache/spark/pull/59933#discussion_r4417361040). The
sentence reads as if adaptive query execution caused the conversion of a query
with a shuffle. But `ShuffleExchangeExec` is not columnar, so Spark also
inserts `ColumnarToRow` before the shuffle when adaptive query execution is
disabled, and the write gets rows. What adaptive query execution adds is that
the adaptive plan as a whole is not columnar. The PR description has the same
framing ("Under adaptive query execution, ... the batches only reach the write
when the query has no shuffle"). For example:
```suggestion
schema, such as `utf8_view` strings. The data of a query in which Spark also
filters, projects or
shuffles the rows is still converted to rows, and with adaptive query
execution, so is the data of
a query with a subquery.
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala:
##########
@@ -701,18 +729,49 @@ trait WritingSparkTask[W <: DataWriter[InternalRow]]
extends Logging with Serial
iter: Iterator[InternalRow],
useCommitCoordinator: Boolean,
customMetrics: Map[String, SQLMetric]): DataWritingSparkTaskResult = {
+ val dataWriter =
+ writerFactory.createWriter(context.partitionId(),
context.taskAttemptId()).asInstanceOf[W]
+
+ val iterWithMetrics = IteratorWithMetrics(iter, dataWriter, customMetrics)
+
+ writeAndCommit(dataWriter, context, useCommitCoordinator, customMetrics) {
+ write(dataWriter, iterWithMetrics)
+ iterWithMetrics.count
+ }
+ }
+
+ /**
+ * Writes the columnar batches of a query, for a write whose data writers
take them (see
+ * [[SupportsColumnarInput]]). Only the task that writes the rows as they
are does.
+ */
+ def runColumnar(
+ writerFactory: DataWriterFactory,
+ context: TaskContext,
+ iter: Iterator[ColumnarBatch],
+ useCommitCoordinator: Boolean,
+ customMetrics: Map[String, SQLMetric]): DataWritingSparkTaskResult = {
+ throw SparkException.internalError(s"${getClass.getName} does not write
columnar batches.")
+ }
+
+ /**
+ * Writes the data of the task with `write`, which returns the number of
rows that it wrote, and
+ * commits the data writer, once the output commit coordinator authorizes it
if
+ * `useCommitCoordinator` is true. Aborts the data writer if either fails,
and closes it.
+ */
+ protected def writeAndCommit(
+ dataWriter: DataWriter[_],
+ context: TaskContext,
+ useCommitCoordinator: Boolean,
+ customMetrics: Map[String, SQLMetric])(write: => Long):
DataWritingSparkTaskResult = {
Review Comment:
**5. [Nit] The block of `writeAndCommit` has the name of the method `write`**
The by-name parameter `write` has the name of the abstract method
`write(writer, iter)` of this trait, so `val numRows = write` at L774 evaluates
the block, while the block that `run` passes calls the method. It works, but a
reader can take L774 for a call of the method, and an edit of `writeAndCommit`
that means to call the method gets a confusing compile error. Another name,
such as `writeData`, would avoid this, with `val numRows = writeData` at L774:
```suggestion
customMetrics: Map[String, SQLMetric])(writeData: => Long):
DataWritingSparkTaskResult = {
```
--
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]
