sam-1112 opened a new pull request, #5559: URL: https://github.com/apache/datafusion-comet/pull/5559
## Which issue does this PR close? Related to #5409 and #5354. ## Rationale for this change DataFusion 54.1.0's `UnnestExec` can emit batches containing more rows than `spark.comet.batchSize`. This violates the runtime batch-size boundary expected by downstream native operators. In particular, a downstream projection that broadcasts a scalar into an Arrow `Utf8` array can overflow its 32-bit offsets. This PR wraps native `UnnestExec` with `BatchSplitExec`, which splits oversized output batches according to the runtime `TaskContext` batch size before passing them downstream. This is a temporary workaround until Comet upgrades to a DataFusion version containing [datafusion#24384](https://github.com/apache/datafusion/pull/24384). The change was split out of #5409 so that the broader explode runtime behavior can be reviewed independently. ## What changes are included in this PR? - Add `BatchSplitExec` after native `UnnestExec`. - Preserve row order while splitting oversized batches. - Forward the existing `UnnestExec` metrics through the wrapper. - Expose `batches_split` on `CometExplodeExec`. DataFusion increments this counter once for each output slice produced from an oversized input batch. - Keep `output_batches` as the original `UnnestExec` output-batch count from before splitting, rather than changing the meaning of the existing metric. ## Limitations - This bounds the batch size observed by downstream operators. It does not reduce peak memory while `UnnestExec` constructs the expanded batch. - `elapsed_compute` remains the `UnnestExec` compute time and does not include the wrapper's slicing cost. - `batches_split` exposes split activity, but not the time spent slicing. - The wrapper should be removed after upgrading to a DataFusion version containing datafusion#24384. ## Performance Measured using a release native build (`-Ctarget-cpu=native`) with 5 warm-up runs and 10 measured runs. Values below are medians in ns per input row. | Case | Without wrapper | With wrapper | Difference(with − without) | | -------- | --------------: | -----------: | ---------: | | No split | 389 ns/row | 395 ns/row | +6 ns/row | | Split | 487 ns/row | 444 ns/row | −43 ns/row | Environment: - Queries against temp view `src` = `spark.range(131072)`, `spark.sql.leafNodeDefaultParallelism=1` - No split: `SELECT explode(array(id)) AS e FROM src` (array size 1) - Split: `SELECT explode(array(id, id)) AS e FROM src` (array size 2) - Rows: 131072 input rows - `spark.comet.batchSize`: 8192 - Spark 4.1.3, JDK 21.0.12, DataFusion 54.1.0 - Hardware: Apple M4 - Timed against `076c092a4`. Later commit `db10f73db` does not change the explode execution path. Both arms used the same source revision and release build settings. For the without-wrapper arm, the planner was changed locally to return the original `UnnestExec` directly; no other code was changed. The observed no-split pass-through overhead was approximately 6 ns per input row (about 1.5%). The split arm was 43 ns per input row faster in this end-to-end measurement, possibly because downstream `collect()` received smaller batches. This result does not establish that splitting itself improves performance. No regression was observed in this workload. ## How are these changes tested? The Rust unit test verifies that: - a 10-row batch with a runtime batch size of 4 becomes `[4, 4, 2]`; - row order is preserved; - the wrapped plan's `output_rows` metric is forwarded; and - `batches_split == 3`. The Spark integration test verifies that: - a single-partition 16-row input expands to 32 rows through native `explode`; - the plan contains `CometExplodeExec`; - results match Spark; - `input_rows == 16`; - `output_rows == 32`; and - `batches_split > 0` is exposed through the Spark SQL metric map. ```bash (cd native && \ cargo test -p datafusion-comet \ splits_batches_at_the_runtime_batch_size_and_preserves_order) ./mvnw test -Dtest=none \ -Dsuites="org.apache.comet.exec.CometGenerateExecSuite native explode splits" ``` -- 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]
