Examples with Materializations org.apache.beam.sdk.transforms.Materializations used on opensource projects

Example 1 with Materializations

use of org.apache.beam.sdk.transforms.Materializations in project beam by apache.

the class FnApiStateAccessor method get.

@Override
@Nullable
public <T> T get(PCollectionView<T> view, BoundedWindow window) {
    TupleTag<?> tag = view.getTagInternal();
    SideInputSpec sideInputSpec = sideInputSpecMap.get(tag);
    checkArgument(sideInputSpec != null, "Attempting to access unknown side input %s.", view);
    ByteString.Output encodedWindowOut = ByteString.newOutput();
    try {
        sideInputSpec.getWindowCoder().encode(sideInputSpec.getWindowMappingFn().getSideInputWindow(window), encodedWindowOut);
    } catch (IOException e) {
        throw new IllegalStateException(e);
    }
    ByteString encodedWindow = encodedWindowOut.toByteString();
    StateKey.Builder cacheKeyBuilder = StateKey.newBuilder();
    switch(sideInputSpec.getAccessPattern()) {
        case Materializations.ITERABLE_MATERIALIZATION_URN:
            cacheKeyBuilder.getIterableSideInputBuilder().setTransformId(ptransformId).setSideInputId(tag.getId()).setWindow(encodedWindow);
            break;
        case Materializations.MULTIMAP_MATERIALIZATION_URN:
            checkState(sideInputSpec.getCoder() instanceof KvCoder, "Expected %s but received %s.", KvCoder.class, sideInputSpec.getCoder().getClass());
            cacheKeyBuilder.getMultimapKeysSideInputBuilder().setTransformId(ptransformId).setSideInputId(tag.getId()).setWindow(encodedWindow);
            break;
        default:
            throw new IllegalStateException(String.format("This SDK is only capable of dealing with %s materializations " + "but was asked to handle %s for PCollectionView with tag %s.", ImmutableList.of(Materializations.ITERABLE_MATERIALIZATION_URN, Materializations.MULTIMAP_MATERIALIZATION_URN), sideInputSpec.getAccessPattern(), tag));
    }
    return (T) stateKeyObjectCache.computeIfAbsent(cacheKeyBuilder.build(), key -> {
        switch(sideInputSpec.getAccessPattern()) {
            case Materializations.ITERABLE_MATERIALIZATION_URN:
                return sideInputSpec.getViewFn().apply(new IterableSideInput<>(getCacheFor(key), beamFnStateClient, processBundleInstructionId.get(), key, sideInputSpec.getCoder()));
            case Materializations.MULTIMAP_MATERIALIZATION_URN:
                return sideInputSpec.getViewFn().apply(new MultimapSideInput<>(getCacheFor(key), beamFnStateClient, processBundleInstructionId.get(), key, ((KvCoder) sideInputSpec.getCoder()).getKeyCoder(), ((KvCoder) sideInputSpec.getCoder()).getValueCoder()));
            default:
                throw new IllegalStateException(String.format("This SDK is only capable of dealing with %s materializations " + "but was asked to handle %s for PCollectionView with tag %s.", ImmutableList.of(Materializations.ITERABLE_MATERIALIZATION_URN, Materializations.MULTIMAP_MATERIALIZATION_URN), sideInputSpec.getAccessPattern(), tag));
        }
    });
}