Examples with OutputManager org.apache.beam.runners.core.DoFnRunners.OutputManager used on opensource projects

Example 6 with OutputManager

use of org.apache.beam.runners.core.DoFnRunners.OutputManager in project beam by apache.

the class ParDoEvaluator method create.

public static <InputT, OutputT> ParDoEvaluator<InputT> create(EvaluationContext evaluationContext, PipelineOptions options, DirectStepContext stepContext, AppliedPTransform<?, ?, ?> application, Coder<InputT> inputCoder, WindowingStrategy<?, ? extends BoundedWindow> windowingStrategy, DoFn<InputT, OutputT> fn, StructuralKey<?> key, List<PCollectionView<?>> sideInputs, TupleTag<OutputT> mainOutputTag, List<TupleTag<?>> additionalOutputTags, Map<TupleTag<?>, PCollection<?>> outputs, DoFnSchemaInformation doFnSchemaInformation, Map<String, PCollectionView<?>> sideInputMapping, DoFnRunnerFactory<InputT, OutputT> runnerFactory) {
    BundleOutputManager outputManager = createOutputManager(evaluationContext, key, outputs);
    ReadyCheckingSideInputReader sideInputReader = evaluationContext.createSideInputReader(sideInputs);
    Map<TupleTag<?>, Coder<?>> outputCoders = outputs.entrySet().stream().collect(Collectors.toMap(e -> e.getKey(), e -> e.getValue().getCoder()));
    PushbackSideInputDoFnRunner<InputT, OutputT> runner = runnerFactory.createRunner(options, fn, sideInputs, sideInputReader, outputManager, mainOutputTag, additionalOutputTags, stepContext, inputCoder, outputCoders, windowingStrategy, doFnSchemaInformation, sideInputMapping);
    return create(runner, stepContext, application, outputManager);
}

Example 7 with OutputManager

use of org.apache.beam.runners.core.DoFnRunners.OutputManager in project beam by apache.

the class GroupAlsoByWindowsParDoFn method createRunner.

/**
 * Composes and returns a {@link DoFnRunner} based on the parameters.
 *
 * <p>A {@code SimpleOldDoFnRunner} executes the {@link GroupAlsoByWindowFn}.
 *
 * <p>A {@link LateDataDroppingDoFnRunner} handles late data dropping for a {@link
 * StreamingGroupAlsoByWindowViaWindowSetFn}.
 *
 * <p>A {@link StreamingSideInputDoFnRunner} handles streaming side inputs.
 *
 * <p>A {@link StreamingKeyedWorkItemSideInputDoFnRunner} handles streaming side inputs for a
 * {@link StreamingGroupAlsoByWindowViaWindowSetFn}.
 */
private DoFnRunner<InputT, KV<K, Iterable<V>>> createRunner() {
    OutputManager outputManager = new OutputManager() {

        @Override
        public <T> void output(TupleTag<T> tag, WindowedValue<T> output) {
            checkState(tag.equals(mainOutputTag), "Must only output to main output tag (%s), but was %s", tag, mainOutputTag);
            try {
                receiver.process(output);
            } catch (Throwable t) {
                throw new RuntimeException(t);
            }
        }
    };
    boolean hasStreamingSideInput = options.as(StreamingOptions.class).isStreaming() && !sideInputReader.isEmpty();
    DoFnRunner<InputT, KV<K, Iterable<V>>> basicRunner = new GroupAlsoByWindowFnRunner<>(options, doFn, sideInputReader, outputManager, mainOutputTag, stepContext);
    if (doFn instanceof StreamingGroupAlsoByWindowViaWindowSetFn) {
        DoFnRunner<KeyedWorkItem<K, V>, KV<K, Iterable<V>>> streamingGABWRunner = (DoFnRunner<KeyedWorkItem<K, V>, KV<K, Iterable<V>>>) basicRunner;
        if (hasStreamingSideInput) {
            @SuppressWarnings("unchecked") WindmillKeyedWorkItem.FakeKeyedWorkItemCoder<K, V> keyedWorkItemCoder = (WindmillKeyedWorkItem.FakeKeyedWorkItemCoder<K, V>) inputCoder;
            StreamingSideInputFetcher<V, W> sideInputFetcher = new StreamingSideInputFetcher<>(sideInputViews, keyedWorkItemCoder.getElementCoder(), windowingStrategy, (StreamingModeExecutionContext.StreamingModeStepContext) stepContext);
            streamingGABWRunner = new StreamingKeyedWorkItemSideInputDoFnRunner<>(streamingGABWRunner, keyedWorkItemCoder.getKeyCoder(), sideInputFetcher, stepContext);
        }
        return (DoFnRunner<InputT, KV<K, Iterable<V>>>) DoFnRunners.<K, V, Iterable<V>, W>lateDataDroppingRunner(streamingGABWRunner, stepContext.timerInternals(), windowingStrategy);
    } else {
        if (hasStreamingSideInput) {
            return new StreamingSideInputDoFnRunner<>(basicRunner, new StreamingSideInputFetcher<>(sideInputViews, inputCoder, windowingStrategy, (StreamingModeExecutionContext.StreamingModeStepContext) stepContext));
        } else {
            return basicRunner;
        }
    }
}