viirya opened a new pull request, #59331:
URL: https://github.com/apache/spark/pull/59331
### What changes were proposed in this pull request?
`TaskMemoryManager` tracks spillable consumers in a
`HashSet<MemoryConsumer>`, which deduplicates by `equals`/`hashCode`. This PR
switches it to an identity-based set (`Collections.newSetFromMap(new
IdentityHashMap<>())`), so every consumer instance is tracked regardless of how
it defines equality.
The set is only used for `add`, iteration and `clear` (no
`contains`/`remove`), and the spill heuristic already compares the requesting
consumer by reference (`c == requestingConsumer`), so the switch does not
affect any other logic. This is the only hash-based collection keyed by
`MemoryConsumer` in the codebase.
### Why are the changes needed?
With value-based equality, an equal but distinct consumer is silently
dropped from the set. That instance:
- is never offered for spilling, so a memory request can fail to free memory
that is actually spillable; and
- is missing from the per-consumer memory breakdown in
`UNABLE_TO_ACQUIRE_MEMORY` errors, where its memory shows up as "not attributed
to a specific consumer".
This affects master today: `HybridRowQueue` is a Scala case class extending
`MemoryConsumer`, so two queues created with the same arguments in the same
task are deduplicated. #59282 fixes this for `HybridRowQueue` specifically; as
suggested by @dongjoon-hyun in that review, this PR fixes it in one place for
all current and future consumers. The two PRs are independent and can be merged
in either order.
### Does this PR introduce _any_ user-facing change?
Yes, as a bug fix: distinct memory consumers that compare equal (currently
`HybridRowQueue` instances with identical arguments in the same task) are now
all eligible for spilling and are all listed in the OOM memory breakdown. There
is no API or configuration change.
### How was this patch tested?
Added two tests to `TaskMemoryManagerSuite` using a test consumer with
value-based `equals`/`hashCode`:
- `equalButDistinctConsumersAreAllSpilled`: before this change, the second
equal consumer was not spilled (50 B left).
- `equalButDistinctConsumersAreAllInMemoryConsumptionBreakdown`: before this
change, the second consumer's memory was reported as not attributed to a
specific consumer.
Both fail without the fix and pass with it. Also ran the
`org.apache.spark.memory` suites, `UnsafeExternalSorterSuite`,
`ShuffleExternalSorterSuite`, `BytesToBytesMap{On,Off}HeapSuite`,
`ExternalSorterSpillSuite` and `ExternalAppendOnlyMapSuite`, plus
`core/checkstyle` and `core/Test/checkstyle`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
This pull request and its description were written by Isaac.
--
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]
