andygrove opened a new pull request, #5632:
URL: https://github.com/apache/datafusion-comet/pull/5632
## Which issue does this PR close?
Part of #5625. That issue proposes handing Spark's writers a zero-copy row
view instead of a
materialized `UnsafeRow`, and restricts the idea to unpartitioned,
unbucketed writes. #5626
implements it that way. This PR is an alternative that lifts the restriction.
**Draft / RFC.** Off by default. The implementation is complete and tested,
but there are no
measurements yet, and the measurements are what should decide whether this
lands at all. See
"What is still owed" at the bottom.
## Rationale for this change
The restriction in #5625 is not arbitrary. A partitioned or bucketed write
never reaches
`SingleDirectoryDataWriter`; it goes through
`BaseDynamicPartitionDataWriter`, which projects
every row a second time before the `OutputWriter` sees it:
```scala
protected val getOutputRow =
UnsafeProjection.create(description.dataColumns, description.allColumns)
// :266
protected def writeRecord(record: InternalRow): Unit = {
val outputRow = getOutputRow(record) // strips partition / bucket columns
currentWriter.write(outputRow)
...
}
```
So simply planting a row view under the write does not help a partitioned
write. Today the cost is
a slow field-by-field build of nested `UnsafeRow` / `UnsafeArrayData` in the
columnar-to-row
transition, followed by a cheap second projection that hits
`GenerateUnsafeProjection`'s
`instanceof UnsafeRow` / `UnsafeArrayData` / `UnsafeMapData` bulk-copy
branches
(`GenerateUnsafeProjection.scala:72,193,227`). Feeding the writer a
`ColumnarBatchRow` would delete
the cheap pass and leave the expensive one, because those branches stop
matching.
The projection exists for one reason: to drop the partition and bucket
columns. Spark needs a full
row copy to do that because it is row-at-a-time. On a `ColumnarBatch` the
same pruning is
`new ColumnarBatch(subsetOfVectors, numRows)`, which copies nothing. Getting
that requires
replacing the writer, not the transition.
Spark 4.0 provides the seam. `V1WritesUtils.getWriteFilesOpt` matches the
`WriteFilesExecBase`
trait, so a Comet node extending it is driven through
`FileFormatWriter.executeWrite` -> `SparkPlan.executeWrite` ->
`doExecuteWrite`, and everything
above the per-task write stays with Spark: SaveMode, the commit protocol,
`_SUCCESS`, dynamic
partition overwrite, stats tracker aggregation, catalog updates. This is the
same seam #5293 opens
for native writes, and the two are consistent by design (see below).
## Relationship to #5293 and #5626
- **#5293** (`refactor: hook native Parquet writes into Spark's
WriteFilesExec seam`) puts a
`CometWriteFilesExec` at this seam that writes Parquet with the **native
Rust** writer, currently
unpartitioned only. This PR reuses that PR's `ShimCometWriteFilesExec`
traits **verbatim**, so
whichever lands second is a trivial rebase. The node here is named
`CometRowViewWriteFilesExec` to avoid a file-level collision.
- **#5626** implements #5625 the other way, as a
`CometColumnarToRowViewExec` transition planted
under the write. If this PR is the direction people prefer, #5626 should
be closed in favour of
it, because this one subsumes it: it covers the unpartitioned case too,
and with a stronger
safety argument.
On that argument: #5626's main caveat is that a reused mutable row becomes
an operator's output
and could reach a consumer that retains it. Here the reused rows never leave
`CometRowViewWriteFilesExec.doExecuteWrite`; they are created and consumed
inside one method whose
only consumer is Spark's `OutputWriter`. Extending `WriteFilesExecBase` is
also what stops AQE from
re-inserting a second `WriteFilesExec` above the node, which is the hazard
#5293 documents.
## What changes are included in this PR?
`CometRowViewWriteFilesExec` replaces `WriteFilesExec` and drives Spark's
own writers from
`child.executeColumnar()`. Its `executeTask` is a direct port of
`FileFormatWriter.executeTask`,
differing in two places:
- the partitioned and bucketed case builds a
`CometRowViewDynamicPartitionWriter`, a subclass of
`DynamicPartitionDataSingleWriter` that overrides only `writeRecord` to
substitute a pruned view
for `getOutputRow(record)`. Partition-change detection, writer renewal and
`maxRecordsPerFile`
are inherited unchanged and still see the full row.
- rows come from two iterators advanced in lockstep over the same batch: one
over all columns,
which Spark's writer uses to compute partition values and bucket ids, and
one over the pruned
batch, which is what reaches the `OutputWriter`. Both are views over the
same Arrow vectors, so
the pair costs one object per batch rather than a copy per row.
The unpartitioned case needs no subclass: `SingleDirectoryDataWriter` hands
the row straight to the
`OutputWriter`, so the batch row goes in as-is.
`EliminateRedundantTransitions` plants the node, gated on:
- **Spark 4.0+.** On 3.4 / 3.5 `getWriteFilesOpt` matches the concrete
`WriteFilesExec` case class,
so a replacement node is invisible and the write would silently take a
path that calls
`doExecute` on it. Same reasoning and same shim as #5293.
- **`spark.sql.maxConcurrentOutputFileWriters` at its default of 0.** Above
0,
`V1WritesUtils.getSortOrder` plants no sort and `FileFormatWriter` picks
`DynamicPartitionDataConcurrentWriter`, which spills through
`UnsafeKVExternalSorter`. This gate
is load-bearing in both directions: it keeps that writer out, and it is
what guarantees
`DynamicPartitionDataSingleWriter` gets the sorted input it requires.
- **one of Spark's own `FileFormat`s**, whose `OutputWriter`s encode each
row on the spot. A
third-party format is free to buffer the `InternalRow` it is handed.
- **a complex type among the data columns.** Partition and bucket columns do
not count; they never
reach the `OutputWriter`. On flat schemas the projection this removes is a
generated fixed-width
copy that #5625 measured inside the noise floor.
No check for an intervening `SortExec` is needed. The rule only rewrites
`w.child`, so a write
whose required ordering was satisfied by a Spark `SortExec` rather than a
Comet one simply does not
match. This is worth stating because #5625 and #5626 both cite
`UnsafeExternalSorter` as a reason
to exclude partitioned writes; that exclusion was already enforced
structurally.
Behind `spark.comet.exec.write.rowView.enabled`, default false.
## How are these changes tested?
New `CometWriteRowViewSuite`, 15 tests, registered in both PR build
workflows and gated on
`isSpark40Plus`. The bar for each is that enabling the config changes
nothing observable but the
plan, with the baseline written by the same Comet plan with the config off:
- same data and row count for: unpartitioned; one dynamic partition column;
two partition columns;
deeply nested types (four levels, nulls at every level) with partitions;
null and empty-string
partition values, which exercise `__HIVE_DEFAULT_PARTITION__` and the
`Empty2Null` projection;
`maxRecordsPerFile` set and unset across partition boundaries; ORC as well
as Parquet
- identical partition directory layout
- the node is present exactly once for unpartitioned, dynamically
partitioned and bucketed writes,
absent by default, absent for a flat schema, absent when the only complex
column is the partition
column, and absent when `maxConcurrentOutputFileWriters` is raised
Ran alongside `CometParquetWriterSuite`: 48 tests pass on Spark 4.1.
`CometExecSuite` passes as a
regression sweep. Compiles against Spark 3.4, 3.5, 4.0 and 4.1.
## What is still owed
1. **Measurements.** There are none yet, and they are the whole question.
The reasoning above says
a partitioned write should now save two projections per row rather than
the one #5626 saves, but
reasoning is not a number. A partitioned arm in
`CometParquetWriteBenchmark` (complex schema,
low-cardinality partition column, pre-sorted input) is the next step, and
this should not land
without it. It is also possible that the complex-type gate should be
relaxed for partitioned
writes, since `getOutputRow` is removed there even on flat schemas.
2. **The coupling.** `writeRecord`, `currentWriter`, `statsTrackers` and
`recordsInFile` are
`protected` members of Spark's writer, and this depends on their current
structure across
4.0 / 4.1 / 4.2. A future refactor there would be a compile error rather
than silent corruption,
which is the good case, but the dependency is real and worth a second
opinion.
3. `WriteTaskStatsTracker.newRow` now receives a reused row rather than a
materialized one. That is
correct for `BasicWriteTaskStatsTracker`, which ignores it, and it is
strictly better than
#5293's native path, which passes `InternalRow.empty`. A third-party
tracker that retains rows
would still be wrong. There is a format gate but no tracker gate;
`statsTrackers` only exists at
execution time.
--
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]
