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

Example 1 with SdkComponents

use of org.apache.beam.runners.core.construction.SdkComponents 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 2 with SdkComponents

use of org.apache.beam.runners.core.construction.SdkComponents 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)));
}

Example 3 with SdkComponents

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

the class WindowMergingFnRunnerTest method createMergeTransformForWindowFn.

private static <W extends BoundedWindow> RunnerApi.PTransform createMergeTransformForWindowFn(WindowFn<?, W> windowFn) throws Exception {
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("test"));
    RunnerApi.FunctionSpec functionSpec = RunnerApi.FunctionSpec.newBuilder().setUrn(WindowMergingFnRunner.URN).setPayload(WindowingStrategyTranslation.toProto(windowFn, components).toByteString()).build();
    return RunnerApi.PTransform.newBuilder().setSpec(functionSpec).build();
}

Example 4 with SdkComponents

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

the class PubSubReadPayloadTranslationTest method testTranslateSourceToFunctionSpec.

@Test
public void testTranslateSourceToFunctionSpec() throws Exception {
    PCollection<byte[]> output = pipeline.apply(readFromPubSub);
    AppliedPTransform<?, ?, Read.Unbounded<byte[]>> appliedPTransform = AppliedPTransform.of("ReadFromPubsub", PValues.expandInput(pipeline.begin()), PValues.expandOutput(output), readFromPubSub, ResourceHints.create(), pipeline);
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("java"));
    RunnerApi.FunctionSpec spec = sourceTranslator.translate((AppliedPTransform) appliedPTransform, components);
    assertEquals(PTransformTranslation.PUBSUB_READ, spec.getUrn());
    PubSubReadPayload result = PubSubReadPayload.parseFrom(spec.getPayload());
    assertEquals(pubsubReadPayload, result);
}

Example 5 with SdkComponents

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

the class DoFnFunction method initTransient.

/**
 * Method used to initialize the transient variables that were sent over as byte arrays or proto
 * buffers.
 */
private void initTransient() {
    if (isInitialized) {
        return;
    }
    try {
        SdkComponents components = SdkComponents.create();
        pipelineOptions = new SerializablePipelineOptions(serializedOptions).get();
        DoFnWithExecutionInformation doFnWithExecutionInformation = (DoFnWithExecutionInformation) SerializableUtils.deserializeFromByteArray(doFnwithExBytes, "Custom Coder Bytes");
        this.doFn = (DoFn<InputT, OutputT>) doFnWithExecutionInformation.getDoFn();
        this.mainOutput = (TupleTag<OutputT>) doFnWithExecutionInformation.getMainOutputTag();
        this.sideInputMapping = doFnWithExecutionInformation.getSideInputMapping();
        this.doFnSchemaInformation = doFnWithExecutionInformation.getSchemaInformation();
        inputCoder = (Coder<InputT>) SerializableUtils.deserializeFromByteArray(coderBytes, "Custom Coder Bytes");
        windowStrategyProto = RunnerApi.MessageWithComponents.parseFrom(windowBytes);
        windowingStrategy = (WindowingStrategy<?, ?>) WindowingStrategyTranslation.fromProto(windowStrategyProto.getWindowingStrategy(), RehydratedComponents.forComponents(components.toComponents()));
        sideInputs = new HashMap<>();
        for (Map.Entry<String, byte[]> entry : sideInputBytes.entrySet()) {
            windowStrategyProto = RunnerApi.MessageWithComponents.parseFrom(entry.getValue());
            sideInputs.put(new TupleTag<>(entry.getKey()), WindowingStrategyTranslation.fromProto(windowStrategyProto.getWindowingStrategy(), RehydratedComponents.forComponents(components.toComponents())));
        }
    } catch (InvalidProtocolBufferException e) {
        LOG.info(e.getMessage());
    }
    outputCoders = new HashMap<>();
    for (Map.Entry<String, byte[]> entry : outputCodersBytes.entrySet()) {
        outputCoders.put(new TupleTag<>(entry.getKey()), (Coder<?>) SerializableUtils.deserializeFromByteArray(entry.getValue(), "Custom Coder Bytes"));
    }
    sideOutputs = new ArrayList<>();
    for (String sideOutput : serializedSideOutputs) {
        sideOutputs.add(new TupleTag<>(sideOutput));
    }
    outputMap = new HashMap<>();
    for (Map.Entry<String, Integer> entry : serializedOutputMap.entrySet()) {
        outputMap.put(new TupleTag<>(entry.getKey()), entry.getValue());
    }
    outputManager = new DoFnOutputManager(this.outputMap);
    this.isInitialized = true;
}