[SPARK-59168][CORE] Avoid NoSuchElementException on missing keys in KVStore reads and writes by ulysses-you · Pull Request #58466 · apache/spark

ulysses-you · 2026-09-02T07:24:25Z

What changes were proposed in this pull request?

Keep RocksDB.get and LevelDB.get throwing NoSuchElementException for a missing key, and add a null-returning getOrNull sibling for the callers where a missing key is expected:

The store constructor (type aliases), getMetadata, updateBatch, and delete use getOrNull instead of catching NoSuchElementException.
delete() drops catch blocks that nothing inside their try bodies can throw; deleting a missing entry remains a silent no-op, and the KVStore.delete javadoc now says so instead of claiming a NoSuchElementException.

Why are the changes needed?

updateBatch looks up the existing value via get before writing. During an event-log rebuild most writes are new entries, so the key is absent and get throws NoSuchElementException on essentially every write. That exception is immediately caught, but constructing it still fills in a full stack trace, using exceptions for ordinary control flow. In a History Server CPU profile this showed up as a measurable share of CPU spent in Throwable.fillInStackTrace under RocksDB.get. getOrNull avoids building the exception entirely.

Does this PR introduce any user-facing change?

No. The public KVStore behavior is unchanged: read still throws NoSuchElementException for a missing key, and getMetadata still returns null when no metadata is present.

How was this patch tested?

Added to both RocksDBSuite and LevelDBSuite:

testGetOrNullMissingKey: getOrNull returns null for a missing key, while read still throws NoSuchElementException for it.
testDeleteEdgeCases: deleting a never-written type or a missing key is a silent no-op.

Existing kvstore test suites pass.

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

Yes. Generated-by: Claude Code 2.1.259

ulysses-you · 2026-09-02T11:35:49Z

cc @LuciferYang @dongjoon-hyun thank you

LuciferYang

Reviewed the kvstore changes. Overall this looks clean: the enumeration of the package-private get(byte[], Class) call sites is complete, the public contract (read() throwing NoSuchElementException, getMetadata() returning null) is preserved and matches InMemoryStore, and the perf motivation holds since both write() and writeAll() go through updateBatch on the replay path. Only two minor non-blocking nits below.

dongjoon-hyun · 2026-09-02T13:44:30Z

Thank you for the PR. I went through the diff and the surrounding code paths in common/kvstore and the History Server callers. I didn't find any correctness issue: every removed try/catch is re-established by a null check, and read() / iterator next() still throw the same NoSuchElementException with the same message as before. A few non-blocking comments below.

Design / simplification

Consider keeping the throwing get() and adding a null-returning sibling instead. Only three callers ever swallowed the NoSuchElementException (constructor, getMetadata, updateBatch), and updateBatch is the hot path. If we keep get() as-is and add e.g. getOrNull(byte[], Class) for those three, read(), RocksDBIterator and LevelDBIterator need no change at all, and the missing-key exception is built in one place per store instead of four (RocksDB.read, LevelDB.read, RocksDBIterator.next, LevelDBIterator.next). It also drops the two new UTF_8 static imports from the iterators.
delete() still has a dead catch (NoSuchElementException nse) { // Ignore. } (RocksDB.java:271, LevelDB.java:239). Nothing in that try body can throw it (it uses the raw db().get(key) with a null guard). It was already dead before this PR, but since this PR removes every other NSEE catch in these files and documents that get() returns null, it would be nice to drop it here too.
testNextAfterEntityDelete hardens "hasNext() is true, then next() throws NoSuchElementException" as the contract. That is the pre-existing behavior, so no objection to the test itself, just noting that NSEE from Iterator.next() conventionally means exhaustion, and in the live-UI RocksDB store case a concurrent delete by ElementTrackingStore would surface to the REST layer as a 404. Skipping dangling index entries in loadNext() (or using a distinct exception type) could be a follow-up.

Efficiency (pre-existing, on the same write path; fine as follow-ups)

updateBatch computes the natural key three times per write: naturalIndex.entityKey(null, value), then naturalIndex.toKey(naturalIndex.getValue(value)), then again inside Index.addOrRemove for the natural index. Computing naturalKey first and building the lookup key with ti.buildKey(false, naturalIndex.keyPrefix(null), naturalKey) (the same construction RocksDBIterator.next uses) would remove two reflective accessor calls and two encodings per write.
updateBatch deserializes the existing entity even for natural-index-only types (AppSummary, PoolData, SparkPlanGraphWrapper, ...), where existing is only used as a presence flag. A raw db().get null check would suffice there.
In next(), ti.naturalIndex().keyPrefix(null) allocates a fresh prefix on every element although it is constant for the iterator's lifetime; the constructor already caches indexKeyPrefix the same way.
Outside this PR: FsHistoryProvider still uses listing.read(...) + catch NoSuchElementException as the "not tracked yet" check for every newly discovered log (checkForLogs, checkAndCleanLog, cleanDriverLogs, addListing). It is orders of magnitude smaller than the updateBatch cost fixed here, but a null-returning read exposed through KVStore would remove it. Probably a separate JIRA.

PR description

The "How was this patch tested?" section still lists only testGetMissingKeyReturnsNull and does not mention testNextAfterEntityDelete or the iterator message change from the second commit.

ulysses-you · 2026-09-03T01:58:32Z

Thanks @dongjoon-hyun for the detailed review.

Adopted in b5ac972: get() keeps throwing, and a null-returning getOrNull() now serves the constructor, getMetadata and updateBatch. read() and both iterators stay at the base shape, so the exception is built in one place per store.
Verified dead (nothing inside the try bodies can throw NoSuchElementException) and dropped in b5ac972; testDeleteEdgeCases pins down that deleting a missing entry is a silent no-op.
Agreed; keeping the existing behavior. Skipping dangling index entries or using a distinct exception type works as a follow-up.

4./5./6./7. Agreed, left for a follow-up to keep this PR focused.

dongjoon-hyun

Thank you for the update. I re-reviewed b5ac972 end to end, including the serializer.deserialize paths under KVStoreScalaSerializer / KVStoreProtobufSerializer, and I still don't see any correctness issue: with Jackson's default WRAP_EXCEPTIONS a NoSuchElementException from a nested deserializer is wrapped into JsonMappingException, so the old catch (NoSuchElementException) blocks never covered anything beyond get()'s own throw, which getOrNull now expresses directly.

A few remaining non-blocking comments, plus two that can't be attached inline:

KVStore.delete Javadoc (KVStore.java:111) still says @throws java.util.NoSuchElementException If an element with the given key does not exist. None of the three implementations throw, every caller in core relies on the no-op, and this PR now pins the no-op with testDeleteEdgeCases while removing the only catch that ever hinted at the documented behavior. It would be good to fix that @throws line in the same PR (e.g. "Deleting a missing key is a no-op.").
PR description: .github/PULL_REQUEST_TEMPLATE asks for Generated-by: followed by the tool name and its version; the description currently has only Generated-by: Claude Code.

dongjoon-hyun · 2026-09-03T03:46:00Z

@@ -240,8 +238,6 @@ public void delete(Class<?> type, Object naturalKey) throws Exception {
          db().write(batch);


delete() (a few lines above, at byte[] data = db().get(key); if (data != null) { Object existing = serializer.deserialize(data, type); ... }) is exactly the body of the new getOrNull. Since this PR already touches delete(), how about

Object existing = getOrNull(key, type); if (existing != null) {

so the read-or-null sequence has a single home per store? It is behavior-identical (same synchronized (ti) block; the Class<?> capture already compiles the same way in updateBatch).

Done in 1330280: Object existing = getOrNull(key, type).

dongjoon-hyun · 2026-09-03T03:46:00Z

@@ -272,8 +270,6 @@ public void delete(Class<?> type, Object naturalKey) throws Exception {
          db().write(writeOptions, writeBatch);


Same as LevelDB.delete: db().get(key) + serializer.deserialize(data, type) here can be Object existing = getOrNull(key, type); if (existing != null) { ... }.

Same, done in 1330280.

dongjoon-hyun · 2026-09-03T03:46:01Z

+   * callers where a missing key is expected do not pay the cost of throwing and filling in an
+   * exception stack trace.
+   */
+  @VisibleForTesting


nit: @VisibleForTesting reads as "widened only for tests", but getOrNull is a production helper with four callers in this class (constructor, getMetadata, get, updateBatch) and its visibility is the same package-private as the unannotated get() right above. I'd drop the annotation to match get().

Dropped in 1330280.

dongjoon-hyun · 2026-09-03T03:46:01Z

+   * callers where a missing key is expected do not pay the cost of throwing and filling in an
+   * exception stack trace.
+   */
+  @VisibleForTesting


nit: same as LevelDB.getOrNull; @VisibleForTesting can be dropped.

Dropped in 1330280.

dongjoon-hyun · 2026-09-03T03:46:01Z

  }

+  @Test
+  public void testGetOrNullMissingKey() throws Exception {


Only the assertNull(db.getOrNull(missingKey, ...)) line is new coverage here. The get()-throws half and the present-key round trip are already asserted by testObjectWriteReadDelete through db.read / db.write, and read() builds exactly this key (getTypeInfo(klass).naturalIndex().start(null, naturalKey)). Hand-building the key twice per suite couples the test to the key layout. I'd keep the assertNull and use the public API for the rest:

assertThrows(NoSuchElementException.class, () -> db.read(CustomType1.class, "missing"));

Done in 1330280: only the assertNull keeps the hand-built key; the throws and present-key assertions go through db.read.

dongjoon-hyun · 2026-09-03T03:46:01Z

  }

+  @Test
+  public void testGetOrNullMissingKey() throws Exception {


Same as LevelDBSuite.testGetOrNullMissingKey: only the assertNull is new; the other assertions can go through db.read.

Same, done in 1330280.

dongjoon-hyun · 2026-09-03T03:46:01Z

+  }
+
+  @Test
+  public void testDeleteEdgeCases() throws Exception {


Just noting for the record: all three cases short-circuit at the pre-existing if (data != null) guard in delete(), so this test passes identically at the merge base and at this commit. It documents the no-op contract (which is useful given the KVStore Javadoc mismatch), but it isn't a regression guard for the removed catch. Fine to keep as-is.

Kept as-is; it now documents the same no-op contract the corrected KVStore.delete javadoc states.

dongjoon-hyun · 2026-09-03T03:46:01Z

+    db.delete(CustomType1.class, "missing");
+    assertEquals(1L, db.count(CustomType1.class));
+
+    // Mismatched key type: the encoded lookup key misses, nothing is removed.


nit: Integer is a legal key type in Index.toKey (sign marker + hex), so this is not really a "mismatched key type"; it's another never-written key in a different encoding. Maybe reword the comment, or drop the case since the one above already covers a missing key.

Dropped the case in 1330280.

dongjoon-hyun · 2026-09-03T03:46:01Z

+    db.delete(CustomType1.class, "missing");
+    assertEquals(1L, db.count(CustomType1.class));
+
+    // Mismatched key type: the encoded lookup key misses, nothing is removed.


nit: same as LevelDBSuite; 42 is a valid Integer key, so "mismatched key type" is misleading.

Same, dropped in 1330280.

ulysses-you · 2026-09-03T04:00:51Z

Both done in 1330280:

The @throws line is removed; the KVStore.delete javadoc now states that deleting a missing key is a no-op.
Updated to Generated-by: Claude Code 2.1.259.

dongjoon-hyun

+1, LGTM. Thank you, @ulysses-you .

…and writes RocksDB.get and LevelDB.get now return null instead of throwing NoSuchElementException for a missing key, letting callers decide whether a missing key is an error. updateBatch previously threw and filled in an exception stack trace on every write of a new entry, and most writes during an event log rebuild are new entries, so this removes that overhead. read keeps throwing NoSuchElementException to preserve the KVStore contract; getMetadata, the constructor and the secondary-index iterator path handle the null directly. Assisted-by: Claude Opus 4.8

The secondary-index iterator paths resolve entries through a live get() and now include the missing key in the NoSuchElementException message, matching read(). Adds testNextAfterEntityDelete to both DB suites. Assisted-by: Claude Opus 4.8

Keep get() throwing NoSuchElementException and add a null-returning getOrNull() for callers where a missing key is expected: the store constructor, getMetadata and updateBatch. Also drop the catch blocks in delete() that nothing inside their try bodies can throw. Assisted-by: Claude Opus 4.8

@VisibleForTesting

delete() reads the existing entry through getOrNull, getOrNull drops the @VisibleForTesting annotation, and KVStore.delete javadoc states that deleting a missing key is a no-op. The new tests assert through the public API. Assisted-by: Claude Opus 4.8

uros-b · 2026-09-03T14:02:35Z

+1, thank you @ulysses-you and @LuciferYang @dongjoon-hyun!

LuciferYang reviewed Sep 2, 2026

View reviewed changes

LuciferYang approved these changes Sep 2, 2026

View reviewed changes

dongjoon-hyun reviewed Sep 3, 2026

View reviewed changes

dongjoon-hyun approved these changes Sep 3, 2026

View reviewed changes

ulysses-you added 4 commits September 3, 2026 15:57

ulysses-you force-pushed the worktree-rockdb branch from 1330280 to db615c0 Compare September 3, 2026 07:57

uros-b approved these changes Sep 3, 2026

View reviewed changes

		@@ -240,8 +238,6 @@ public void delete(Class<?> type, Object naturalKey) throws Exception {
		db().write(batch);

		@@ -272,8 +270,6 @@ public void delete(Class<?> type, Object naturalKey) throws Exception {
		db().write(writeOptions, writeBatch);

Uh oh!

Conversation

ulysses-you commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

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

Uh oh!

ulysses-you commented Sep 2, 2026

Uh oh!

LuciferYang left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dongjoon-hyun commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ulysses-you commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

dongjoon-hyun left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ulysses-you commented Sep 3, 2026

Uh oh!

dongjoon-hyun left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

uros-b commented Sep 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ulysses-you commented Sep 2, 2026 •

edited

Loading

dongjoon-hyun commented Sep 2, 2026 •

edited

Loading

ulysses-you commented Sep 3, 2026 •

edited

Loading

dongjoon-hyun left a comment •

edited

Loading