Examples with Timer org.apache.beam.runners.core.construction.Timer used on opensource projects

Example 11 with Timer

use of org.apache.beam.runners.core.construction.Timer in project beam by apache.

the class CommonCoderTest method verifyDecodedValue.

private void verifyDecodedValue(CommonCoder coder, Object expectedValue, Object actualValue) {
    String s = coder.getUrn();
    if (s.equals(getUrn(StandardCoders.Enum.BYTES))) {
        assertThat(expectedValue, equalTo(actualValue));
    } else if (s.equals(getUrn(StandardCoders.Enum.BOOL))) {
        assertEquals(expectedValue, actualValue);
    } else if (s.equals(getUrn(StandardCoders.Enum.STRING_UTF8))) {
        assertEquals(expectedValue, actualValue);
    } else if (s.equals(getUrn(StandardCoders.Enum.KV))) {
        assertThat(actualValue, instanceOf(KV.class));
        verifyDecodedValue(coder.getComponents().get(0), ((KV) expectedValue).getKey(), ((KV) actualValue).getKey());
        verifyDecodedValue(coder.getComponents().get(0), ((KV) expectedValue).getValue(), ((KV) actualValue).getValue());
    } else if (s.equals(getUrn(StandardCoders.Enum.VARINT))) {
        assertEquals(expectedValue, actualValue);
    } else if (s.equals(getUrn(StandardCoders.Enum.INTERVAL_WINDOW))) {
        assertEquals(expectedValue, actualValue);
    } else if (s.equals(getUrn(StandardCoders.Enum.ITERABLE)) || s.equals(getUrn(StandardCoders.Enum.STATE_BACKED_ITERABLE))) {
        assertThat(actualValue, instanceOf(Iterable.class));
        CommonCoder componentCoder = coder.getComponents().get(0);
        Iterator<Object> expectedValueIterator = ((Iterable<Object>) expectedValue).iterator();
        for (Object value : (Iterable<Object>) actualValue) {
            verifyDecodedValue(componentCoder, expectedValueIterator.next(), value);
        }
        assertFalse(expectedValueIterator.hasNext());
    } else if (s.equals(getUrn(StandardCoders.Enum.TIMER))) {
        assertEquals((Timer) expectedValue, (Timer) actualValue);
    } else if (s.equals(getUrn(StandardCoders.Enum.GLOBAL_WINDOW))) {
        assertEquals(expectedValue, actualValue);
    } else if (s.equals(getUrn(StandardCoders.Enum.WINDOWED_VALUE))) {
        assertEquals(expectedValue, actualValue);
    } else if (s.equals(getUrn(StandardCoders.Enum.PARAM_WINDOWED_VALUE))) {
        assertEquals(expectedValue, actualValue);
    } else if (s.equals(getUrn(StandardCoders.Enum.DOUBLE))) {
        assertEquals(expectedValue, actualValue);
    } else if (s.equals(getUrn(StandardCoders.Enum.ROW))) {
        assertEquals(expectedValue, actualValue);
    } else if (s.equals(getUrn(StandardCoders.Enum.SHARDED_KEY))) {
        assertEquals(expectedValue, actualValue);
    } else if (s.equals(getUrn(StandardCoders.Enum.CUSTOM_WINDOW))) {
        assertEquals(expectedValue, actualValue);
    } else {
        throw new IllegalStateException("Unknown coder URN: " + coder.getUrn());
    }
}

Example 12 with Timer

use of org.apache.beam.runners.core.construction.Timer in project beam by apache.

the class SparkExecutableStageFunction method call.

@Override
public Iterator<RawUnionValue> call(Iterator<WindowedValue<InputT>> inputs) throws Exception {
    SparkPipelineOptions options = pipelineOptions.get().as(SparkPipelineOptions.class);
    // Register standard file systems.
    FileSystems.setDefaultPipelineOptions(options);
    // Otherwise, this may cause validation errors (e.g. ParDoTest)
    if (!inputs.hasNext()) {
        return Collections.emptyIterator();
    }
    try (ExecutableStageContext stageContext = contextFactory.get(jobInfo)) {
        ExecutableStage executableStage = ExecutableStage.fromPayload(stagePayload);
        try (StageBundleFactory stageBundleFactory = stageContext.getStageBundleFactory(executableStage)) {
            ConcurrentLinkedQueue<RawUnionValue> collector = new ConcurrentLinkedQueue<>();
            StateRequestHandler stateRequestHandler = getStateRequestHandler(executableStage, stageBundleFactory.getProcessBundleDescriptor());
            if (executableStage.getTimers().size() == 0) {
                ReceiverFactory receiverFactory = new ReceiverFactory(collector, outputMap);
                processElements(stateRequestHandler, receiverFactory, null, stageBundleFactory, inputs);
                return collector.iterator();
            }
            // Used with Batch, we know that all the data is available for this key. We can't use the
            // timer manager from the context because it doesn't exist. So we create one and advance
            // time to the end after processing all elements.
            final InMemoryTimerInternals timerInternals = new InMemoryTimerInternals();
            timerInternals.advanceProcessingTime(Instant.now());
            timerInternals.advanceSynchronizedProcessingTime(Instant.now());
            ReceiverFactory receiverFactory = new ReceiverFactory(collector, outputMap);
            TimerReceiverFactory timerReceiverFactory = new TimerReceiverFactory(stageBundleFactory, (Timer<?> timer, TimerInternals.TimerData timerData) -> {
                currentTimerKey = timer.getUserKey();
                if (timer.getClearBit()) {
                    timerInternals.deleteTimer(timerData);
                } else {
                    timerInternals.setTimer(timerData);
                }
            }, windowCoder);
            // Process inputs.
            processElements(stateRequestHandler, receiverFactory, timerReceiverFactory, stageBundleFactory, inputs);
            // Finish any pending windows by advancing the input watermark to infinity.
            timerInternals.advanceInputWatermark(BoundedWindow.TIMESTAMP_MAX_VALUE);
            // Finally, advance the processing time to infinity to fire any timers.
            timerInternals.advanceProcessingTime(BoundedWindow.TIMESTAMP_MAX_VALUE);
            timerInternals.advanceSynchronizedProcessingTime(BoundedWindow.TIMESTAMP_MAX_VALUE);
            // itself)
            while (timerInternals.hasPendingTimers()) {
                try (RemoteBundle bundle = stageBundleFactory.getBundle(receiverFactory, timerReceiverFactory, stateRequestHandler, getBundleProgressHandler())) {
                    PipelineTranslatorUtils.fireEligibleTimers(timerInternals, bundle.getTimerReceivers(), currentTimerKey);
                }
            }
            return collector.iterator();
        }
    }
}