andygrove opened a new pull request, #6926: URL: https://github.com/apache/datafusion-comet/pull/6926
## Which issue does this PR close? Closes #6913. ## Rationale for this change When a native Iceberg write sits directly on a Comet shuffle, `CometIcebergWriteExec` read the shuffle the way a Spark operator would. The JVM fetched and decoded every block, exported the batches over Arrow FFI, and the write's native `Scan` imported them again and re-validated their strings as UTF-8. Shuffle direct read (`spark.comet.shuffle.directRead.enabled`, on by default) exists to skip that round trip, but only native blocks built by `CometNativeExec` used it, because the write builds its own native plan. #6913 has the profile. The shape is common. Under AQE it is the plan for a repartitioned insert, and for a hash-distributed write that needs no sort, such as one to a partitioned table without a sort order, including copy-on-write DML. ## What changes are included in this PR? - `CometIcebergNativeWrite.convert` takes the child's `ShuffleScan` as the leaf of the write's native plan when the child is a sink that reads its shuffle directly (`CometExchangeSink`), and builds the FFI `Scan` as before otherwise. The write follows the decision every other native operator follows, so it reads directly under AQE only, and honours `spark.comet.shuffle.directRead.enabled`. - `CometIcebergWriteExec.doExecuteColumnar` returns a new `CometIcebergWriteRDD`, which resolves the write's one input slot with `CometExecRDD.resolveInputObjects`, as every native block does: a `ShuffleScan` leaf gets a `CometShuffleBlockIterator` over the raw blocks, and a `Scan` leaf gets the child's batches as an Arrow stream. The per-task `partition_id` and `task_attempt_id` stamping moves into its `compute`. - No native changes. The planner already accepts a `ShuffleScan` under `IcebergWrite`, and `ShuffleScanExec` reconciles each batch with the declared schema the way `ScanExec` does. - The contributor guide's `iceberg-writes.md` and `native_shuffle.md` describe the new input path. ## How are these changes tested? A new test in `CometIcebergWriteActionSuite` writes over a shuffle stage (`REPARTITION(8, id)`), a coalesced read of one (`REPARTITION(id)`), a JVM columnar shuffle (`REBALANCE`), and a shuffle with AQE off, each with direct read on and off. It checks which leaf the write's native plan reads its input with (`ShuffleScan` under AQE with direct read on, `Scan` otherwise), that the rows, which include string, struct, array and map columns with nulls, match the source, and that both input paths write data files with the same record counts. It fails without the change, and passes on the Spark 3.4, 3.5, 4.0 and 4.1 profiles. `CometIcebergWriteActionSuite`, `CometIcebergWriteDetectionSuite`, `CometIcebergRewriteActionSuite`, `CometIcebergSystemFunctionSuite`, `RevertNativeForTransitionHeavyStagesSuite` and `CometScanSchemeFallbackSuite` pass. With this change their AQE runs of partitioned writes and copy-on-write DML read their shuffle directly as well. The Iceberg Spark test diffs enable the native writer, so `run-iceberg-tests` is applied. ### Benchmark `CometIcebergEndToEndBenchmark` from #6910 (2M rows, `local[5]`, `INSERT ... SELECT /*+ REPARTITION(10, id) */` into an unpartitioned table, release build, Apple M3 Ultra). I added a local arm that runs the native write with `spark.comet.shuffle.directRead.enabled=false`, which is how the write read its input before this change, and ran the two arms in both orders, because whichever arm runs first in a group pays some JIT warm-up. The table gives the average statement time over four runs and the process CPU time per statement over two: | Type group | Wall, FFI | Wall, direct | Change | CPU, FFI | CPU, direct | Change | | ---------- | --------: | -----------: | -----: | -------: | ----------: | -----: | | primitives | 554 ms | 546 ms | -1.4% | 2177 ms | 2051 ms | -5.8% | | structs | 416 ms | 407 ms | -2.2% | 1464 ms | 1380 ms | -5.8% | | arrays | 604 ms | 598 ms | -1.0% | 2336 ms | 2300 ms | -1.5% | | maps | 571 ms | 562 ms | -1.6% | 2323 ms | 2236 ms | -3.8% | | nested | 1193 ms | 1173 ms | -1.7% | 5260 ms | 5128 ms | -2.5% | The gain is small, as the profile predicted, and close to the noise of a single run. Direct read was faster in 17 of the 20 wall-time comparisons, and the other three were within 2 ms. The benchmark is not part of this PR. -- 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]
