[cloud-fan]

What changes were proposed in this pull request?

Harden BytesToBytesMap when reset cannot allocate a replacement pointer array. Reset clears the map data, page, and growth state before allocation. If allocation fails, the map remains empty and enters an explicit failed-reset state. Array-dependent operations fail fast in that state; only a later successful reset returns the map to service.

Use a single lifecycle state to reject operations consistently after failed reset, destructive iteration, or free.

Add regression coverage for configured key operations in a terminal lifecycle state, and document the lifecycle exceptions on the iterator entry points.

Why are the changes needed?

Previously, a failed reset could leave a null pointer array and stale page state. A caller that caught SparkOutOfMemoryError and reused the map could hit NullPointerException or append into a freed page.

Lazy recovery from lookup also introduced allocation and SparkOutOfMemoryError into a read path, and could let a caller continue with an empty map while separately spilled aggregation data was lost. Requiring an explicit successful reset keeps the failure visible and makes recovery deliberate. Current in-tree callers still propagate SparkOutOfMemoryError from reset.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

build/sbt "core/testOnly *BytesToBytesMap*Suite"

Passed 46 tests total: 23 on-heap and 23 off-heap.

./dev/lint-java

Passed.

git diff --check

Passed.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Isaac (version not recorded), Codex (GPT-6).

[cloud-fan]

cc @LuciferYang @peter-toth

[peter-toth]

Thanks for the PR, @cloud-fan!

reset() freed the pointer array before reallocating it, so an OOM on that reallocation left longArray == null plus a stale currentPage. The fix moves the empty-state bookkeeping ahead of the allocation and has the next lookup re-allocate at initialCapacity. I re-derived the publication argument and it holds. A reader that observes longArrayRecoveryRequired == false is ordered after the allocate() that preceded that write, and a reader that observes true goes through the monitor instead. All 50 tests pass here, and the test counts in the description check out.

Blocking

• 1. Description claims a recovery that no in-tree caller can reach: The only production reset() caller is UnsafeKVExternalSorter.java:179, and every path into it lets the SparkOutOfMemoryError propagate (TungstenAggregationIterator.scala:247, :282, HashAggregateExec.scala:344, and HashAggregateExec.spillHashMapToSorter). SparkOutOfMemoryError extends OutOfMemoryError, so no NonFatal or case e: Exception up the stack catches it either, and nothing can resume the aggregation loop on the same map afterwards. Either name the caller I am missing, or set "Does this PR introduce any user-facing change?" to No and describe the change as hardening BytesToBytesMap's own contract for callers that do recover. This one is from reading rather than a repro. The experiment that would settle it is a test that forces map.reset() to fail during a spilling HashAggregateExec and asserts the query still completes.

Non-blocking

• 2. The destructiveIterator == null guard is untested: I deleted it and all 25 tests in each suite stayed green. A short test pins it. inline: BytesToBytesMap.java:616
• 3. The new contract is missing from the public javadoc: safeLookup still says the map must not be modified concurrently, while the lookup itself now allocates and can throw SparkOutOfMemoryError. reset() does not say that a failed reallocation leaves the map empty and self-healing. inline: BytesToBytesMap.java:654
• 4. Allocating under the map's monitor adds the reverse of SPARK-26265's lock order: It does not deadlock today, but the file documents that hazard for MapIterator and the new invariant is silent. inline: BytesToBytesMap.java:615
• 5. Self-healing is wired only into safeLookup: getArray() and maxNumKeysIndex() still hand out or dereference the null array in the failed-reset state, which is the same NPE on a different entry point. inline: BytesToBytesMap.java:609

Minor

• 6. The recovery flag is armed five statements after the invariant it guards breaks: Anything that throws in between leaves a null array with the recovery disarmed. inline: BytesToBytesMap.java:1185
• 7. lookupAfterFreeDoesNotRestoreArray pins -ea-only behaviour: Without assertions the same lookup throws a bare NPE. inline: AbstractBytesToBytesMapSuite.java:899

[peter-toth]

Finding 2. This guard has no test. I removed destructiveIterator == null from the condition and re-ran both suites on your head commit: all 25 tests in each stayed green, so nothing currently pins it.

This closes it, alongside the existing lookupAfterFreeDoesNotRestoreArray:

  @Test
  public void lookupAfterDestructiveIterationDoesNotRestoreArray() {
    memoryManager.limit(5000);
    BytesToBytesMap map =
      new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager, 256, 0.5, 4000);
    final long[] key = new long[]{1L};
    try {
      memoryManager.markExecutionAsOutOfMemoryOnce();
      assertThrows(SparkOutOfMemoryError.class, map::reset);
      map.destructiveIterator();
      memoryManager.limit(PAGE_SIZE_BYTES);
      assertThrows(
        AssertionError.class,
        () -> map.lookup(key, Platform.LONG_ARRAY_OFFSET, 8));
      assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
    } finally {
      map.free();
    }
  }

I ran it both ways. On your head it passes. With destructiveIterator == null deleted it fails with Expected java.lang.AssertionError to be thrown, but nothing was thrown, and the memory assertion would catch the stray array as well.

[cloud-fan]

Addressed in a380c19. I added lookupAfterDestructiveIterationDoesNotRestoreArray; the recovery path now fails explicitly with IllegalStateException and the test also verifies that no array memory was allocated.

[peter-toth]

Finding 3. Two javadoc gaps on public methods, both about behaviour this line introduces.

safeLookup's doc right above says "The map must not be modified concurrently", but this call makes the lookup itself modify the map. It allocates a LongArray, registers it with the TaskMemoryManager, and rewrites mask and growthThreshold. It can also throw SparkOutOfMemoryError, which nobody expects from a read method. UnsafeHashedRelation.get and getValue reach safeLookup on the broadcast-join probe path (sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala:235), so it is worth stating even though those maps are never reset.

reset()'s doc at :1160 gained the free() sentence but not the one that matters most here. A failed reallocation now leaves the map empty and the next lookup restores it. That contract currently lives only on the private restoreArrayAfterFailedReset, where a caller will not look for it.

[cloud-fan]

Addressed in a380c19. The reset() and both safeLookup() overloads now document failed-reset recovery and the possible SparkOutOfMemoryError. I also documented recovery on the two public array accessors.

[peter-toth]

Finding 4. This block holds the map's monitor across allocate(), which goes through TaskMemoryManager.allocatePage and can spill other consumers to disk. That establishes the lock order BytesToBytesMap -> TaskMemoryManager. TaskMemoryManager.spillConsumer uses the opposite order, calling consumerToSpill.spill(...) while holding the TaskMemoryManager monitor (core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java:313).

It does not deadlock today, and I checked why. BytesToBytesMap.spill is not synchronized, and MapIterator.spill takes the iterator's monitor rather than the map's, so nothing reachable from spillConsumer ever waits on this monitor. avoidDeadlock still passes with reset() and free() synchronized.

That is the same hazard advanceToNextPage spells out for MapIterator under SPARK-26265, and three new things now depend on it quietly: the synchronized here, on reset(), and on free(). A comment beside this synchronized recording the invariant ("nothing invoked under the TaskMemoryManager monitor may take this monitor") would stop the next change from breaking it silently.

[cloud-fan]

Addressed in a380c19. I added the lock-order invariant beside the synchronized recovery block: allocation enters TaskMemoryManager while holding the map monitor, so spill() and anything reachable from it must not acquire the map monitor.

[peter-toth]

Finding 5. The recovery only runs on the safeLookup paths, so the failed-reset state still reaches two other public methods with longArray == null.

• getArray(): assert(longArray != null) is a no-op without -ea, so it returns null. UnsafeKVExternalSorter calls it first thing (sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java:109) and then pointerArray.size(), which is the NPE this PR is fixing on a different entry point.
• maxNumKeysIndex(): dereferences longArray directly.

iterator(), destructiveIterator() and iteratorWithKeyIndex() are safe, because numValues == 0 after the reset keeps them from touching the array.

Calling restoreArrayAfterFailedReset() at the top of getArray() and maxNumKeysIndex() covers those two. The alternative is to state in the javadoc that only lookups self-heal, so a caller knows to do a lookup first.

[cloud-fan]

Addressed in a380c19. Both getArray() and maxNumKeysIndex() now invoke failed-reset recovery. arrayAccessorsRecoverAfterFailedReset independently exercises each entry point after a failed reset.

[peter-toth]

Finding 6. The invariant the whole recovery rests on is "longArray == null means recovery is required". It breaks at :1180 and is only re-established here, five statements later. If anything throws in between (a freeArray or freePage failure), the map is left with a null array and the flag still false, which disarms the recovery for good.

Arming it before the array is dropped removes the window:

    longArrayRecoveryRequired = true;
    if (longArray != null) {
      freeArray(longArray);
    }
    longArray = null;
    while (dataPages.size() > 0) {
      MemoryBlock dataPage = dataPages.removeLast();
      freePage(dataPage);
    }
    allocate(initialCapacity);
    longArrayRecoveryRequired = false;

[cloud-fan]

I kept the flag placement unchanged intentionally. Recovery is safe only after all old data pages have been removed. If freePage() unexpectedly failed midway, setting the flag earlier could let a later lookup allocate a fresh pointer array while stale pages remain in dataPages; subsequent iteration could then expose stale records. Since freeArray() and freePage() are not normal failure points, I prefer not to mark a partially cleaned map as recoverable.

[peter-toth]

Finding 7. This pins a JVM assertion, so it only describes the -ea build. Production runs without -ea, where the same lookup falls through to longArray.get(pos * 2) and throws a bare NPE. That is the failure mode the PR set out to remove.

Now that freed is tracked, making it explicit costs nothing. restoreArrayAfterFailedReset already reads freed on the fast path, so it can just be:

    if (freed) {
      throw new IllegalStateException("BytesToBytesMap has already been freed");
    }
    if (longArrayRecoveryRequired) {
      ...
    }

and the test asserts IllegalStateException instead. If you would rather not add a throw on the lookup path, a comment saying the AssertionError only happens under -ea would at least record what production does.

[cloud-fan]

Addressed in a380c19. Lookup after free() now throws an explicit IllegalStateException, and the regression test no longer depends on JVM assertions.

[cloud-fan]

Thanks for tracing the production callers. Addressed in a380c19: I updated the title and description to frame this as low-level BytesToBytesMap contract hardening, changed the user-facing answer to No, and explicitly noted that no in-tree caller currently catches the SparkOutOfMemoryError from reset().

[peter-toth]

Re-checked through a380c19. Findings 1, 2, 3, 4, 5 and 7 are resolved. The new lookupAfterDestructiveIterationDoesNotRestoreArray fails once the destructiveIterator guard is removed, so it pins it. Both suites pass here, 27 tests each.

I withdraw 6. Your stale-pages argument holds: with the flag armed early, a lookup could restore the array while old pages are still in dataPages, and iterator() walks those pages first.

Non-blocking

• 8. reset() after destructiveIterator() still brings the map back (late catch): The description says destructive iteration stays terminal, but only the recovery gate checks it. A reset() whose allocation succeeds makes the map fully usable again. inline: BytesToBytesMap.java:1194

[peter-toth]

Finding 8. The description says destructive iteration stays terminal, but only restoreArrayAfterFailedReset checks destructiveIterator, and only when a recovery is pending. reset() never checks it. After destructiveIterator(), a reset() that allocates successfully leaves longArrayRecoveryRequired == false, so the gate never runs and the map is fully usable again.

I checked it on this head: destructiveIterator(), then reset(), lookup() and append() all succeed. On master the same reset() failed with an NPE in freeArray(null). The new longArray != null check below is what lets it through now.

It also leaves destructiveIterator set. From reading, a spill() triggered by another consumer would then still go to the old iterator, which spills and frees the reused map's data pages.

Failing the same way as the recovery gate closes it:

Suggested change

	if (freed) {
	return;
	}
	if (freed) {
	return;
	}
	if (destructiveIterator != null) {
	throw new IllegalStateException(
	"BytesToBytesMap cannot be used after destructiveIterator() has been called");
	}

And a test for it:

  @Test
  public void resetAfterDestructiveIterationFails() {
    memoryManager.limit(PAGE_SIZE_BYTES);
    BytesToBytesMap map =
      new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager, 256, 0.5, 4000);
    try {
      map.destructiveIterator();
      assertThrows(IllegalStateException.class, map::reset);
      assertEquals(0L, taskMemoryManager.getMemoryConsumptionForThisTask());
    } finally {
      map.free();
    }
  }

I ran it both ways. It passes with the check above (56 tests green across both suites). Without the check it fails with Expected java.lang.IllegalStateException to be thrown, but nothing was thrown.

[cloud-fan]

Addressed in 64eaea7. reset() now throws IllegalStateException once destructive iteration has begun, before any reset bookkeeping or allocation. resetAfterDestructiveIterationFails pins the terminal state. Both map suites pass: 28 tests each, 56 total.

[LuciferYang]

Thanks for the rework — getArray()/maxNumKeysIndex() joining the recovery hook and the ISE on a freed/destructive map both look right, and they cover the two states I was going to raise (an assertion-only guard, and the accessors outside safeLookup). Three comments left, none blocking.

The two on BytesToBytesMap.java:676 and :1215 are about who the recovery is for rather than the mechanism, which reads fine. I could not find a production caller that resumes using the map after reset() throws, nor a map instance that sees both concurrent safeLookup and reset(), so I would like to know which caller you had in mind. I have not run the tests; this is all from reading the code.

[LuciferYang]

In production reset() has a single caller, UnsafeKVExternalSorter:179, and its call sites fall into two shapes, neither of which reaches reuse after a failed reset. The two mid-loop spills do re-look-up right after spilling (TungstenAggregationIterator:247 then :258; generated HashAggregateExec:993 then :997), but SparkOutOfMemoryError is an Error, NonFatal does not match it, and neither path has a catch, so the re-lookup line never runs. The two drain sites (TungstenAggregationIterator:282, HashAggregateExec:344) call free() immediately after spilling and never look up again.

The recovery logic itself reads fine to me. What I cannot match is the stated user-facing change, "a hash aggregation that reaches this failed-reset state can continue": with today's callers the task ends with SparkOutOfMemoryError either way. Did I miss a caller that keeps looking up after reset() throws, or is this intended as hardening? If it is hardening, one sentence in the description covers it.

[cloud-fan]

You did not miss a caller. This is intended as hardening of BytesToBytesMap's contract for callers that recover; current in-tree hash aggregation still propagates SparkOutOfMemoryError from reset(). The PR description now says there is no in-tree production caller that catches it and marks the user-facing change as No.

[LuciferYang]

Thanks — hardening is the right framing, and marking the user-facing change as No, with the note that no in-tree caller catches the error, resolves this from my side.

[LuciferYang]

I cannot find the combination this concurrent recovery is for. The only map that is ever reset() is the aggregation map, which is task-private and reached only through lookup(). The only map read by concurrent safeLookup is UnsafeHashedRelation.binaryMap, where asReadOnlyCopy hands each thread a wrapper around one shared map, and that one is only ever free()d, never reset. Point me at a usage I did not scan.

The flag is also published after the array is released: :1210 nulls longArray and only :1215 sets the flag, so a safeLookup landing in between gets neither. The contract at :647 rules this concurrency out anyway, so it cannot happen today.

Either narrow the comment to callers arriving after reset() returns, or move :1215 above :1207. The move costs nothing on the success path, since reset() holds the monitor throughout.

[cloud-fan]

Agreed that there is no current in-tree caller combining reset() with concurrent safeLookup(). The intent is to preserve safeLookup's thread-safe contract for callers that catch a failed reset. In 64eaea7 I clarified that concurrent retry is supported only after reset() has completed exceptionally; overlap with reset() remains prohibited. I kept the flag after page cleanup because moving it earlier could make a partially cleaned map recoverable if freePage() unexpectedly failed, allowing stale pages to remain visible.

[LuciferYang]

Agreed on both counts: the narrowed contract (concurrent retry only after reset() has returned exceptionally) covers the concern, and keeping the flag after page cleanup is the safer choice given a mid-cleanup freePage() failure. Withdrawing the move suggestion.

[LuciferYang]

The TODO block starting at :155 is about the 32-bit hashcode kept in longArray's odd slots ("we're wasting 32 bits of space here", and later "this array will fit in cache"); it used to sit directly under the field at :142. With the two new flags inserted in between, that "here" now reads as if it were about freed.

Moving the two flags below the TODO block fixes it. AGENTS.md puts it as: place a new member with the code it belongs with, rather than splitting a section of code that belonged together.

[cloud-fan]

Addressed in 64eaea7. The longArray TODO is directly below the field again, and the two new state flags now follow the complete TODO block.

[LuciferYang]

Confirmed on 64eaea7 — the TODO sits under longArray again with the two flags below it. Resolved.

[dongjoon-hyun]

I left a few inline comments.

[dongjoon-hyun]

This contract (the map stays empty and the next lookup recovers it) can silently drop data when combined with the only in-tree reset() caller. UnsafeKVExternalSorter's constructor first spills every map record into a local sorter via createWithExistingInMemorySorter(...) and calls map.reset() last. If reset() throws SparkOutOfMemoryError there, the constructor fails and the sorter holding the spilled groups never reaches TungstenAggregationIterator / HashAggregateExec. A caller that follows this documented contract (catch the OOM and keep using the map) then aggregates into a freshly recovered empty map, and every pre-spill group is missing from the result. Before this PR, the next lookup failed with an NPE and the task failed.

No in-tree caller catches this OOM today, so this is latent. But the new contract turns a loud failure into a silent wrong result for the one existing reset() call site. Could we either not advertise reuse after a failed reset(), or document that UnsafeKVExternalSorter callers must not recover?

[cloud-fan]

Addressed in 60eea74. Failed reset no longer self-heals through lookup. The map remains empty and unusable until the caller explicitly retries reset(), so catching the reset OOM and continuing with lookup fails fast instead of silently discarding sorter-held data.

[dongjoon-hyun]

With this call, lookup() / safeLookup() can now allocate memory, trigger spills of other consumers, and throw SparkOutOfMemoryError. UnsafeFixedWidthAggregationMap.getAggregationBufferFromUnsafeRow assumes memory exhaustion shows up only as loc.append(...) returning false (it returns null so the caller can call destructAndCreateExternalSorter()). In recovery mode the OOM escapes from the lookup instead and bypasses that spill fallback. Also, only the safeLookup javadocs gained @throws; the two lookup(...) overloads do not document it.

[cloud-fan]

Addressed in 60eea74. lookup() and safeLookup() no longer allocate during recovery or throw SparkOutOfMemoryError from that recovery path. They perform a state check and fail with IllegalStateException after a failed reset; only reset() allocates.

[dongjoon-hyun]

This destructiveIterator != null check only runs when longArrayRecoveryRequired is true, so the same illegal call behaves differently depending on history. After destructiveIterator() on a map that never had a failed reset, lookup() falls through to assert(longArray != null) (an NPE with assertions off), and getArray() silently returns null. After a failed reset, the same calls throw this IllegalStateException. lookupAfterDestructiveIterationDoesNotRestoreArray covers only the second case. Should the check sit outside the recovery-only branch, or be dropped from here?

[cloud-fan]

Addressed in 60eea74. A single lifecycle state is checked by all array-dependent operations regardless of history. Lookups, accessors, and iterators after destructive iteration now consistently throw IllegalStateException, and operationsAfterDestructiveIterationFail covers the normal destructive path.

[dongjoon-hyun]

reset() after free() is now a silent no-op, while lookup / getArray / maxNumKeysIndex throw IllegalStateException on a freed map. Before this PR, reset() after free() failed immediately (NPE in freeArray(null)). Now a caller that resets a freed map by mistake gets no signal at the point of the mistake, and numKeys() / numValues() stay stale after a "successful" reset. Also, because this check runs first, destructiveIterator(); free(); reset(); returns normally, which contradicts the new @throws IllegalStateException if destructive iteration has begun. Would throwing here be more consistent?

[cloud-fan]

Addressed in 60eea74. reset() after free() now throws IllegalStateException, as does reset() after destructive iteration. operationsAfterFreeFail covers reset and the array-dependent accessors.

[dongjoon-hyun]

The root cause seems to be that reset() discards the old pointer array before securing a replacement. How about fixing the ordering instead of accepting a null array and recovering lazily? Keep the counter/page reset as this PR does, then:

LongArray old = longArray;
try {
  allocate(initialCapacity);
  freeArray(old);
} catch (SparkOutOfMemoryError e) {
  old.zeroOut(); // or skip allocation entirely when old is already initial capacity
}

allocate() assigns longArray / growthThreshold / mask only after allocateArray returns, so a failed allocation leaves the old fields consistent. At reset() time the array borrowed by UnsafeKVExternalSorter is no longer referenced (inMemSorter = null after the spill, and that in-memory sorter has a null consumer). The map would always stay valid, and the two volatile flags, restoreArrayAfterFailedReset() with its four call sites, synchronized reset() / free(), the TestMemoryManager hook, and the ThreadMXBean test would all be unnecessary. The cost is briefly holding the old array plus a small initial-capacity array.

[cloud-fan]

I used the explicit-reset alternative in 60eea74. Rather than retaining and zeroing a potentially large old pointer array, which would work against reset releasing memory, reset frees the old storage and enters RESET_FAILED if replacement allocation fails. Only a later explicit reset can recover. This also removes all lazy-recovery machinery.

[dongjoon-hyun]

shutdownNow() does not wait for the workers. If an assertion above fails (e.g. assertTrue(contended, ...) on a loaded CI), the released second worker can already be in the lock-free probe loop of safeLookup when map.free() frees and nulls longArray. It then hits an NPE, or in the off-heap suite reads freed native memory. The success path is fine because both futures are get()-ed first. Adding executor.awaitTermination(...) before map.free() would close this window.

[cloud-fan]

Obsolete after 60eea74. The executor-based concurrent recovery test was removed with the lazy-recovery design, so the shutdown race no longer exists.

[dongjoon-hyun]

This setup append is what creates the data page the stale-currentPage regression check depends on, but its return value is ignored. With limit(5000), the 4096-byte pointer array leaves 904 bytes, so it only succeeds because TestMemoryManager grants a partial page. If that ever changes and append returns false, no page exists before reset(), and all later assertions still pass while the test no longer covers the currentPage / pageCursor reordering. Could we wrap it in assertTrue(...)?

[cloud-fan]

Addressed in 60eea74. The setup append is now wrapped in assertTrue(), preserving the stale-page regression coverage if memory-grant behavior changes.

[dongjoon-hyun]

This comment describes the pre-fix behavior in the present tense. After this PR, the bookkeeping (currentPage = null; pageCursor = 0;) runs before allocate() and safeLookup recovers first, so "the trailing reset() bookkeeping never runs" and "A subsequent lookup then NPEs" are no longer true. Maybe rephrase as "Before this fix, ..." or just state the invariant being tested?

[cloud-fan]

Addressed in 60eea74. The stale pre-fix narrative was replaced with the invariant being tested: failed reset clears old page state and lookup must not allocate implicitly.

[dongjoon-hyun]

This KeyOperationsFactory (constant hash 0, case-insensitive first byte) and the 'A' / 'a' fixtures duplicate supportsSemanticKeyOperations. The limit(5000) / new BytesToBytesMap(..., 256, 0.5, 4000) / markExecutionAsOutOfMemoryOnce() / assertThrows(SparkOutOfMemoryError.class, map::reset) block is also repeated in about eight new tests, with 256/4000 implicitly tied to the expected 512L and 256 elsewhere. A shared factory plus a small newMapWithFailedReset() helper would remove a good amount of duplication.

[cloud-fan]

Addressed in 60eea74. The duplicated semantic-key case and repeated failed-reset fixtures were removed. Coverage is consolidated into two failed-reset tests plus focused terminal-state tests, reducing the added test code to 141 lines.

[dongjoon-hyun]

Nit: this hunk only adds a blank line between longArray and the TODO that documents it. It looks like a leftover from moving the new fields; could we drop it so the TODO stays directly under its field?

[cloud-fan]

Addressed in 60eea74. The extra blank line is removed, so the TODO is directly below longArray again.

[dongjoon-hyun]

Thank you for working on this, @cloud-fan. Moving the canGrowArray / currentPage / pageCursor reset ahead of the allocation is a good fix for the stale-page issue.

Here is a summary of my review comments above.

Design

• The root cause looks like reset() discarding the old pointer array before securing a replacement. Allocating first and falling back to old.zeroOut() on OOM would keep the map always valid. It would also remove the lazy-recovery machinery: the volatile flags, restoreArrayAfterFailedReset(), synchronized reset() / free(), the new map -> TaskMemoryManager lock order, and the TestMemoryManager hook.
• The concurrent-recovery contract has no in-tree caller. The maps probed concurrently (broadcast UnsafeHashedRelation) are never reset.
• The only reset() caller is UnsafeKVExternalSorter. If a caller there recovered from the OOM, it would silently lose the spilled groups. This is latent, since no in-tree caller catches that OOM today.

Behavior

• lookup() / safeLookup() can now allocate memory and throw OOM. That bypasses the spill fallback in UnsafeFixedWidthAggregationMap, which expects append() to return false instead.
• The destructiveIterator check runs only in recovery mode.
• reset() after free() is a silent no-op. That is inconsistent with the other methods and with the new @throws.

Performance

• Every lookup on the hash aggregation / hash join hot paths now does two volatile reads, where one would be enough.

Tests

• safeLookupAfterFailedReset calls map.free() without an awaitTermination first.
• The setup append in reuseMapAfterFailedReset is unchecked.
• One test comment describes the pre-fix behavior.
• Some fixtures are duplicated.
• The TestMemoryManager hook's doc is incomplete, and the existing anonymous TaskMemoryManager subclass pattern could be reused instead.

Minor

• There is a stray blank line in BytesToBytesMap.
• The PR description gives the test count as 54, but it should be 56.

[dongjoon-hyun]

+1, LGTM.

Minor: the PR description (which becomes the commit message) says "This removes lazy allocation from lookup, the concurrent-recovery synchronization and volatile flags, and the shared TestMemoryManager test hook." None of these exist in master; they only existed in earlier iterations of this PR. Could we drop that sentence?

[dongjoon-hyun]

Location.append() is not guarded. With a Location obtained before a failed reset() or free(), append() only has assert (longArray != null), so without -ea it acquires a new data page first and then hits an NPE at longArray.get(...), leaving that page in dataPages. Since append() runs only for new keys, adding ensureReady() there should be cheap. Alternatively, documenting that a stale Location is unsupported would be enough.

[cloud-fan]

Addressed in 50c8f34. Location.append() now calls ensureReady() before allocation or memory access. I also added coverage for stale Location handles after failed reset and after free.

[dongjoon-hyun]

This (and the javadoc) still advertises retrying reset(). For the only in-tree caller, UnsafeKVExternalSorter, a failed reset() also drops the sorter holding the spilled groups, so a caller that retries reset() and keeps aggregating would silently produce a result without the pre-spill groups. Could we add a sentence that a caller retrying reset() is responsible for the validity of any state kept outside the map (e.g., spilled data)?

[cloud-fan]

Addressed in 50c8f34. The reset() javadoc now clarifies that a successful retry restores only the map, and callers remain responsible for external state such as spilled data.

[dongjoon-hyun]

nit. The @throws wording is inconsistent: the iterator methods say "if a reset failed, destructive iteration has begun, or the map has been freed", while lookup / safeLookup / getArray / maxNumKeysIndex say "if the map is not ready for use". Could we use one form?

[cloud-fan]

Addressed in 50c8f34. The array-dependent operations now use the same explicit lifecycle-state wording in their IllegalStateException documentation.

[dongjoon-hyun]

nit. This memoryManager.limit(PAGE_SIZE_BYTES) does not seem to be needed for this test.

[cloud-fan]

Addressed in 50c8f34. I removed the unnecessary memoryManager limit from operationsAfterFreeFail.

[dongjoon-hyun]

nit. This anonymous KeyOperations duplicates the pattern of the existing caseInsensitiveWithConstantHash fixture in this suite. It may be worth sharing a fixture, although the existing one is test-local.

[cloud-fan]

Thanks. I left this one unchanged. The existing fixture also tracks creation and thread ownership for its concurrency assertions, while this test only needs a configured factory whose operations must never be reached. Keeping the small fixture local avoids coupling those test concerns.

[cloud-fan]

@dongjoon-hyun I also removed the earlier-iteration sentence from the PR description. The lifecycle follow-up is in 50c8f34; Java checkstyle and git diff --check pass.

[dongjoon-hyun]

Could you fix the compilation failure, @cloud-fan ?

[cloud-fan]

the k8s failure is unrelated, thanks for the review, merging to master/4.x!

[cloud-fan]

Merge Summary:

• merged into master 5a05bff
• merged into branch-4.x e0fc8de

Posted by merge_spark_pr.py