HyukjinKwon commented on code in PR #59933:
URL: https://github.com/apache/spark/pull/59933#discussion_r4418950903
##########
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:
Done in b3afefd: the write takes the batches only when the query supports
columnar output and not rows (`V2TableWriteExec.writesColumnar`), in
`ApplyColumnarRulesAndInsertTransitions` and `writeWithV2`. A test writes a
cached DataFrame, checks that its `InMemoryTableScanExec` supports both, and
that it is written as rows.
##########
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:
Done in b3afefd: kept the copy per batch, and documented it. The header, the
guide, the contract of `ColumnarWriter` and the doc of
`ArrowBatchingDataWriter` say that the batches that Spark copies end with those
of the scan, so a library may get smaller batches than `maxRecordsPerBatch`
rows. The PR description is updated too.
##########
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:
Done in b3afefd: the guide now uses the suggested sentence, and the PR
description is updated.
##########
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:
Done in b3afefd: the `withSQLConf` also sets `spark.sql.adaptive.enabled` to
true and `spark.sql.adaptive.forceApply` to false.
##########
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:
Done in b3afefd: renamed the block to `writeData`.
--
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]
