Search in sources :

Example 1 with StateInternalsProxy

use of org.apache.beam.runners.apex.translation.utils.StateInternalsProxy in project beam by apache.

the class ApexParDoOperator method fireTimer.

@Override
public void fireTimer(Object key, Collection<TimerData> timerDataSet) {
    pushbackDoFnRunner.startBundle();
    @SuppressWarnings("unchecked") Coder<Object> keyCoder = (Coder) currentKeyStateInternals.getKeyCoder();
    ((StateInternalsProxy) currentKeyStateInternals).setKey(key);
    currentKeyTimerInternals.setContext(key, keyCoder, new Instant(this.currentInputWatermark), new Instant(this.currentOutputWatermark));
    for (TimerData timerData : timerDataSet) {
        StateNamespace namespace = timerData.getNamespace();
        checkArgument(namespace instanceof WindowNamespace);
        BoundedWindow window = ((WindowNamespace<?>) namespace).getWindow();
        pushbackDoFnRunner.onTimer(timerData.getTimerId(), window, timerData.getTimestamp(), timerData.getDomain());
    }
    pushbackDoFnRunner.finishBundle();
}
Also used : WindowedValueCoder(org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder) KeyedWorkItemCoder(org.apache.beam.runners.core.KeyedWorkItemCoder) ListCoder(org.apache.beam.sdk.coders.ListCoder) KvCoder(org.apache.beam.sdk.coders.KvCoder) Coder(org.apache.beam.sdk.coders.Coder) StringUtf8Coder(org.apache.beam.sdk.coders.StringUtf8Coder) VoidCoder(org.apache.beam.sdk.coders.VoidCoder) StateInternalsProxy(org.apache.beam.runners.apex.translation.utils.StateInternalsProxy) Instant(org.joda.time.Instant) WindowNamespace(org.apache.beam.runners.core.StateNamespaces.WindowNamespace) TimerData(org.apache.beam.runners.core.TimerInternals.TimerData) BoundedWindow(org.apache.beam.sdk.transforms.windowing.BoundedWindow) StateNamespace(org.apache.beam.runners.core.StateNamespace)

Example 2 with StateInternalsProxy

use of org.apache.beam.runners.apex.translation.utils.StateInternalsProxy in project beam by apache.

the class ApexParDoOperator method processElementInReadyWindows.

private Iterable<WindowedValue<InputT>> processElementInReadyWindows(WindowedValue<InputT> elem) {
    try {
        pushbackDoFnRunner.startBundle();
        if (currentKeyStateInternals != null) {
            InputT value = elem.getValue();
            final Object key;
            final Coder<Object> keyCoder;
            @SuppressWarnings({ "rawtypes", "unchecked" }) WindowedValueCoder<InputT> wvCoder = (WindowedValueCoder) inputCoder;
            if (value instanceof KeyedWorkItem) {
                key = ((KeyedWorkItem) value).key();
                @SuppressWarnings({ "rawtypes", "unchecked" }) KeyedWorkItemCoder<Object, ?> kwiCoder = (KeyedWorkItemCoder) wvCoder.getValueCoder();
                keyCoder = kwiCoder.getKeyCoder();
            } else {
                key = ((KV) value).getKey();
                @SuppressWarnings({ "rawtypes", "unchecked" }) KvCoder<Object, ?> kwiCoder = (KvCoder) wvCoder.getValueCoder();
                keyCoder = kwiCoder.getKeyCoder();
            }
            ((StateInternalsProxy) currentKeyStateInternals).setKey(key);
            currentKeyTimerInternals.setContext(key, keyCoder, new Instant(this.currentInputWatermark), new Instant(this.currentOutputWatermark));
        }
        Iterable<WindowedValue<InputT>> pushedBack = pushbackDoFnRunner.processElementInReadyWindows(elem);
        pushbackDoFnRunner.finishBundle();
        return pushedBack;
    } catch (UserCodeException ue) {
        if (ue.getCause() instanceof AssertionError) {
            ApexRunner.ASSERTION_ERROR.set((AssertionError) ue.getCause());
        }
        throw ue;
    }
}
Also used : StateInternalsProxy(org.apache.beam.runners.apex.translation.utils.StateInternalsProxy) KeyedWorkItemCoder(org.apache.beam.runners.core.KeyedWorkItemCoder) Instant(org.joda.time.Instant) KvCoder(org.apache.beam.sdk.coders.KvCoder) KeyedWorkItem(org.apache.beam.runners.core.KeyedWorkItem) WindowedValueCoder(org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder) UserCodeException(org.apache.beam.sdk.util.UserCodeException) WindowedValue(org.apache.beam.sdk.util.WindowedValue) OutputWindowedValue(org.apache.beam.runners.core.OutputWindowedValue)

Aggregations

StateInternalsProxy (org.apache.beam.runners.apex.translation.utils.StateInternalsProxy)2 KeyedWorkItemCoder (org.apache.beam.runners.core.KeyedWorkItemCoder)2 KvCoder (org.apache.beam.sdk.coders.KvCoder)2 WindowedValueCoder (org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder)2 Instant (org.joda.time.Instant)2 KeyedWorkItem (org.apache.beam.runners.core.KeyedWorkItem)1 OutputWindowedValue (org.apache.beam.runners.core.OutputWindowedValue)1 StateNamespace (org.apache.beam.runners.core.StateNamespace)1 WindowNamespace (org.apache.beam.runners.core.StateNamespaces.WindowNamespace)1 TimerData (org.apache.beam.runners.core.TimerInternals.TimerData)1 Coder (org.apache.beam.sdk.coders.Coder)1 ListCoder (org.apache.beam.sdk.coders.ListCoder)1 StringUtf8Coder (org.apache.beam.sdk.coders.StringUtf8Coder)1 VoidCoder (org.apache.beam.sdk.coders.VoidCoder)1 BoundedWindow (org.apache.beam.sdk.transforms.windowing.BoundedWindow)1 UserCodeException (org.apache.beam.sdk.util.UserCodeException)1 WindowedValue (org.apache.beam.sdk.util.WindowedValue)1