Search in sources :

Example 1 with PDone

use of org.apache.beam.sdk.values.PDone in project beam by apache.

the class WriteWithShardingFactoryTest method withNoShardingSpecifiedReturnsNewTransform.

@Test
public void withNoShardingSpecifiedReturnsNewTransform() {
    ResourceId outputDirectory = LocalResources.fromString("/foo", true);
    FilenamePolicy policy = DefaultFilenamePolicy.constructUsingStandardParameters(StaticValueProvider.of(outputDirectory), DefaultFilenamePolicy.DEFAULT_SHARD_TEMPLATE, "");
    WriteFiles<Object> original = WriteFiles.to(new FileBasedSink<Object>(StaticValueProvider.of(outputDirectory), policy) {

        @Override
        public WriteOperation<Object> createWriteOperation() {
            throw new IllegalArgumentException("Should not be used");
        }
    });
    @SuppressWarnings("unchecked") PCollection<Object> objs = (PCollection) p.apply(Create.empty(VoidCoder.of()));
    AppliedPTransform<PCollection<Object>, PDone, WriteFiles<Object>> originalApplication = AppliedPTransform.of("write", objs.expand(), Collections.<TupleTag<?>, PValue>emptyMap(), original, p);
    assertThat(factory.getReplacementTransform(originalApplication).getTransform(), not(equalTo((Object) original)));
}
Also used : DefaultFilenamePolicy(org.apache.beam.sdk.io.DefaultFilenamePolicy) FilenamePolicy(org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy) PCollection(org.apache.beam.sdk.values.PCollection) ResourceId(org.apache.beam.sdk.io.fs.ResourceId) PDone(org.apache.beam.sdk.values.PDone) WriteFiles(org.apache.beam.sdk.io.WriteFiles) Test(org.junit.Test)

Example 2 with PDone

use of org.apache.beam.sdk.values.PDone in project beam by apache.

the class DirectGraphVisitorTest method getStepNamesContainsAllTransforms.

@Test
public void getStepNamesContainsAllTransforms() {
    PCollection<String> created = p.apply(Create.of("1", "2", "3"));
    PCollection<String> transformed = created.apply(ParDo.of(new DoFn<String, String>() {

        @ProcessElement
        public void processElement(DoFn<String, String>.ProcessContext<String, String> c) throws Exception {
            c.output(Integer.toString(c.element().length()));
        }
    }));
    PDone finished = transformed.apply(new PTransform<PInput, PDone>() {

        @Override
        public PDone expand(PInput input) {
            return PDone.in(input.getPipeline());
        }
    });
    p.traverseTopologically(visitor);
    DirectGraph graph = visitor.getGraph();
    assertThat(graph.getStepName(graph.getProducer(created)), equalTo("s0"));
    assertThat(graph.getStepName(graph.getProducer(transformed)), equalTo("s1"));
// finished doesn't have a producer, because it's not a PValue.
// TODO: Demonstrate that PCollectionList/Tuple and other composite PValues are either safe to
// use, or make them so.
}
Also used : DoFn(org.apache.beam.sdk.transforms.DoFn) PDone(org.apache.beam.sdk.values.PDone) PInput(org.apache.beam.sdk.values.PInput) Test(org.junit.Test)

Aggregations

PDone (org.apache.beam.sdk.values.PDone)2 Test (org.junit.Test)2 DefaultFilenamePolicy (org.apache.beam.sdk.io.DefaultFilenamePolicy)1 FilenamePolicy (org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy)1 WriteFiles (org.apache.beam.sdk.io.WriteFiles)1 ResourceId (org.apache.beam.sdk.io.fs.ResourceId)1 DoFn (org.apache.beam.sdk.transforms.DoFn)1 PCollection (org.apache.beam.sdk.values.PCollection)1 PInput (org.apache.beam.sdk.values.PInput)1