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

Example 31 with SdkComponents

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

the class AssignWindowsRunnerTest method multipleInputWindowsAsMapFnSucceeds.

@Test
public void multipleInputWindowsAsMapFnSucceeds() throws Exception {
    WindowFn<Object, BoundedWindow> windowFn = new WindowFn<Object, BoundedWindow>() {

        @Override
        public Collection<BoundedWindow> assignWindows(AssignContext c) {
            c.window();
            return ImmutableSet.of(GlobalWindow.INSTANCE, new IntervalWindow(new Instant(-500), Duration.standardMinutes(3)));
        }

        @Override
        public void mergeWindows(MergeContext c) {
            throw new UnsupportedOperationException();
        }

        @Override
        public WindowMappingFn<BoundedWindow> getDefaultWindowMappingFn() {
            throw new UnsupportedOperationException();
        }

        @Override
        public boolean isCompatible(WindowFn<?, ?> other) {
            throw new UnsupportedOperationException();
        }

        @Override
        public Coder<BoundedWindow> windowCoder() {
            throw new UnsupportedOperationException();
        }
    };
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("java"));
    RunnerApi.PCollection pCollection = RunnerApi.PCollection.newBuilder().setUniqueName("input").setCoderId("coder-id").build();
    RunnerApi.Coder coder = CoderTranslation.toProto(VarIntCoder.of()).getCoder();
    PTransformRunnerFactoryTestContext context = PTransformRunnerFactoryTestContext.builder("ptransform", PTransform.newBuilder().putInputs("in", "input").putOutputs("out", "output").setSpec(FunctionSpec.newBuilder().setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN).setPayload(WindowIntoPayload.newBuilder().setWindowFn(WindowingStrategyTranslation.toProto(windowFn, components)).build().toByteString())).build()).pCollections(Collections.singletonMap("input", pCollection)).coders(Collections.singletonMap("coder-id", coder)).build();
    Collection<WindowedValue<?>> outputs = new ArrayList<>();
    context.addPCollectionConsumer("output", outputs::add, VarIntCoder.of());
    MapFnRunners.forWindowedValueMapFnFactory(new AssignWindowsMapFnFactory<>()).createRunnerForPTransform(context);
    WindowedValue<Integer> value = WindowedValue.of(2, new Instant(-10L), ImmutableList.of(new IntervalWindow(new Instant(-22L), Duration.standardMinutes(5L)), new IntervalWindow(new Instant(-120000L), Duration.standardMinutes(3L))), PaneInfo.ON_TIME_AND_ONLY_FIRING);
    context.getPCollectionConsumer("input").accept(value);
    assertThat(outputs, containsInAnyOrder(WindowedValue.of(2, new Instant(-10L), ImmutableSet.of(GlobalWindow.INSTANCE, new IntervalWindow(new Instant(-500), Duration.standardMinutes(3))), PaneInfo.ON_TIME_AND_ONLY_FIRING), WindowedValue.of(2, new Instant(-10L), ImmutableSet.of(GlobalWindow.INSTANCE, new IntervalWindow(new Instant(-500), Duration.standardMinutes(3))), PaneInfo.ON_TIME_AND_ONLY_FIRING)));
}

Example 32 with SdkComponents

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

the class CombineRunnersTest method createPipeline.

@Before
public void createPipeline() throws Exception {
    // Create pipeline with an input pCollection, combine, and output pCollection.
    TestCombineFn combineFn = new TestCombineFn();
    Combine.PerKey<String, String, Integer> combine = Combine.perKey(combineFn);
    Pipeline p = Pipeline.create();
    PCollection<KV<String, String>> inputPCollection = p.apply(Create.of(KV.of("unused", "0")));
    inputPCollection.setCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
    PCollection<KV<String, Integer>> outputPCollection = inputPCollection.apply(TEST_COMBINE_ID, combine);
    outputPCollection.setCoder(KvCoder.of(StringUtf8Coder.of(), BigEndianIntegerCoder.of()));
    // Create FnApi protos needed for the runner.
    SdkComponents sdkComponents = SdkComponents.create(p.getOptions());
    pProto = PipelineTranslation.toProto(p, sdkComponents);
    inputPCollectionId = sdkComponents.registerPCollection(inputPCollection);
    outputPCollectionId = sdkComponents.registerPCollection(outputPCollection);
    pTransform = pProto.getComponents().getTransformsOrThrow(TEST_COMBINE_ID);
}

Example 33 with SdkComponents

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

the class ToStringFnRunnerTest method testPrimitiveToString.

@Test
public void testPrimitiveToString() throws Exception {
    String pTransformId = "pTransformId";
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("java"));
    RunnerApi.PTransform pTransform = RunnerApi.PTransform.newBuilder().build();
    ThrowingFunction<KV<String, Integer>, KV<String, String>> toStringFunction = ToStringFnRunner.createToStringFunctionForPTransform(pTransformId, pTransform);
    KV<String, Integer> input = KV.of("key", 12345);
    assertEquals(KV.of("key", "12345"), toStringFunction.apply(input));
}

Example 34 with SdkComponents

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

the class ToStringFnRunnerTest method testToStringOverride.

@Test
public void testToStringOverride() throws Exception {
    class ClassWithToStringOverride {

        @Override
        public String toString() {
            return "Some string";
        }
    }
    String pTransformId = "pTransformId";
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("java"));
    RunnerApi.PTransform pTransform = RunnerApi.PTransform.newBuilder().build();
    ThrowingFunction<KV<String, ClassWithToStringOverride>, KV<String, String>> toStringFunction = ToStringFnRunner.createToStringFunctionForPTransform(pTransformId, pTransform);
    KV<String, ClassWithToStringOverride> input = KV.of("key", new ClassWithToStringOverride());
    assertEquals(KV.of("key", "Some string"), toStringFunction.apply(input));
}

Example 35 with SdkComponents

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

the class PubSubWritePayloadTranslationTest method testTranslateSinkWithTopicOverridden.

@Test
public void testTranslateSinkWithTopicOverridden() throws Exception {
    ValueProvider<TopicPath> runtimeProvider = pipeline.newProvider(TOPIC);
    PubsubUnboundedSink pubsubUnboundedSinkSink = new PubsubUnboundedSink(null, runtimeProvider, TIMESTAMP_ATTRIBUTE, ID_ATTRIBUTE, 0, 0, 0, Duration.ZERO, null);
    PubsubSink pubsubSink = new PubsubSink(pubsubUnboundedSinkSink);
    PCollection<byte[]> input = pipeline.apply(Create.of(new byte[0]));
    PDone output = input.apply(pubsubSink);
    AppliedPTransform<?, ?, PubsubSink> appliedPTransform = AppliedPTransform.of("sink", PValues.expandInput(input), PValues.expandOutput(output), pubsubSink, ResourceHints.create(), pipeline);
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("java"));
    RunnerApi.FunctionSpec spec = sinkTranslator.translate(appliedPTransform, components);
    assertEquals(PTransformTranslation.PUBSUB_WRITE, spec.getUrn());
    PubSubWritePayload payload = PubSubWritePayload.parseFrom(spec.getPayload());
    assertEquals(((NestedValueProvider) runtimeProvider).propertyName(), payload.getTopicRuntimeOverridden());
    assertTrue(payload.getTopic().isEmpty());
    assertEquals(TIMESTAMP_ATTRIBUTE, payload.getTimestampAttribute());
    assertEquals(ID_ATTRIBUTE, payload.getIdAttribute());
}