Examples with TimerFamilySpec org.apache.beam.model.pipeline.v1.RunnerApi.TimerFamilySpec used on opensource projects

Example 1 with TimerFamilySpec

use of org.apache.beam.model.pipeline.v1.RunnerApi.TimerFamilySpec in project beam by apache.

the class ProcessBundleDescriptors method forTimerSpecs.

private static Map<String, Map<String, TimerSpec>> forTimerSpecs(ExecutableStage stage, Components.Builder components) throws IOException {
    ImmutableTable.Builder<String, String, TimerSpec> idsToSpec = ImmutableTable.builder();
    for (TimerReference timerReference : stage.getTimers()) {
        RunnerApi.ParDoPayload payload = RunnerApi.ParDoPayload.parseFrom(timerReference.transform().getTransform().getSpec().getPayload());
        RunnerApi.TimerFamilySpec timerFamilySpec = payload.getTimerFamilySpecsOrThrow(timerReference.localName());
        org.apache.beam.sdk.state.TimerSpec spec;
        switch(timerFamilySpec.getTimeDomain()) {
            case EVENT_TIME:
                spec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
                break;
            case PROCESSING_TIME:
                spec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
                break;
            default:
                throw new IllegalArgumentException(String.format("Unknown or unsupported time domain %s", timerFamilySpec.getTimeDomain()));
        }
        for (WireCoderSetting wireCoderSetting : stage.getWireCoderSettings()) {
            if (wireCoderSetting.hasTimer() && wireCoderSetting.getTimer().getTransformId().equals(timerReference.transform().getId()) && wireCoderSetting.getTimer().getLocalName().equals(timerReference.localName())) {
                throw new UnsupportedOperationException("WireCoderSetting for timer is yet to be supported.");
            }
        }
        String originalTimerCoderId = timerFamilySpec.getTimerFamilyCoderId();
        String sdkCoderId = LengthPrefixUnknownCoders.addLengthPrefixedCoder(originalTimerCoderId, components, false);
        String runnerCoderId = LengthPrefixUnknownCoders.addLengthPrefixedCoder(originalTimerCoderId, components, true);
        Coder<?> timerCoder = RehydratedComponents.forComponents(components.build()).getCoder(runnerCoderId);
        checkArgument(timerCoder instanceof Timer.Coder, "Expected a timer coder but received %s.", timerCoder);
        RunnerApi.FunctionSpec.Builder updatedSpec = components.getTransformsOrThrow(timerReference.transform().getId()).toBuilder().getSpecBuilder();
        RunnerApi.ParDoPayload.Builder updatedPayload = RunnerApi.ParDoPayload.parseFrom(updatedSpec.getPayload()).toBuilder();
        updatedPayload.putTimerFamilySpecs(timerReference.localName(), updatedPayload.getTimerFamilySpecsOrThrow(timerReference.localName()).toBuilder().setTimerFamilyCoderId(sdkCoderId).build());
        updatedSpec.setPayload(updatedPayload.build().toByteString());
        components.putTransforms(timerReference.transform().getId(), // and not the original
        components.getTransformsOrThrow(timerReference.transform().getId()).toBuilder().setSpec(updatedSpec).build());
        idsToSpec.put(timerReference.transform().getId(), timerReference.localName(), TimerSpec.of(timerReference.transform().getId(), timerReference.localName(), spec, (Coder) timerCoder));
    }
    return idsToSpec.build().rowMap();
}