Skip to content

[SPARK-59168][CORE] Avoid NoSuchElementException on missing keys in KVStore reads and writes #58466

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ulysses-you wants to merge 4 commits into
base: master
Choose a base branch
from
+91 −51
Open

[SPARK-59168][CORE] Avoid NoSuchElementException on missing keys in KVStore reads and writes #58466

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
b7ed268
[CORE] Avoid NoSuchElementException on missing keys in KVStore reads …
ulysses-you Sep 2, 2026
71ecbd3
[CORE] Carry missing key in KVStore iterator exceptions
ulysses-you Sep 2, 2026
f7a120d
[CORE] Rework KVStore missing-key handling around getOrNull
ulysses-you Sep 3, 2026
db615c0
[CORE] Route delete() through getOrNull and fix KVStore.delete javadoc
ulysses-you Sep 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion common/kvstore/src/main/java/org/apache/spark/util/kvstore/KVStore.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,11 +104,11 @@ public interface KVStore extends Closeable {

/**
* Removes an object and all data related to it, like index entries, from the store.
* Deleting a missing key is a no-op.
*
* @param type The object's type.
* @param naturalKey The object's "natural key", which uniquely identifies it. Null keys
* are not allowed.
* @throws java.util.NoSuchElementException If an element with the given key does not exist.
*/
void delete(Class<?> type, Object naturalKey) throws Exception;

Expand Down
44 changes: 19 additions & 25 deletions common/kvstore/src/main/java/org/apache/spark/util/kvstore/LevelDB.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,24 +98,16 @@ public LevelDB(File path, KVStoreSerializer serializer) throws Exception {
db().put(STORE_VERSION_KEY, serializer.serialize(STORE_VERSION));
}

Map<String, byte[]> aliases;
try {
aliases = get(TYPE_ALIASES_KEY, TypeAliases.class).aliases;
} catch (NoSuchElementException e) {
aliases = new HashMap<>();
}
TypeAliases aliasesValue = getOrNull(TYPE_ALIASES_KEY, TypeAliases.class);
Map<String, byte[]> aliases = aliasesValue != null ? aliasesValue.aliases : new HashMap<>();
typeAliases = new ConcurrentHashMap<>(aliases);

iteratorTracker = new ConcurrentLinkedQueue<>();
}

@Override
public <T> T getMetadata(Class<T> klass) throws Exception {
try {
return get(METADATA_KEY, klass);
} catch (NoSuchElementException nsee) {
return null;
}
return getOrNull(METADATA_KEY, klass);
}

@Override
Expand All @@ -128,11 +120,21 @@ public void setMetadata(Object value) throws Exception {
}

<T> T get(byte[] key, Class<T> klass) throws Exception {
byte[] data = db().get(key);
if (data == null) {
T value = getOrNull(key, klass);
if (value == null) {
throw new NoSuchElementException(new String(key, UTF_8));
}
return serializer.deserialize(data, klass);
return value;
}

/**
* Returns the value for the given key, or {@code null} if the key is not present, so that
* callers where a missing key is expected do not pay the cost of throwing and filling in an
* exception stack trace.
*/
<T> T getOrNull(byte[] key, Class<T> klass) throws Exception {
byte[] data = db().get(key);
return data != null ? serializer.deserialize(data, klass) : null;
}

private void put(byte[] key, Object value) throws Exception {
Expand Down Expand Up @@ -207,12 +209,7 @@ private void updateBatch(
Class<?> klass,
LevelDBTypeInfo.Index naturalIndex,
Collection<LevelDBTypeInfo.Index> indices) throws Exception {
Object existing;
try {
existing = get(naturalIndex.entityKey(null, value), klass);
} catch (NoSuchElementException e) {
existing = null;
}
Object existing = getOrNull(naturalIndex.entityKey(null, value), klass);

PrefixCache cache = new PrefixCache(value);
byte[] naturalKey = naturalIndex.toKey(naturalIndex.getValue(value));
Expand All @@ -229,9 +226,8 @@ public void delete(Class<?> type, Object naturalKey) throws Exception {
LevelDBTypeInfo ti = getTypeInfo(type);
byte[] key = ti.naturalIndex().start(null, naturalKey);
synchronized (ti) {
byte[] data = db().get(key);
if (data != null) {
Object existing = serializer.deserialize(data, type);
Object existing = getOrNull(key, type);
if (existing != null) {
PrefixCache cache = new PrefixCache(existing);
byte[] keyBytes = ti.naturalIndex().toKey(ti.naturalIndex().getValue(existing));
for (LevelDBTypeInfo.Index idx : ti.indices()) {
Expand All @@ -240,8 +236,6 @@ public void delete(Class<?> type, Object naturalKey) throws Exception {
db().write(batch);

@dongjoon-hyun dongjoon-hyun Sep 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

@ulysses-you ulysses-you Sep 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

}
}
} catch (NoSuchElementException nse) {
// Ignore.
}
}

Expand Down
44 changes: 19 additions & 25 deletions common/kvstore/src/main/java/org/apache/spark/util/kvstore/RocksDB.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,24 +131,16 @@ public RocksDB(File path, KVStoreSerializer serializer) throws Exception {
db().put(STORE_VERSION_KEY, serializer.serialize(STORE_VERSION));
}

Map<String, byte[]> aliases;
try {
aliases = get(TYPE_ALIASES_KEY, TypeAliases.class).aliases;
} catch (NoSuchElementException e) {
aliases = new HashMap<>();
}
TypeAliases aliasesValue = getOrNull(TYPE_ALIASES_KEY, TypeAliases.class);
Map<String, byte[]> aliases = aliasesValue != null ? aliasesValue.aliases : new HashMap<>();
typeAliases = new ConcurrentHashMap<>(aliases);

iteratorTracker = new ConcurrentLinkedQueue<>();
}

@Override
public <T> T getMetadata(Class<T> klass) throws Exception {
try {
return get(METADATA_KEY, klass);
} catch (NoSuchElementException nsee) {
return null;
}
return getOrNull(METADATA_KEY, klass);
}

@Override
Expand All @@ -161,11 +153,21 @@ public void setMetadata(Object value) throws Exception {
}

<T> T get(byte[] key, Class<T> klass) throws Exception {
byte[] data = db().get(key);
if (data == null) {
T value = getOrNull(key, klass);
if (value == null) {
throw new NoSuchElementException(new String(key, UTF_8));
}
return serializer.deserialize(data, klass);
return value;
}

/**
* Returns the value for the given key, or {@code null} if the key is not present, so that
* callers where a missing key is expected do not pay the cost of throwing and filling in an
* exception stack trace.
*/
<T> T getOrNull(byte[] key, Class<T> klass) throws Exception {
byte[] data = db().get(key);
return data != null ? serializer.deserialize(data, klass) : null;
}

private void put(byte[] key, Object value) throws Exception {
Expand Down Expand Up @@ -239,12 +241,7 @@ private void updateBatch(
Class<?> klass,
RocksDBTypeInfo.Index naturalIndex,
Collection<RocksDBTypeInfo.Index> indices) throws Exception {
Object existing;
try {
existing = get(naturalIndex.entityKey(null, value), klass);
} catch (NoSuchElementException e) {
existing = null;
}
Object existing = getOrNull(naturalIndex.entityKey(null, value), klass);

PrefixCache cache = new PrefixCache(value);
byte[] naturalKey = naturalIndex.toKey(naturalIndex.getValue(value));
Expand All @@ -261,9 +258,8 @@ public void delete(Class<?> type, Object naturalKey) throws Exception {
RocksDBTypeInfo ti = getTypeInfo(type);
byte[] key = ti.naturalIndex().start(null, naturalKey);
synchronized (ti) {
byte[] data = db().get(key);
if (data != null) {
Object existing = serializer.deserialize(data, type);
Object existing = getOrNull(key, type);
if (existing != null) {
PrefixCache cache = new PrefixCache(existing);
byte[] keyBytes = ti.naturalIndex().toKey(ti.naturalIndex().getValue(existing));
for (RocksDBTypeInfo.Index idx : ti.indices()) {
Expand All @@ -272,8 +268,6 @@ public void delete(Class<?> type, Object naturalKey) throws Exception {
db().write(writeOptions, writeBatch);

@dongjoon-hyun dongjoon-hyun Sep 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@ulysses-you ulysses-you Sep 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same, done in 1330280.

}
}
} catch (NoSuchElementException nse) {
// Ignore.
}
}

Expand Down
26 changes: 26 additions & 0 deletions common/kvstore/src/test/java/org/apache/spark/util/kvstore/LevelDBSuite.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,32 @@ public void testObjectWriteReadDelete() throws Exception {
assertEquals(0, countKeys(t.getClass()));
}

@Test
public void testGetOrNullMissingKey() throws Exception {

@dongjoon-hyun dongjoon-hyun Sep 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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"));

@ulysses-you ulysses-you Sep 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

// getOrNull() returns null for a missing key so expected misses (e.g. the write path
// looking up an existing entry) skip the cost of building an exception, while read()
// still surfaces a missing key as NoSuchElementException.
byte[] missingKey = db.getTypeInfo(CustomType1.class).naturalIndex().start(null, "missing");
assertNull(db.getOrNull(missingKey, CustomType1.class));
assertThrows(NoSuchElementException.class, () -> db.read(CustomType1.class, "missing"));

CustomType1 t = createCustomType1(1);
db.write(t);
assertEquals(t, db.read(CustomType1.class, t.key));
}

@Test
public void testDeleteEdgeCases() throws Exception {

@dongjoon-hyun dongjoon-hyun Sep 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@ulysses-you ulysses-you Sep 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

// Never-written type: type info is created on the fly, lookup misses, nothing happens.
db.delete(CustomType1.class, "missing");
assertEquals(0L, db.count(CustomType1.class));

// Never-written key of a written type.
db.write(createCustomType1(1));
db.delete(CustomType1.class, "missing");
assertEquals(1L, db.count(CustomType1.class));
}

@Test
public void testMultipleObjectWriteReadDelete() throws Exception {
CustomType1 t1 = createCustomType1(1);
Expand Down
26 changes: 26 additions & 0 deletions common/kvstore/src/test/java/org/apache/spark/util/kvstore/RocksDBSuite.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,32 @@ public void testObjectWriteReadDelete() throws Exception {
assertEquals(0, countKeys(t.getClass()));
}

@Test
public void testGetOrNullMissingKey() throws Exception {

@dongjoon-hyun dongjoon-hyun Sep 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@ulysses-you ulysses-you Sep 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same, done in 1330280.

// getOrNull() returns null for a missing key so expected misses (e.g. the write path
// looking up an existing entry) skip the cost of building an exception, while read()
// still surfaces a missing key as NoSuchElementException.
byte[] missingKey = db.getTypeInfo(CustomType1.class).naturalIndex().start(null, "missing");
assertNull(db.getOrNull(missingKey, CustomType1.class));
assertThrows(NoSuchElementException.class, () -> db.read(CustomType1.class, "missing"));

CustomType1 t = createCustomType1(1);
db.write(t);
assertEquals(t, db.read(CustomType1.class, t.key));
}

@Test
public void testDeleteEdgeCases() throws Exception {
// Never-written type: type info is created on the fly, lookup misses, nothing happens.
db.delete(CustomType1.class, "missing");
assertEquals(0L, db.count(CustomType1.class));

// Never-written key of a written type.
db.write(createCustomType1(1));
db.delete(CustomType1.class, "missing");
assertEquals(1L, db.count(CustomType1.class));
}

@Test
public void testMultipleObjectWriteReadDelete() throws Exception {
CustomType1 t1 = createCustomType1(1);
Expand Down