sunchao opened a new pull request, #5628: URL: https://github.com/apache/datafusion-comet/pull/5628
## Which issue does this PR close? Closes #5617. ## Rationale for this change Native shuffle currently remembers spilled Arrow allocations only while processing one outer input batch. A producer such as partial HashAggregate can instead emit separate zero-copy slices of one allocation. With sixteen chunks of a 128 KiB Int64 buffer and sixteen spills, the issue reports 2,228,224 bytes for separately delivered chunks versus 262,144 bytes when the writer slices one incoming batch itself. The new regression reproduces the inflated 2,228,224-byte result on the original implementation. Extending the address set across input calls would need to keep allocations alive to prevent address reuse, or require an allocation-lifetime handle that Arrow does not expose. Retaining payloads solely for metrics would interfere with the memory the spill is meant to release. ## What changes are included in this PR? Measure the buffer footprint of the batches that the existing partition iterator materializes for spilling, before compression, and add the partition-index allocation capacities released by the spill. The metric observes existing batches; it does not copy payloads or retain them. This removes the per-input spilled-address set, repeated-buffer subtraction, and special adjustment for a rejected reservation. Physical input reservations, spill triggers, data_size, and disk-byte accounting remain unchanged. This deliberately defines memory_spilled_bytes as the cumulative footprint of materialized spill batches and their partition indices, not globally unique input allocations or literal process memory freed. Dictionary and string-view output can still share backing storage across emitted batches. Both internal and external delivery of the same chunks now use the same output-footprint measurement; repeated data continues contributing cumulatively. On a write error, indices are released and counted, and data-buffer bytes are counted only for batches already materialized for the writer. Existing cleanup coverage now checks failure before and after consuming a batch and verifies that neither path retains input buffers. ## How are these changes tested? The regression compares one incoming batch with sixteen externally sliced inputs under both the maximum-buffer and rejected-reservation spill triggers. It verifies equal memory metrics, equal spill counts, byte-identical shuffle output, independently allocated chunks, and cumulative repeated input. Additional coverage exercises nullable string views, dictionaries, and lists. Local validation on Linux with JDK 17: - `cargo test --locked -p datafusion-comet-shuffle --lib -- --test-threads=1`: **94 passed**, none failed or ignored. - `cargo clippy --locked -p datafusion-comet-shuffle --all-targets -- -D warnings`: passed. - `cargo fmt --all -- --check` and `git diff --check`: passed. - `make core`: passed. - Root-reactor Spark **4.0.4** validation with `./mvnw -B -Pspark-4.0 test spotless:check -Dtest=none '-Dsuites=org.apache.spark.sql.comet.CometTaskMetricsSuite memory and disk spill metrics'`: **2 tests passed in 1 suite**, full reactor **BUILD SUCCESS**, including Spotless. The two Spark tests verify successful shuffle reporting and preservation of memory/disk spill metrics on failed attempts. Native builds used eight jobs with dev/test debug symbols disabled; Spark ran with `SPARK_LOCAL_IP=127.0.0.1` and `SPARK_LOCAL_HOSTNAME=localhost`. Other Spark versions and benchmarks were not rerun. -- 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]
