HyukjinKwon opened a new pull request, #59933:
URL: https://github.com/apache/spark/pull/59933
### What changes were proposed in this pull request?
A follow-up to SPARK-60101 (#59424), discussed there with @zhouyuan and
@zhztheplayer. When the query of a batch or micro-batch write to a native data
source produces columnar batches, the data writers now get the batches. Before,
a `ColumnarToRowExec` converted them to rows first. A scan of a native data
source is one such query. This PR is based on #59931 (version 1 of the
interface), whose commit it includes until that PR is merged.
- **Planning and execution.** An internal mixin, `SupportsColumnarInput` in
`org.apache.spark.sql.internal.connector`, marks:
- a `Write`, for appends and overwrites;
- the `StreamingWrite` of a micro-batch.
Its data writers are `ColumnarInputDataWriter`s, which also take columnar
batches.
- `V2TableWriteExec.supportsColumnarInput` tells whether a write node
writes the batches of its query. Row-level operations do not: their tasks
project the rows.
- `ApplyColumnarRulesAndInsertTransitions` then inserts no transition to
rows above a columnar query, as it already does for planned V1 writes.
- `writeWithV2` runs `DataWritingSparkTask.runColumnar`, which shares the
commit protocol of `WritingSparkTask` (moved into `writeAndCommit`). It
releases each batch with `closeIfFreeable` once written, as `ColumnarToRowExec`
does.
- **Copying the batches.** The data writer of the columnar adapter
(`ArrowBatchingDataWriter`) copies the Arrow vectors of a batch buffer by
buffer with `VectorAppender` when they are physically congruent with the Arrow
schema of the write (`ArrowUtils.isCompatibleWithDeclaredField`), as the
vectors of a native scan are. It converts other columns value by value with
`ArrowWriter`, without going through `UnsafeRow`s. Either way, the batches end
at `spark.sql.execution.arrow.maxRecordsPerBatch` rows, or once they hold
`spark.sql.execution.arrow.maxBytesPerBatch` bytes, as the header promises. A
batch is copied rather than passed as it is for two reasons:
- The native data writer may move the arrays that it gets, to keep them
after the call.
- The producer of a `ColumnarBatch` may overwrite its vectors once the
batch is consumed, as `ArrowWriter.reset()` does.
Arrow vectors are sliced before they are copied, because `VectorAppender`
assumes that the offsets of a vector start at 0, which they need not in an
array imported through the C data interface.
Under adaptive query execution, the query of a write is an
`AdaptiveSparkPlanExec`, which is not columnar, so the batches only reach the
write when the query has no shuffle, as for a plain copy from one native data
source to another. A columnar rule can mark the adaptive plan columnar, as
`CachedBatchSerializer.convertToColumnarPlanIfPossible` does.
### Why are the changes needed?
The batches of a native scan were converted to rows by `ColumnarToRowExec`,
and then back to Arrow by the data writer. With this change, a copy from one
native data source to another stays in Arrow, at the cost of a copy of the
buffers. A native engine can also write its batches to a native sink without
converting them to rows.
### Does this PR introduce _any_ user-facing change?
No. The written data is the same. The plan of such a write no longer has a
`ColumnarToRow` node.
### How was this patch tested?
New tests in `ColumnarDataSourceSuite`:
- "write the columnar batches of a query" checks the plan and the batch
sizes for on-heap and Arrow batches. It also checks a data writer that keeps
the data of its batches, through the C data interface, while the scan
overwrites its vectors with each batch.
- "copy the Arrow columns of every type, or convert the columns of a batch"
covers every type that Spark converts to Arrow, nulls, slices at offsets that
are not aligned to bytes, the limits in rows and in bytes, Arrow types other
than those of the schema, and on-heap vectors. Each case checks the data that
the data writer kept after the producer overwrote the batch.
The plans of the batch and streaming writes of `NativeDataSourceSuite`, and
of "streaming read and write" in `ColumnarDataSourceSuite`, are now also
checked.
Replacing the copy with the slices of the batch makes the tests of kept data
fail. `DataSourceV2Suite`, `SparkSessionExtensionSuite`, `ColumnarRulesSuite`,
`DeltaBasedUpdateTableSuite`, `DeltaBasedMergeIntoTableSuite`,
`StreamingDataSourceV2Suite` and `V1WriteFallbackSuite` pass.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
This pull request and its description were written by Isaac.
--
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]
