Skip to content

Streamline non-cached state backed iterable. #34746

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
robertwb merged 3 commits into from
May 7, 2025
Merged

Streamline non-cached state backed iterable. #34746

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
deeb4c4
Streamline non-cached state backed iterable.
robertwb Apr 25, 2025
6b8ef13
Address reviewer comments.
robertwb May 6, 2025
877f111
Remove tests wrongly asserting that reiteration is still cached.
robertwb May 6, 2025
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
4 changes: 1 addition & 3 deletions sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateBackedIterable.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,6 @@
import java.util.concurrent.ThreadLocalRandom;
import java.util.function.Supplier;
import org.apache.beam.fn.harness.Cache;
import org.apache.beam.fn.harness.Caches;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest;
import org.apache.beam.sdk.coders.IterableLikeCoder;
Expand Down Expand Up @@ -90,8 +89,7 @@ public StateBackedIterable(
StateRequest.newBuilder().setInstructionId(instructionId).setStateKey(stateKey).build();
this.prefix = prefix;
this.suffix =
StateFetchingIterators.readAllAndDecodeStartingFrom(
Caches.subCache(cache, stateKey), beamFnStateClient, request, elemCoder);
StateFetchingIterators.readAllAndDecodeStartingFrom(beamFnStateClient, request, elemCoder);
this.elemCoder = elemCoder;
}

Expand Down
79 changes: 79 additions & 0 deletions sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;

import com.google.auto.value.AutoValue;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
Expand All @@ -46,6 +48,7 @@
import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.AbstractIterator;

/**
* Adapters which convert a logical series of chunks using continuation tokens over the Beam Fn
Expand Down Expand Up @@ -91,6 +94,82 @@ public static <T> CachingStateIterable<T> readAllAndDecodeStartingFrom(
valueCoder);
}

public static <T> UncachedStateIterable<T> readAllAndDecodeStartingFrom(
BeamFnStateClient beamFnStateClient,
StateRequest stateRequestForFirstChunk,
Coder<T> valueCoder) {
return new UncachedStateIterable<>(beamFnStateClient, stateRequestForFirstChunk, valueCoder);
}

static class UncachedStateIterable<T> extends PrefetchableIterables.Default<T> {
Copy link
Contributor Author

@robertwb robertwb Apr 25, 2025

Choose a reason for hiding this comment

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

If this isn't used for state anywhere (I don't think it is) we probably don't even need to use the Prefethable* interfaces at all, simplifying things further.

Copy link
Contributor

@scwhittle scwhittle May 6, 2025

Choose a reason for hiding this comment

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

It is used for BagUserState, MultimapState etc. https://github.com/apache/beam/blob/master/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java#L76

Copy link
Contributor Author

@robertwb robertwb May 6, 2025

Choose a reason for hiding this comment

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

I meant to refer to the StateBackedIterable class itself. StateFetchingIterators are definitely used elsewhere.

Copy link
Contributor

@scwhittle scwhittle May 7, 2025

Choose a reason for hiding this comment

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

Sorry still not sure I'm understanding. Are you suggesting just inlining this class in StateBackedIterable to avoid having it use the interfaces?

Copy link
Contributor Author

@robertwb robertwb May 7, 2025

Choose a reason for hiding this comment

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

I meant, e.g. we could probably change declaration at https://github.com/apache/beam/blob/release-2.60.0/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateBackedIterable.java#L77 to be just an Iterable, changing the constructor for StateBackedIterable[Coder] to not take a Cache at all, etc. But it's also not too much work to keep it as a Prefetchable iterable given the underlying iterable is.

private final BeamFnStateClient beamFnStateClient;
private final StateRequest stateRequestForFirstChunk;
private final Coder<T> valueCoder;

public UncachedStateIterable(
BeamFnStateClient beamFnStateClient,
StateRequest stateRequestForFirstChunk,
Coder<T> valueCoder) {
this.beamFnStateClient = beamFnStateClient;
this.stateRequestForFirstChunk = stateRequestForFirstChunk;
this.valueCoder = valueCoder;
}

@Override
public PrefetchableIterator<T> createIterator() {
return new DecodingIterator<T>(
new LazyBlockingStateFetchingIterator(beamFnStateClient, stateRequestForFirstChunk),
valueCoder);
}

private static class DecodingIterator<T> extends AbstractIterator<T>
implements PrefetchableIterator<T> {
PrefetchableIterator<ByteString> chunkIterator;
InputStream currentChunk;
Coder<T> valueCoder;

public DecodingIterator(PrefetchableIterator<ByteString> chunkIterator, Coder<T> valueCoder) {
this.chunkIterator = chunkIterator;
this.currentChunk = ByteString.EMPTY.newInput();
this.valueCoder = valueCoder;
}

@Override
protected T computeNext() {
try {
while (currentChunk.available() == 0) {
if (chunkIterator.hasNext()) {
currentChunk = chunkIterator.next().newInput();
} else {
return endOfData();
}
}
return valueCoder.decode(currentChunk);
} catch (IOException exn) {
// Should never get here as ByteString.newInput() returns InputStreams
// that don't do actual IO operations.
throw new IllegalStateException(exn);
}
}

@Override
public boolean isReady() {
try {
return currentChunk.available() > 0 || chunkIterator.isReady();
} catch (IOException exn) {
throw new IllegalStateException(exn);
}
}

@Override
public void prefetch() {
if (!isReady()) {
chunkIterator.prefetch();
}
}
}
}

@VisibleForTesting
static class IterableCacheKey implements Weighted {

Expand Down
Loading