Search in sources :

Example 36 with Components

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.Components in project beam by apache.

the class QueryablePipelineTest method retainOnlyPrimitivesComposites.

@Test
public void retainOnlyPrimitivesComposites() {
    Pipeline p = Pipeline.create();
    p.apply(new org.apache.beam.sdk.transforms.PTransform<PBegin, PCollection<Long>>() {

        @Override
        public PCollection<Long> expand(PBegin input) {
            return input.apply(Impulse.create()).apply(Window.into(FixedWindows.of(Duration.standardMinutes(5L)))).apply(MapElements.into(TypeDescriptors.longs()).via(l -> 1L));
        }
    });
    Components originalComponents = PipelineTranslation.toProto(p).getComponents();
    Collection<String> primitiveComponents = QueryablePipeline.getPrimitiveTransformIds(originalComponents);
    // Read, Window.Assign, ParDo. This will need to be updated if the expansions change.
    assertThat(primitiveComponents, hasSize(3));
    for (String transformId : primitiveComponents) {
        assertThat(originalComponents.getTransformsMap(), hasKey(transformId));
    }
}
Also used : Components(org.apache.beam.model.pipeline.v1.RunnerApi.Components) PCollection(org.apache.beam.sdk.values.PCollection) PBegin(org.apache.beam.sdk.values.PBegin) Pipeline(org.apache.beam.sdk.Pipeline) Test(org.junit.Test)

Example 37 with Components

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.Components in project beam by apache.

the class QueryablePipelineTest method getEnvironmentWithEnvironment.

@Test
public void getEnvironmentWithEnvironment() {
    Pipeline p = Pipeline.create();
    PCollection<Long> longs = p.apply("Impulse", Impulse.create()).apply("ParDo", ParDo.of(new TestFn()));
    longs.apply(WithKeys.of("a")).apply("groupByKey", GroupByKey.create());
    Components components = PipelineTranslation.toProto(p).getComponents();
    QueryablePipeline qp = QueryablePipeline.forPrimitivesIn(components);
    PTransformNode environmentalTransform = PipelineNode.pTransform("ParDo-ParMultiDo-Test-", components.getTransformsOrThrow("ParDo-ParMultiDo-Test-"));
    PTransformNode nonEnvironmentalTransform = PipelineNode.pTransform("groupByKey", components.getTransformsOrThrow("groupByKey"));
    assertThat(qp.getEnvironment(environmentalTransform).isPresent(), is(true));
    assertThat(qp.getEnvironment(environmentalTransform).get().getUrn(), equalTo(Environments.JAVA_SDK_HARNESS_ENVIRONMENT.getUrn()));
    assertThat(qp.getEnvironment(environmentalTransform).get().getPayload(), equalTo(Environments.JAVA_SDK_HARNESS_ENVIRONMENT.getPayload()));
    assertThat(qp.getEnvironment(nonEnvironmentalTransform).isPresent(), is(false));
}
Also used : Components(org.apache.beam.model.pipeline.v1.RunnerApi.Components) PTransformNode(org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode) Pipeline(org.apache.beam.sdk.Pipeline) Test(org.junit.Test)

Example 38 with Components

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.Components in project beam by apache.

the class QueryablePipelineTest method getProducer.

@Test
public void getProducer() {
    Pipeline p = Pipeline.create();
    PCollection<Long> longs = p.apply("BoundedRead", Read.from(CountingSource.upTo(100L)));
    PCollectionList.of(longs).and(longs).and(longs).apply("flatten", Flatten.pCollections());
    Components components = PipelineTranslation.toProto(p).getComponents();
    QueryablePipeline qp = QueryablePipeline.forPrimitivesIn(components);
    String impulseOutputName = getOnlyElement(PipelineNode.pTransform("BoundedRead-Impulse", components.getTransformsOrThrow("BoundedRead-Impulse")).getTransform().getOutputsMap().values());
    PTransformNode impulseProducer = PipelineNode.pTransform("BoundedRead-Impulse", components.getTransformsOrThrow("BoundedRead-Impulse"));
    PCollectionNode impulseOutput = PipelineNode.pCollection(impulseOutputName, components.getPcollectionsOrThrow(impulseOutputName));
    String flattenOutputName = getOnlyElement(PipelineNode.pTransform("flatten", components.getTransformsOrThrow("flatten")).getTransform().getOutputsMap().values());
    PTransformNode flattenProducer = PipelineNode.pTransform("flatten", components.getTransformsOrThrow("flatten"));
    PCollectionNode flattenOutput = PipelineNode.pCollection(flattenOutputName, components.getPcollectionsOrThrow(flattenOutputName));
    assertThat(qp.getProducer(impulseOutput), equalTo(impulseProducer));
    assertThat(qp.getProducer(flattenOutput), equalTo(flattenProducer));
}
Also used : Components(org.apache.beam.model.pipeline.v1.RunnerApi.Components) PTransformNode(org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode) PCollectionNode(org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode) Pipeline(org.apache.beam.sdk.Pipeline) Test(org.junit.Test)

Example 39 with Components

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.Components in project beam by apache.

the class QueryablePipelineTest method fromComponentsWithMalformedComponents.

@Test
public void fromComponentsWithMalformedComponents() {
    Components components = Components.newBuilder().putTransforms("root", PTransform.newBuilder().setSpec(FunctionSpec.newBuilder().setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN).build()).putOutputs("output", "output.out").build()).build();
    thrown.expect(IllegalArgumentException.class);
    QueryablePipeline.forPrimitivesIn(components).getComponents();
}
Also used : Components(org.apache.beam.model.pipeline.v1.RunnerApi.Components) Test(org.junit.Test)

Example 40 with Components

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.Components in project beam by apache.

the class QueryablePipelineTest method retainOnlyPrimitivesWithOnlyPrimitivesUnchanged.

@Test
public void retainOnlyPrimitivesWithOnlyPrimitivesUnchanged() {
    Pipeline p = Pipeline.create();
    p.apply("Impulse", Impulse.create()).apply("multi-do", ParDo.of(new TestFn()).withOutputTags(new TupleTag<>(), TupleTagList.empty()));
    Components originalComponents = PipelineTranslation.toProto(p).getComponents();
    Collection<String> primitiveComponents = QueryablePipeline.getPrimitiveTransformIds(originalComponents);
    assertThat(primitiveComponents, equalTo(originalComponents.getTransformsMap().keySet()));
}
Also used : Components(org.apache.beam.model.pipeline.v1.RunnerApi.Components) TupleTag(org.apache.beam.sdk.values.TupleTag) Pipeline(org.apache.beam.sdk.Pipeline) Test(org.junit.Test)

Aggregations

Components (org.apache.beam.model.pipeline.v1.RunnerApi.Components)49 Test (org.junit.Test)37 PTransform (org.apache.beam.model.pipeline.v1.RunnerApi.PTransform)19 PTransformNode (org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode)18 RunnerApi (org.apache.beam.model.pipeline.v1.RunnerApi)12 PCollection (org.apache.beam.model.pipeline.v1.RunnerApi.PCollection)10 PCollectionNode (org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode)10 Map (java.util.Map)9 Pipeline (org.apache.beam.sdk.Pipeline)9 Environment (org.apache.beam.model.pipeline.v1.RunnerApi.Environment)7 Collection (java.util.Collection)6 Collectors (java.util.stream.Collectors)6 List (java.util.List)5 ExecutableStagePayload (org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload)5 FunctionSpec (org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec)5 DeduplicationResult (org.apache.beam.runners.core.construction.graph.OutputDeduplicator.DeduplicationResult)5 ArrayList (java.util.ArrayList)4 Pipeline (org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline)4 PCollection (org.apache.beam.sdk.values.PCollection)4 ImmutableList (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList)4