Search in sources :

Example 1 with LateDataDroppingDoFnRunner

use of org.apache.beam.runners.core.LateDataDroppingDoFnRunner 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;
        }
    }
}
Also used : TupleTag(org.apache.beam.sdk.values.TupleTag) WindowedValue(org.apache.beam.sdk.util.WindowedValue) KV(org.apache.beam.sdk.values.KV) DoFnRunner(org.apache.beam.runners.core.DoFnRunner) LateDataDroppingDoFnRunner(org.apache.beam.runners.core.LateDataDroppingDoFnRunner) KV(org.apache.beam.sdk.values.KV) KeyedWorkItem(org.apache.beam.runners.core.KeyedWorkItem) OutputManager(org.apache.beam.runners.core.DoFnRunners.OutputManager)

Aggregations

DoFnRunner (org.apache.beam.runners.core.DoFnRunner)1 OutputManager (org.apache.beam.runners.core.DoFnRunners.OutputManager)1 KeyedWorkItem (org.apache.beam.runners.core.KeyedWorkItem)1 LateDataDroppingDoFnRunner (org.apache.beam.runners.core.LateDataDroppingDoFnRunner)1 WindowedValue (org.apache.beam.sdk.util.WindowedValue)1 KV (org.apache.beam.sdk.values.KV)1 TupleTag (org.apache.beam.sdk.values.TupleTag)1