Skip to content

[BEAM-40264] Fix Spark MapState and SetState isEmpty after removals #40398

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

Merged
tvalentyn merged 5 commits into from
Oct 6, 2026
Merged

[BEAM-40264] Fix Spark MapState and SetState isEmpty after removals #40398

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
294d6c9
[Go SDK] Add portable logical type abstraction
dhruv-dagar Sep 11, 2026
4d82ecb
Fix Spark MapState and SetState emptiness after removals
dhruv-dagar Oct 4, 2026
d4a3aed
Add regression tests for state emptiness after removal
dhruv-dagar Oct 4, 2026
70a7a92
Merge branch 'apache:master' into fix/spark-state-is-empty
dhruv-dagar Oct 4, 2026
4b466bd
Remove unrelated upstream file from BEAM-40264 branch
dhruv-dagar Oct 4, 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
10 changes: 10 additions & 0 deletions runners/core-java/src/test/java/org/apache/beam/runners/core/StateInternalsTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -225,6 +225,11 @@ public void testSet() throws Exception {
assertThat(later.read(), hasItems("C", "D"));
assertFalse(later.contains("A").read());

value.remove("B");
value.remove("C");
value.remove("D");
assertTrue(value.isEmpty().read());

// clear
value.clear();
assertThat(value.read(), Matchers.emptyIterable());
Expand Down Expand Up @@ -389,6 +394,11 @@ public void testMap() throws Exception {
// isEmpty
assertFalse(value.isEmpty().read());

value.remove("B");
value.remove("D");
value.remove("E");
assertTrue(value.isEmpty().read());

// clear
value.clear();
assertThat(value.entries().read(), Matchers.emptyIterable());
Expand Down
12 changes: 10 additions & 2 deletions runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkStateInternals.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -432,7 +432,11 @@ private Map<MapKeyT, MapValueT> readAsMap() {
public void remove(MapKeyT key) {
Map<MapKeyT, MapValueT> sparkMapState = readAsMap();
sparkMapState.remove(key);
writeValue(sparkMapState);
if (sparkMapState.isEmpty()) {
clear();
} else {
writeValue(sparkMapState);
}
}

@Override
Expand Down Expand Up @@ -537,7 +541,11 @@ public ReadableState<Boolean> addIfAbsent(InputT input) {
public void remove(InputT input) {
Set<InputT> sparkSetState = readAsSet();
sparkSetState.remove(input);
writeValue(sparkSetState);
if (sparkSetState.isEmpty()) {

@saimanojsuper saimanojsuper Oct 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

nit: redundancy of the code can be improved using Collection Generic

clear();
} else {
writeValue(sparkSetState);
}
}

@Override
Expand Down
Loading