qingfureal opened a new pull request, #58801:
URL: https://github.com/apache/spark/pull/58801
### What changes were proposed in this pull request?
`UnsafeSorterIterator.getNumRecords()` has two contracts, and the abstract
method documents neither. Four of the five implementations report the *total*
number of records and never change it as the iterator is consumed:
| Implementation | `getNumRecords()` returns | Decrements on `loadNext()`? |
|---|---|---|
| `UnsafeSorterSpillReader` | total from the spill file header | no |
| `UnsafeInMemorySorter.SortedIterator` | total | no |
| `UnsafeSorterSpillMerger` (anonymous) | sum of the counts added | no |
| `UnsafeExternalSorter.ChainedIterator` | sum taken at construction | no |
| `UnsafeExternalSorter.SpillableIterator` | records **remaining** | **yes**
|
`UnsafeSorterSpillMerger.addSpillIfNotEmpty` reads the count *after* calling
`loadNext()`. This PR moves the read to before it:
```java
public void addSpillIfNotEmpty(UnsafeSorterIterator spillReader) throws
IOException {
if (spillReader.hasNext()) {
+ numRecords += spillReader.getNumRecords();
spillReader.loadNext();
priorityQueue.add(spillReader);
- numRecords += spillReader.getNumRecords();
}
}
```
and documents the contract on `UnsafeSorterIterator.getNumRecords()`.
### Why are the changes needed?
`UnsafeExternalSorter.getSortedIterator()` adds its in-memory
`SpillableIterator` to the merger alongside the spill readers. Because the
count was read after `loadNext()` had already decremented it, the merged
iterator reported one record fewer than it will actually produce, per in-memory
iterator. The two readings agree only while an iterator is untouched, so
reading the count first is correct under either contract.
To be clear about the impact: **there is none today.** No caller in the
repository reads a merged iterator's `getNumRecords()` — the value is computed
and never consumed. Iteration itself has always been correct, because the
merged iterator's `hasNext()` is driven by the priority queue rather than the
counter. The other call sites that do read `getNumRecords()` all read fresh,
unconsumed iterators, so none of them is affected. That is also why this has
gone unnoticed since the line was introduced in SPARK-12295 (2016).
It is still a wrong value, and a trap for the next caller who hands a merged
iterator to `UnsafeSorterSpillWriter`, which enforces the declared count and
fails with `IllegalStateException: Number of records written exceeded
numRecordsToWrite`. The undocumented split contract is what produced the bug,
so this PR documents it rather than only fixing the one symptom.
### Does this PR introduce _any_ user-facing change?
No. `getNumRecords()` on the merged iterator now returns the correct value;
no public API, behaviour or output changes.
### How was this patch tested?
Two tests added to `UnsafeExternalSorterSuite`:
- `testGetNumRecordsCountsUnspilledRecords` — two spills plus five records
still in memory; asserts `getSortedIterator().getNumRecords()` is 15 and that
the iterator yields exactly that many records.
- `testGetNumRecordsWithoutUnspilledRecords` — the all-spilled case, which
was already correct, pinned against regression.
Verified the first test actually catches the bug. With the one-line change
reverted:
```
[ERROR] UnsafeExternalSorterSuite.testGetNumRecordsCountsUnspilledRecords
org.opentest4j.AssertionFailedError: expected: <15> but was: <14>
```
`testGetNumRecordsWithoutUnspilledRecords` still passes with the fix
reverted, confirming the miscount is specific to the in-memory iterator path.
With the fix, `UnsafeExternalSorterSuite` (28),
`UnsafeExternalSorterRadixSortSuite` (28) and `UnsafeInMemorySorterSuite` (3)
all pass, and `dev/lint-java` is clean.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5)
--
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]
