On Spark 4.0+, an adaptive cached plan that ends in a Spark row operator is converted into Spark column vectors and then copied into Arrow #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 OnHeapColumnVectors. 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 (55ba50d) 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 OnHeapColumnVectors:

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:

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.