andygrove opened a new issue, #6952:
URL: https://github.com/apache/datafusion-comet/issues/6952
### What is the problem the feature request solves?
Part of #6834.
From Spark 4.0, when the cache serializer accepts columnar input,
`InMemoryRelation` marks a cached `AdaptiveSparkPlanExec` as columnar
(`adaptive.copy(supportsColumnar = true)`, in
`CachedBatchSerializer.convertToColumnarPlanIfPossible` on 4.1 and 4.2 and
`InMemoryRelation.convertToColumnarIfPossible` on 4.0). AQE then plans the
final stage for columnar output, and when that stage ends in a Spark row
operator, Spark puts a `RowToColumnarExec` on top of it. Comet's serializer
accepts columnar input for every schema it can store (`supportsColumnarInput =
supportsSchema(schema)`, `ArrowCachedBatchSerializer.scala:503`), so this
happens to any cached query with an exchange whose last operator Comet does not
run, including queries Comet does not accelerate at all.
`RowToColumnarExec` fills Spark's `OnHeapColumnVector`s. `encodeBatches`
then copies those into Arrow through
`CometArrowConverters.columnarBatchToArrowBatch`, because they are not
Arrow-backed (`ArrowCachedBatchSerializer.scala:418` and `:424`). Without AQE,
and on Spark 3.x, the same plan takes the row path, which writes the rows into
Arrow once (`convertInternalRowToCachedBatch`, `:655`, through
`CometArrowConverters.rowToArrowBatchIter`). Comet leaves that
`RowToColumnarExec` in place. `EliminateRedundantTransitions` removes one only
below a Comet columnar shuffle.
On main (55ba50d451) with Spark 4.1, the cached relation for
`spark.read.parquet(src).repartition(8, $"a").selectExpr("id + 1 as id1", "a",
"b * 2 as b2", "s").cache()` has this final plan, and the batches that reach
the serializer hold `OnHeapColumnVector`s:
```
ResultQueryStage 1
+- RowToColumnar
+- *(2) Project [(id#4L + 1) AS id1#9L, a#5, (b#6 * 2.0) AS b2#10, s#7]
+- ShuffleQueryStage 0
+- Exchange hashpartitioning(a#5, 8), REPARTITION_BY_NUM,
[plan_id=64]
```
When the scan and the shuffle run natively and only the root is Spark's
(here with `spark.comet.exec.project.enabled=false`), the same `RowToColumnar`
sits above `CometColumnarToRow`, so the rows go from Arrow to `UnsafeRow`, to
Spark's vectors and back to Arrow.
Measured with that query run entirely by Spark
(`spark.comet.exec.enabled=false` and `spark.comet.scan.enabled=false`, so no
native code runs, with the cache still in Comet's format), over 5M rows of four
columns (bigint, int, double, string) in 8 partitions, with `local[1]` on an
Apple M3 Max, median of twelve interleaved runs:
| | AQE on (two copies) | AQE off (one copy) |
| --- | ---: | ---: |
| Cache build | 1994 ms | 1796 ms |
| Same query uncached, `noop` write | 1301 ms | 1284 ms |
The extra copy costs about 11% of the build here, about 40 ns per row, and
AQE itself accounts for almost none of the gap. An earlier run of six gave 2126
against 1914 ms.
### Describe the potential solution
Replace a `RowToColumnarExec` at the root of an adaptive plan's final stage
with `CometSparkToColumnarExec`, which writes Arrow straight from the rows, so
that `encodeBatches` gets batches it can write as they are. Spark's
`RowToColumnarExec` would stay for schemas `CometSparkToColumnarExec` does not
support.
### Additional context
Spark 3.4 and 3.5 have no adaptive case in `convertToColumnarIfPossible`, so
they always take the row path. That is the epic's item about 3.x building
through rows, the other side of the same switch.
--
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]
