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

Example 21 with Components

use of org.apache.beam.model.pipeline.v1.RunnerApi.Components 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();
}

Example 22 with Components

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

the class PipelineTranslatorUtils method getWindowedValueCoder.

public static <T> WindowedValue.WindowedValueCoder<T> getWindowedValueCoder(String pCollectionId, RunnerApi.Components components) {
    RunnerApi.PCollection pCollection = components.getPcollectionsOrThrow(pCollectionId);
    PipelineNode.PCollectionNode pCollectionNode = PipelineNode.pCollection(pCollectionId, pCollection);
    WindowedValue.WindowedValueCoder<T> coder;
    try {
        coder = (WindowedValue.WindowedValueCoder) WireCoders.instantiateRunnerWireCoder(pCollectionNode, components);
    } catch (IOException e) {
        throw new RuntimeException(e);
    }
    return coder;
}

Example 23 with Components

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

the class LengthPrefixUnknownCoders method addForModelCoder.

private static String addForModelCoder(String coderId, RunnerApi.Components.Builder components, boolean replaceWithByteArrayCoder) {
    Coder coder = components.getCodersOrThrow(coderId);
    RunnerApi.Coder.Builder builder = coder.toBuilder().clearComponentCoderIds();
    for (String componentCoderId : coder.getComponentCoderIdsList()) {
        builder.addComponentCoderIds(addLengthPrefixedCoder(componentCoderId, components, replaceWithByteArrayCoder));
    }
    return addCoder(builder.build(), components, coderId + "-length_prefix");
}

Example 24 with Components

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

the class AssignWindowsRunnerTest method factoryCreatesFromJavaWindowFn.

@Test
public void factoryCreatesFromJavaWindowFn() throws Exception {
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("java"));
    PTransform windowPTransform = PTransform.newBuilder().putInputs("in", "input").putOutputs("out", "output").setSpec(FunctionSpec.newBuilder().setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN).setPayload(WindowIntoPayload.newBuilder().setWindowFn(WindowingStrategyTranslation.toProto(new TestWindowFn(), components)).build().toByteString()).build()).build();
    ThrowingFunction<WindowedValue<?>, WindowedValue<?>> fn = (ThrowingFunction) factory.forPTransform("transform", windowPTransform);
    assertThat(fn.apply(WindowedValue.of(22L, new Instant(5), new IntervalWindow(new Instant(0L), new Instant(20027L)), PaneInfo.ON_TIME_AND_ONLY_FIRING)), equalTo(WindowedValue.of(22L, new Instant(5), new TestWindowFn().assignWindow(new Instant(5)), PaneInfo.ON_TIME_AND_ONLY_FIRING)));
}

Example 25 with Components

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

the class AssignWindowsRunnerTest method factoryCreatesFromKnownWindowFn.

@Test
public void factoryCreatesFromKnownWindowFn() throws Exception {
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("java"));
    PTransform windowPTransform = PTransform.newBuilder().putInputs("in", "input").putOutputs("out", "output").setSpec(FunctionSpec.newBuilder().setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN).setPayload(WindowIntoPayload.newBuilder().setWindowFn(WindowingStrategyTranslation.toProto(Sessions.withGapDuration(Duration.standardMinutes(12L)), components)).build().toByteString()).build()).build();
    ThrowingFunction<WindowedValue<?>, WindowedValue<?>> fn = (ThrowingFunction) factory.forPTransform("transform", windowPTransform);
    WindowedValue<?> output = fn.apply(WindowedValue.of(22L, new Instant(5), new IntervalWindow(new Instant(0L), new Instant(20027L)), PaneInfo.ON_TIME_AND_ONLY_FIRING));
    assertThat(output, equalTo(WindowedValue.of(22L, new Instant(5), new IntervalWindow(new Instant(5L), Duration.standardMinutes(12L)), PaneInfo.ON_TIME_AND_ONLY_FIRING)));
}