sunchao opened a new pull request, #7258: URL: https://github.com/apache/datafusion-comet/pull/7258
## Which issue does this PR close? Closes #7257. ## Rationale for this change A broadcast exchange fed by native shuffle can receive thousands of tiny batches. Each batch currently crosses Arrow FFI and becomes a separate serialized stream before the driver collects and combines it. Combining small batches on the producer removes repeated export, serialization and driver stream-opening work. ## What changes are included in this PR? Add the experimental, default-off `spark.comet.exec.broadcastExchange.producerCoalesce.enabled` setting. It groups decoded batches for direct local Comet native shuffle inputs to broadcast exchange, while preserving AQE partition ranges and shared shuffle dependencies. Eligible schemas contain only INT, BIGINT, BOOLEAN, STRING and BINARY columns. One task-owned decoder groups across fetched blocks before Arrow export. Groups stop at `min(spark.comet.batchSize, 65536)` rows, 1,024 inputs or an 8 MiB decoded-buffer budget. Singletons, oversized inputs and batches at least half the row target pass through without concatenation. The 8 MiB limit bounds grouped inputs, **not total reader memory or RSS**: concatenation also owns its output and may retain a decoded lookahead. Execution stays native; this introduces no Spark fallback. Other input paths and driver broadcast assembly retain their existing behavior. ## How are these changes tested? - Passed 14 focused Rust tests: 12 coalescer tests and two JNI ownership/frame-validation tests. - Passed strict Spark 3.5/Delta compilation, Rust formatting, JVM formatting/style, CI suite registration, and Clippy for both changed native crates with warnings denied. - Passed 168 Spark 3.5 tests across `CometNativeShuffleSuite`, `CometCelebornShuffleReaderSuite`, `CometNativeShuffleInputRDDSuite`, `UtilsSuite` and `CometBlockStoreShuffleReaderSuite`; six existing Spark 4-only cases were canceled by their version guards. The tests cover actual batching/stream reduction with AQE ON/OFF, unchanged native routing for unsupported schemas, shared shuffle dependencies, and decoder cleanup/error paths. - Local native validation uses an isolated source copy with an available alternate Cargo.lock because dependencies including `aws-config 1.13` are unavailable locally. Arrow 59.3.1 and DataFusion 55.2.0 are unchanged; the submitted Cargo.lock is unchanged. These results do not validate resolution/build with the submitted lockfile. ### Local producer benchmark Measured decoding → Arrow FFI export → broadcast serialization with `CometBroadcastProducerBenchmark`, using one optimized native artifact, Spark 3.5.9 / JDK 17, a 4 GiB JVM heap, and four pinned CPUs on an AMD EPYC Milan host. Two fresh JVM runs used OFF/ON/OFF and ON/OFF/ON. Each arm ran at least five iterations and two measured seconds after warmup. All fixtures contain 8,192 rows; output values/nulls and stream/byte counts were verified outside timing. The table shows the range of **mean producer times** across the three measured arms per setting. Sub-millisecond means are calculated from the reported elapsed time and iteration count. | Input rows per frame | OFF mean (ms) | ON mean (ms) | Output streams OFF → ON | |---|---:|---:|---:| | 1, INT + BIGINT | 1,162–1,416 | 15.2–20.7 | 8,192 → 8 | | 8, INT + BIGINT | 144–151 | 2.07–2.96 | 1,024 → 1 | | 8,192, full-batch control | 0.327–0.330 | 0.324–0.498 | 1 → 1 | | 64, plus a 6 KiB STRING | 45.3–46.8 | 29.3–59.7 | 128 → 7 | The tiny-batch gain persists in both orders. Full-batch means include one slower ON arm; the other two are close to OFF. The wide-string control also includes a slower ON arm and substantial variance, so the optimization remains opt-in. Its repeated `x` payload is a byte-budget control, not representative JSON throughput. For one-row inputs, serialized bytes fell from 2,295,384 to 64,784. This measures the producer only; transport, driver assembly and whole-query performance are excluded. The alternate native dependency set described above also applies to these measurements (release optimization, LTO disabled, 16 codegen units). Reproduce with: ```sh make benchmark-org.apache.spark.sql.benchmark.CometBroadcastProducerBenchmark BENCH_HEAP=4g make benchmark-org.apache.spark.sql.benchmark.CometBroadcastProducerBenchmark BENCH_HEAP=4g -- --on-first ``` ### Historical whole-query measurements The original internal native-Celeborn implementation recorded these results; **they do not measure this OSS local-shuffle adaptation**. Runs used 64 four-core executors, one warmup and three measured iterations per application, with matching artifacts, image, SQL and settings except batching. The optional memory guard was off. Times are medians from SQL setup through complete result collection; complete results and iteration fingerprints matched. | Workload | Batching OFF | Batching ON | Time reduction | |---|---:|---:|---:| | TPC-H Q20, 1 TB | 46.17 s | 26.71 s | 42.1% | | TPC-H Q19, 10 TB | 85.12 s | 54.66 s | 35.8% | | Wide rows: key + 128 INT columns | 20.82 s | 20.03 s | 3.8% | | Large JSON, OFF-first | 7.67 s | 7.91 s | −3.1% | | Same JSON, ON-first | 7.67 s | 7.34 s | 4.3% | Q19 incoming streams fell from 2,047,700 to 2,083 and driver build time from 32.64 s to 0.89 s. Its slower third ON iteration (72.99 s) was retained. JSON controls showed no repeatable direction; wide-row transport bytes increased 4.0%. These historical results motivate the optimization but do not establish its benefit for every workload or for the submitted dependency set. -- 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]
