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

Example 76 with TupleTag

use of org.apache.beam.sdk.values.TupleTag in project beam by apache.

the class ReplacementOutputsTest method singletonSucceeds.

@Test
public void singletonSucceeds() {
    Map<PCollection<?>, ReplacementOutput> replacements = ReplacementOutputs.singleton(PValues.expandValue(ints), replacementInts);
    assertThat(replacements, Matchers.hasKey(replacementInts));
    ReplacementOutput replacement = replacements.get(replacementInts);
    Map.Entry<TupleTag<?>, PValue> taggedInts = Iterables.getOnlyElement(ints.expand().entrySet());
    assertThat(replacement.getOriginal().getTag(), equalTo(taggedInts.getKey()));
    assertThat(replacement.getOriginal().getValue(), equalTo(taggedInts.getValue()));
    assertThat(replacement.getReplacement().getValue(), equalTo(replacementInts));
}

Example 77 with TupleTag

use of org.apache.beam.sdk.values.TupleTag 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());
    }
}

Example 78 with TupleTag

use of org.apache.beam.sdk.values.TupleTag in project beam by apache.

the class SimpleParDoFn method reallyStartBundle.

private void reallyStartBundle() throws Exception {
    checkState(fnRunner == null, "bundle already started (or not properly finished)");
    OutputManager outputManager = new OutputManager() {

        final Map<TupleTag<?>, OutputReceiver> undeclaredOutputs = new HashMap<>();

        @Nullable
        private Receiver getReceiverOrNull(TupleTag<?> tag) {
            Integer receiverIndex = outputTupleTagsToReceiverIndices.get(tag);
            if (receiverIndex != null) {
                return receivers[receiverIndex];
            } else {
                return undeclaredOutputs.get(tag);
            }
        }

        @Override
        public <T> void output(TupleTag<T> tag, WindowedValue<T> output) {
            outputsPerElementTracker.onOutput();
            Receiver receiver = getReceiverOrNull(tag);
            if (receiver == null) {
                // A new undeclared output.
                // TODO: plumb through the operationName, so that we can
                // name implicit outputs after it.
                String outputName = "implicit-" + tag.getId();
                // TODO: plumb through the counter prefix, so we can
                // make it available to the OutputReceiver class in case
                // it wants to use it in naming output counterFactory.  (It
                // doesn't today.)
                OutputReceiver undeclaredReceiver = new OutputReceiver();
                ElementCounter outputCounter = new DataflowOutputCounter(outputName, counterFactory, stepContext.getNameContext());
                undeclaredReceiver.addOutputCounter(outputCounter);
                undeclaredOutputs.put(tag, undeclaredReceiver);
                receiver = undeclaredReceiver;
            }
            try {
                receiver.process(output);
            } catch (RuntimeException | Error e) {
                // via a chain of DoFn's.
                throw e;
            } catch (Exception e) {
                // with other Receivers.
                throw new RuntimeException(e);
            }
        }
    };
    fnInfo = (DoFnInfo) doFnInstanceManager.get();
    fnSignature = DoFnSignatures.getSignature(fnInfo.getDoFn().getClass());
    fnRunner = runnerFactory.createRunner(fnInfo.getDoFn(), options, mainOutputTag, sideOutputTags, fnInfo.getSideInputViews(), sideInputReader, fnInfo.getInputCoder(), fnInfo.getOutputCoders(), fnInfo.getWindowingStrategy(), stepContext, userStepContext, outputManager, doFnSchemaInformation, sideInputMapping);
    fnRunner.startBundle();
}

Example 79 with TupleTag

use of org.apache.beam.sdk.values.TupleTag in project beam by apache.

the class TaskTest method sideOutput.

@Test
public void sideOutput() {
    PCollection<Integer> numbers = testPipeline.apply(Create.of(10, 50, 120, 20, 200, 0));
    TupleTag<Integer> numBelow100Tag = new TupleTag<Integer>() {
    };
    TupleTag<Integer> numAbove100Tag = new TupleTag<Integer>() {
    };
    PCollectionTuple resultsTuple = Task.applyTransform(numbers, numBelow100Tag, numAbove100Tag);
    PAssert.that(resultsTuple.get(numBelow100Tag)).containsInAnyOrder(0, 10, 20, 50);
    PAssert.that(resultsTuple.get(numAbove100Tag)).containsInAnyOrder(120, 200);
    testPipeline.run().waitUntilFinish();
}

Example 80 with TupleTag

use of org.apache.beam.sdk.values.TupleTag in project beam by apache.

the class Task method main.

public static void main(String[] args) {
    PipelineOptions options = PipelineOptionsFactory.fromArgs(args).create();
    Pipeline pipeline = Pipeline.create(options);
    PCollection<Integer> numbers = pipeline.apply(Create.of(10, 50, 120, 20, 200, 0));
    TupleTag<Integer> numBelow100Tag = new TupleTag<Integer>() {
    };
    TupleTag<Integer> numAbove100Tag = new TupleTag<Integer>() {
    };
    PCollectionTuple outputTuple = applyTransform(numbers, numBelow100Tag, numAbove100Tag);
    outputTuple.get(numBelow100Tag).apply(Log.ofElements("Number <= 100: "));
    outputTuple.get(numAbove100Tag).apply(Log.ofElements("Number > 100: "));
    pipeline.run();
}