fix: make native shuffle spill metrics independent of input batching - #5628
Open
sunchao wants to merge 2 commits into
Open
fix: make native shuffle spill metrics independent of input batching#5628sunchao wants to merge 2 commits into
sunchao wants to merge 2 commits into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Closes #5617.
Rationale for this change
Native shuffle deduplicates spilled Arrow allocations only within one outer
insert_batchcall. A producer such as partial HashAggregate can deliver the same zero-copy chunks as separate input batches. For the issue's 128 KiB allocation and sixteen spills, this reports 2,228,224 bytes for separately delivered chunks but only 262,144 bytes when one input is sliced internally. The spill boundaries and pinned allocations are identical, so caller batching should not change the metric.Use cumulative per-spill input accounting: each spill counts the capacity of its buffered input allocations, deduplicated within that spill, plus its partition-index allocations. Buffering the same allocation for a later spill contributes again. Thus both delivery shapes report 2,228,224 bytes. This is cumulative memory released from shuffle buffering, not peak memory, globally unique task allocations, or necessarily a drop in process RSS when another operator still owns the buffer.
This also preserves the metric's input-memory meaning. Counting materialized output instead can turn a unique 1 MiB input allocation containing 1,024 live Int64 rows into only 8 KiB of reported data, or count shared dictionary/view payloads repeatedly within one spill. Four new regression tests reproduced these mismatches before the correction and pass with per-spill input accounting.
What changes are included in this PR?
Remove the input-batch-scoped spilled-address set and repeated-buffer discount. Continue using the existing reservation, which deduplicates input allocations while they are buffered, and include the just-buffered growth if its reservation was rejected. After the spill producer drops its inputs and indices, publish the released reservation plus any rejected growth. Failed writes also count all input memory released during cleanup, regardless of how much output the writer consumed.
No payloads or allocation addresses are retained across spills for metrics. Physical memory reservations, spill triggers, shuffle output,
data_size, and disk-byte accounting remain unchanged. The user-facing metric documentation now explains the cumulative per-spill contract.How are these changes tested?
The batching regression exercises maximum-buffer and rejected-reservation spills, verifies equal metrics and spill counts for internal/external slicing, and compares serialized shuffle output. Independently allocated small chunks correctly report less pinned memory while producing identical output. Repeated spills remain cumulative.
Additional native regressions cover oversized input allocations, aliased columns, shared string-view and dictionary payloads, heterogeneous nullable nested inputs, and cleanup on early or partially consumed writer errors. A mixed-reservation test covers a successful reservation followed by rejected growth, with both shared and independent input buffers and successful/failed writers.
Local validation on Linux with JDK 17:
cargo test --locked -p datafusion-comet-shuffle --lib -- --test-threads=1: 99 passed, none failed or ignored.cargo clippy --locked -p datafusion-comet-shuffle --all-targets -- -D warnings: passed.cargo fmt --all -- --checkandgit diff --check: passed.make core: passed before JVM validation../mvnw -B -Pspark-4.0 test spotless:check -Dtest=none -Dsuites=org.apache.spark.sql.comet.CometTaskMetricsSuite: all 14 tests passed in 1 suite, full reactor BUILD SUCCESS, including Spotless.Native builds used eight jobs with dev/test debug symbols disabled. Spark ran with
SPARK_LOCAL_IP=127.0.0.1andSPARK_LOCAL_HOSTNAME=localhost. Other Spark versions, live Celeborn, and benchmarks were not rerun locally.