Examples with MultimapViewFn2 org.apache.beam.sdk.values.PCollectionViews.MultimapViewFn2 used on opensource projects

Example 1 with MultimapViewFn2

use of org.apache.beam.sdk.values.PCollectionViews.MultimapViewFn2 in project beam by apache.

the class IsmSideInputReader method get.

@Override
public <ViewT> ViewT get(final PCollectionView<ViewT> view, final BoundedWindow window) {
    @SuppressWarnings("rawtypes") final TupleTag tag = view.getTagInternal();
    checkArgument(tagToIsmReaderMap.containsKey(tag), "calling getSideInput() with unknown view");
    // for its use.
    try {
        ViewFn<?, ?> viewFn = view.getViewFn();
        // back to null for the user.
        if (viewFn instanceof SingletonViewFn || viewFn instanceof SingletonViewFn2) {
            ViewT rval = executionContext.<PCollectionViewWindow<ViewT>, ViewT>getLogicalReferenceCache().get(PCollectionViewWindow.of(view, window), () -> {
                @SuppressWarnings("unchecked") ViewT viewT = getSingletonForWindow(tag, (HasDefaultValue<ViewT>) viewFn, window);
                @SuppressWarnings("unchecked") ViewT nullPlaceHolder = (ViewT) NULL_PLACE_HOLDER;
                return viewT == null ? nullPlaceHolder : viewT;
            });
            return rval == NULL_PLACE_HOLDER ? null : rval;
        } else if (singletonMaterializedTags.contains(tag)) {
            checkArgument(viewFn instanceof MapViewFn || viewFn instanceof MapViewFn2 || viewFn instanceof MultimapViewFn || viewFn instanceof MultimapViewFn2, "Unknown view type stored as singleton. Expected one of %s, got %s", KNOWN_SINGLETON_VIEW_TYPES, viewFn.getClass().getName());
            return executionContext.<PCollectionViewWindow<ViewT>, ViewT>getLogicalReferenceCache().get(PCollectionViewWindow.of(view, window), () -> {
                return getMapSingletonForViewAndWindow(tag, window);
            });
        } else {
            return executionContext.<PCollectionViewWindow<ViewT>, ViewT>getLogicalReferenceCache().get(PCollectionViewWindow.of(view, window), () -> {
                if (viewFn instanceof IterableViewFn || viewFn instanceof IterableViewFn2 || viewFn instanceof ListViewFn || viewFn instanceof ListViewFn2) {
                    @SuppressWarnings("unchecked") ViewT viewT = (ViewT) getListForWindow(tag, window);
                    return viewT;
                } else if (viewFn instanceof MapViewFn || viewFn instanceof MapViewFn2) {
                    @SuppressWarnings("unchecked") ViewT viewT = (ViewT) getMapForWindow(tag, window);
                    return viewT;
                } else if (viewFn instanceof MultimapViewFn || viewFn instanceof MultimapViewFn2) {
                    @SuppressWarnings("unchecked") ViewT viewT = (ViewT) getMultimapForWindow(tag, window);
                    return viewT;
                } else if (viewFn instanceof DataflowPortabilityPCollectionView.PortabilityViewFn) {
                    @SuppressWarnings("unchecked") ViewT viewT = (ViewT) getPortabilityMultimapForWindow(tag, window);
                    return viewT;
                }
                throw new IllegalArgumentException("Unknown type of view requested: " + view);
            });
        }
    } catch (ExecutionException e) {
        throw new IllegalStateException(String.format("Failed to materialize view %s for window %s.", view, window), e.getCause());
    }
}