andygrove opened a new pull request, #6858: URL: https://github.com/apache/datafusion-comet/pull/6858
## Which issue does this PR close? Part of #6855. ## Rationale for this change #6855 reports that JVM columnar shuffle with 201 partitions is up to 20x slower than Spark on the fuzz-generated nested schema in `CometShuffleBenchmark` (100 columns, 1,000 rows), and suggests a per-partition cost in the native writer. A profile of the max-depth-6 case shows that the writer is not where the time goes: the native library takes about 10% of the samples. The cost is on the reduce side, and every one of the 201 tasks pays it: - The largest share went to `CometColumnarToRowExec.doExecute`, which calls `UnsafeProjection.create` in every partition. That call generates the projection's Java source each time; Spark caches only the compiled class. The plan runs outside whole-stage codegen because the schema has more nested fields than `spark.sql.codegen.maxFields`. Spark's own scans return rows for such schemas, so Spark rarely runs `ColumnarToRowExec` on one, but Comet's operators always produce batches. Generating the source once per executor saves about 19 ms per partition, half the time of the case. - Most of the rest is importing the wide schema's Arrow arrays for each batch, and wrapping and closing the vectors. This PR does not change that. The same per-partition cost applies to any query whose `CometColumnarToRowExec` runs outside whole-stage codegen, which a wide or deeply nested output always does. In this benchmark, the "Comet (native Shuffle)" arm runs Spark's shuffle. `repartition(n)` is round-robin, which native shuffle declines by default (`spark.comet.shuffle.native.partitioning.roundrobin.enabled`), and the fuzz schema is read by Spark's scan. So that arm's parity with Spark says nothing about native shuffle. ## What changes are included in this PR? - `CometUnsafeProjection` creates the projection that `CometColumnarToRowExec.doExecute` copies rows through. It generates and compiles the class once per executor for each column layout, meaning the column types and nullability plus `spark.sql.codegen.methodSplitThreshold`, the one setting the generated source depends on. It keeps up to 100 classes and returns a new instance on every call. Spark's `GenerateUnsafeProjection.create` returns only an instance, so the class is built from the same template, which is identical in Spark 3.4 through 4.1. It honors `spark.sql.codegen.factoryMode` as `UnsafeProjection.create` does. - `CometColumnarToRowExec.doExecute` uses it. The whole-stage codegen path and the broadcast path, which builds one projection on the driver, are unchanged. The in-memory cache reader (`CachedBatchRowIterator`) also generates its code in every partition and could cache its class the same way. That is left for a follow-up. ## How are these changes tested? `CometUnsafeProjectionSuite` checks that: - rows match `UnsafeProjection.create` for structs, arrays and maps with nulls, including with a method split threshold small enough to split the field writes - a layout's class is generated once, and each call returns a separate projection - a different split threshold generates its own class - the least recently used class is evicted after 100 - `NO_CODEGEN` returns the interpreted projection - a query whose `CometColumnarToRowExec` runs outside whole-stage codegen matches Spark, and a second run generates no class `CometShuffleSuite` and `CometExecSuite` pass, and the change compiles against Spark 3.4, 3.5 (strict warnings), 4.0 and 4.1. `CometShuffleBenchmark`'s fuzz-generated nested schema cases, run on a Ryzen 9 7950X3D with JDK 17, Spark 4.1 and `local[1]`. The times are the best of five for the JVM shuffle arm, with Spark's time from the same run: | Max depth | Partitions | Before | After | Spark (before / after runs) | |---|---|---|---|---| | 2 | 5 | 114 ms | 114 ms | 117 / 127 ms | | 2 | 201 | 827 ms | 404 ms | 93 / 117 ms | | 6 | 5 | 681 ms | 530 ms | 525 / 454 ms | | 6 | 201 | 7,798 ms | 3,992 ms | 369 / 376 ms | The depth-6, 5-partition rows have a standard deviation near 200 ms, so only the 201-partition rows show a real change. The remaining gap at 201 partitions is the per-batch Arrow import described above. -- 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]
