Examples with PTransformOverride org.apache.beam.sdk.runners.PTransformOverride used on opensource projects

Example 1 with PTransformOverride

use of org.apache.beam.sdk.runners.PTransformOverride in project beam by apache.

the class PipelineTest method testReplaceAllIncomplete.

/**
   * Tests that {@link Pipeline#replaceAll(List)} throws when one of the PTransformOverride still
   * matches.
   */
@Test
public void testReplaceAllIncomplete() {
    pipeline.enableAbandonedNodeEnforcement(false);
    pipeline.apply(GenerateSequence.from(0));
    // The order is such that the output of the second will match the first, which is not permitted
    thrown.expect(IllegalStateException.class);
    pipeline.replaceAll(ImmutableList.of(PTransformOverride.of(new PTransformMatcher() {

        @Override
        public boolean matches(AppliedPTransform<?, ?, ?> application) {
            return application.getTransform() instanceof Create.Values;
        }
    }, new CreateValuesToEmptyFlattenOverride()), PTransformOverride.of(new PTransformMatcher() {

        @Override
        public boolean matches(AppliedPTransform<?, ?, ?> application) {
            return application.getTransform() instanceof GenerateSequence;
        }
    }, new GenerateSequenceToCreateOverride())));
}

Example 2 with PTransformOverride

use of org.apache.beam.sdk.runners.PTransformOverride in project beam by apache.

the class DataflowRunner method getOverrides.

private List<PTransformOverride> getOverrides(boolean streaming) {
    ImmutableList.Builder<PTransformOverride> overridesBuilder = ImmutableList.builder();
    // Create is implemented in terms of a Read, so it must precede the override to Read in
    // streaming
    overridesBuilder.add(PTransformOverride.of(PTransformMatchers.flattenWithDuplicateInputs(), DeduplicatedFlattenFactory.create())).add(PTransformOverride.of(PTransformMatchers.emptyFlatten(), EmptyFlattenAsCreateFactory.instance()));
    if (streaming) {
        if (!hasExperiment(options, "enable_custom_pubsub_source")) {
            overridesBuilder.add(PTransformOverride.of(PTransformMatchers.classEqualTo(PubsubUnboundedSource.class), new ReflectiveRootOverrideFactory(StreamingPubsubIORead.class, this)));
        }
        if (!hasExperiment(options, "enable_custom_pubsub_sink")) {
            overridesBuilder.add(PTransformOverride.of(PTransformMatchers.classEqualTo(PubsubUnboundedSink.class), new StreamingPubsubIOWriteOverrideFactory(this)));
        }
        overridesBuilder.add(// must precede it
        PTransformOverride.of(PTransformMatchers.classEqualTo(Read.Bounded.class), new ReflectiveRootOverrideFactory(StreamingBoundedRead.class, this))).add(PTransformOverride.of(PTransformMatchers.classEqualTo(Read.Unbounded.class), new ReflectiveRootOverrideFactory(StreamingUnboundedRead.class, this))).add(PTransformOverride.of(PTransformMatchers.classEqualTo(View.CreatePCollectionView.class), new StreamingCreatePCollectionViewFactory()));
    } else {
        overridesBuilder.add(PTransformOverride.of(PTransformMatchers.stateOrTimerParDoMulti(), BatchStatefulParDoOverrides.multiOutputOverrideFactory())).add(PTransformOverride.of(PTransformMatchers.stateOrTimerParDoSingle(), BatchStatefulParDoOverrides.singleOutputOverrideFactory())).add(PTransformOverride.of(PTransformMatchers.classEqualTo(Combine.GloballyAsSingletonView.class), new BatchCombineGloballyAsSingletonViewFactory(this))).add(PTransformOverride.of(PTransformMatchers.classEqualTo(View.AsMap.class), new ReflectiveOneToOneOverrideFactory(BatchViewOverrides.BatchViewAsMap.class, this))).add(PTransformOverride.of(PTransformMatchers.classEqualTo(View.AsMultimap.class), new ReflectiveOneToOneOverrideFactory(BatchViewOverrides.BatchViewAsMultimap.class, this))).add(PTransformOverride.of(PTransformMatchers.classEqualTo(View.AsSingleton.class), new ReflectiveOneToOneOverrideFactory(BatchViewOverrides.BatchViewAsSingleton.class, this))).add(PTransformOverride.of(PTransformMatchers.classEqualTo(View.AsList.class), new ReflectiveOneToOneOverrideFactory(BatchViewOverrides.BatchViewAsList.class, this))).add(PTransformOverride.of(PTransformMatchers.classEqualTo(View.AsIterable.class), new ReflectiveOneToOneOverrideFactory(BatchViewOverrides.BatchViewAsIterable.class, this)));
    }
    overridesBuilder.add(PTransformOverride.of(PTransformMatchers.classEqualTo(Reshuffle.class), new ReshuffleOverrideFactory())).add(PTransformOverride.of(PTransformMatchers.classEqualTo(Combine.GroupedValues.class), new PrimitiveCombineGroupedValuesOverrideFactory())).add(PTransformOverride.of(PTransformMatchers.classEqualTo(ParDo.SingleOutput.class), new PrimitiveParDoSingleFactory()));
    return overridesBuilder.build();
}

Example 3 with PTransformOverride

use of org.apache.beam.sdk.runners.PTransformOverride in project beam by apache.

the class Pipeline method replace.

private void replace(final PTransformOverride override) {
    final Set<Node> matches = new HashSet<>();
    final Set<Node> freedNodes = new HashSet<>();
    traverseTopologically(new PipelineVisitor.Defaults() {

        @Override
        public CompositeBehavior enterCompositeTransform(Node node) {
            if (!node.isRootNode() && freedNodes.contains(node.getEnclosingNode())) {
                // This node will be freed because its parent will be freed.
                freedNodes.add(node);
                return CompositeBehavior.ENTER_TRANSFORM;
            }
            if (!node.isRootNode() && override.getMatcher().matches(node.toAppliedPTransform(getPipeline()))) {
                matches.add(node);
                // This node will be freed. When we visit any of its children, they will also be freed
                freedNodes.add(node);
            }
            return CompositeBehavior.ENTER_TRANSFORM;
        }

        @Override
        public void visitPrimitiveTransform(Node node) {
            if (freedNodes.contains(node.getEnclosingNode())) {
                freedNodes.add(node);
            } else if (override.getMatcher().matches(node.toAppliedPTransform(getPipeline()))) {
                matches.add(node);
                freedNodes.add(node);
            }
        }
    });
    for (Node freedNode : freedNodes) {
        usedFullNames.remove(freedNode.getFullName());
    }
    for (Node match : matches) {
        applyReplacement(match, override.getOverrideFactory());
    }
}

Example 4 with PTransformOverride

use of org.apache.beam.sdk.runners.PTransformOverride in project beam by apache.

the class FlinkPipelineExecutionEnvironmentTest method shouldUseStreamingTransformOverridesWithUnboundedSources.

@Test
public void shouldUseStreamingTransformOverridesWithUnboundedSources() {
    FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
    // no explicit streaming mode set
    options.setRunner(FlinkRunner.class);
    FlinkPipelineExecutionEnvironment flinkEnv = new FlinkPipelineExecutionEnvironment(options);
    Pipeline p = Mockito.spy(Pipeline.create(options));
    // Add unbounded source which will set the streaming mode to true
    p.apply(GenerateSequence.from(0));
    flinkEnv.translate(p);
    ArgumentCaptor<ImmutableList> captor = ArgumentCaptor.forClass(ImmutableList.class);
    Mockito.verify(p).replaceAll(captor.capture());
    ImmutableList<PTransformOverride> overridesList = captor.getValue();
    assertThat(overridesList, hasItem(PTransformOverride.of(PTransformMatchers.urnEqualTo(PTransformTranslation.CREATE_VIEW_TRANSFORM_URN), CreateStreamingFlinkView.Factory.INSTANCE)));
}

Example 5 with PTransformOverride

use of org.apache.beam.sdk.runners.PTransformOverride in project beam by apache.

the class FlinkPipelineExecutionEnvironmentTest method shouldUseTransformOverrides.

@Test
public void shouldUseTransformOverrides() {
    boolean[] testParameters = { true, false };
    for (boolean streaming : testParameters) {
        FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
        options.setStreaming(streaming);
        options.setRunner(FlinkRunner.class);
        FlinkPipelineExecutionEnvironment flinkEnv = new FlinkPipelineExecutionEnvironment(options);
        Pipeline p = Mockito.spy(Pipeline.create(options));
        flinkEnv.translate(p);
        ArgumentCaptor<ImmutableList> captor = ArgumentCaptor.forClass(ImmutableList.class);
        Mockito.verify(p).replaceAll(captor.capture());
        ImmutableList<PTransformOverride> overridesList = captor.getValue();
        assertThat(overridesList.isEmpty(), is(false));
        assertThat(overridesList.size(), is(FlinkTransformOverrides.getDefaultOverrides(options).size()));
    }
}