Skip to content

fix: make native shuffle spill metrics independent of input batching - #5628

Open
sunchao wants to merge 2 commits into
apache:mainfrom
sunchao:dev/chao/codex/issue5617-shared-spill-metrics
Open

fix: make native shuffle spill metrics independent of input batching#5628
sunchao wants to merge 2 commits into
apache:mainfrom
sunchao:dev/chao/codex/issue5617-shared-spill-metrics

Conversation

@sunchao

@sunchao sunchao commented Sep 1, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #5617.

Rationale for this change

Native shuffle deduplicates spilled Arrow allocations only within one outer insert_batch call. 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 -- --check and git diff --check: passed.
  • make core: passed before JVM validation.
  • 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: 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.1 and SPARK_LOCAL_HOSTNAME=localhost. Other Spark versions, live Celeborn, and benchmarks were not rerun locally.

@sunchao sunchao changed the title fix: measure native shuffle spills from materialized batches fix: make native shuffle spill metrics independent of input batching Sep 1, 2026
@sunchao
sunchao requested a review from andygrove September 2, 2026 16:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Native shuffle memory_spilled_bytes charges a shared Arrow allocation once per input batch

1 participant