Examples with DirectPipelineResult org.apache.beam.runners.direct.DirectRunner.DirectPipelineResult used on opensource projects

Example 1 with DirectPipelineResult

use of org.apache.beam.runners.direct.DirectRunner.DirectPipelineResult in project beam by apache.

the class DirectRunnerTest method wordCountShouldSucceed.

@Test
public void wordCountShouldSucceed() throws Throwable {
    Pipeline p = getPipeline();
    PCollection<KV<String, Long>> counts = p.apply(Create.of("foo", "bar", "foo", "baz", "bar", "foo")).apply(MapElements.via(new SimpleFunction<String, String>() {

        @Override
        public String apply(String input) {
            return input;
        }
    })).apply(Count.perElement());
    PCollection<String> countStrs = counts.apply(MapElements.via(new SimpleFunction<KV<String, Long>, String>() {

        @Override
        public String apply(KV<String, Long> input) {
            return String.format("%s: %s", input.getKey(), input.getValue());
        }
    }));
    PAssert.that(countStrs).containsInAnyOrder("baz: 1", "bar: 2", "foo: 3");
    DirectPipelineResult result = (DirectPipelineResult) p.run();
    result.waitUntilFinish();
}

Example 2 with DirectPipelineResult

use of org.apache.beam.runners.direct.DirectRunner.DirectPipelineResult in project beam by apache.

the class DirectRunner method run.

@Override
public DirectPipelineResult run(Pipeline pipeline) {
    try {
        options = MAPPER.readValue(MAPPER.writeValueAsBytes(options), PipelineOptions.class).as(DirectOptions.class);
    } catch (IOException e) {
        throw new IllegalArgumentException("PipelineOptions specified failed to serialize to JSON.", e);
    }
    performRewrites(pipeline);
    MetricsEnvironment.setMetricsSupported(true);
    try {
        DirectGraphVisitor graphVisitor = new DirectGraphVisitor();
        pipeline.traverseTopologically(graphVisitor);
        @SuppressWarnings("rawtypes") KeyedPValueTrackingVisitor keyedPValueVisitor = KeyedPValueTrackingVisitor.create();
        pipeline.traverseTopologically(keyedPValueVisitor);
        DisplayDataValidator.validatePipeline(pipeline);
        DisplayDataValidator.validateOptions(options);
        ExecutorService metricsPool = Executors.newCachedThreadPool(new ThreadFactoryBuilder().setThreadFactory(MoreExecutors.platformThreadFactory()).setDaemon(// otherwise you say you want to leak, please don't!
        false).setNameFormat("direct-metrics-counter-committer").build());
        DirectGraph graph = graphVisitor.getGraph();
        EvaluationContext context = EvaluationContext.create(clockSupplier.get(), Enforcement.bundleFactoryFor(enabledEnforcements, graph), graph, keyedPValueVisitor.getKeyedPValues(), metricsPool);
        TransformEvaluatorRegistry registry = TransformEvaluatorRegistry.javaSdkNativeRegistry(context, options);
        PipelineExecutor executor = ExecutorServiceParallelExecutor.create(options.getTargetParallelism(), registry, Enforcement.defaultModelEnforcements(enabledEnforcements), context, metricsPool);
        executor.start(graph, RootProviderRegistry.javaNativeRegistry(context, options));
        DirectPipelineResult result = new DirectPipelineResult(executor, context);
        if (options.isBlockOnRun()) {
            try {
                result.waitUntilFinish();
            } catch (UserCodeException userException) {
                throw new PipelineExecutionException(userException.getCause());
            } catch (Throwable t) {
                if (t instanceof RuntimeException) {
                    throw (RuntimeException) t;
                }
                throw new RuntimeException(t);
            }
        }
        return result;
    } finally {
        MetricsEnvironment.setMetricsSupported(false);
    }
}

Example 3 with DirectPipelineResult

use of org.apache.beam.runners.direct.DirectRunner.DirectPipelineResult in project beam by apache.

the class DirectRunnerTest method reusePipelineSucceeds.

@Test
public void reusePipelineSucceeds() throws Throwable {
    Pipeline p = getPipeline();
    changed = new AtomicInteger(0);
    PCollection<KV<String, Long>> counts = p.apply(Create.of("foo", "bar", "foo", "baz", "bar", "foo")).apply(MapElements.via(new SimpleFunction<String, String>() {

        @Override
        public String apply(String input) {
            return input;
        }
    })).apply(Count.perElement());
    PCollection<String> countStrs = counts.apply(MapElements.via(new SimpleFunction<KV<String, Long>, String>() {

        @Override
        public String apply(KV<String, Long> input) {
            return String.format("%s: %s", input.getKey(), input.getValue());
        }
    }));
    counts.apply(ParDo.of(new DoFn<KV<String, Long>, Void>() {

        @ProcessElement
        public void updateChanged(ProcessContext c) {
            changed.getAndIncrement();
        }
    }));
    PAssert.that(countStrs).containsInAnyOrder("baz: 1", "bar: 2", "foo: 3");
    DirectPipelineResult result = (DirectPipelineResult) p.run();
    result.waitUntilFinish();
    DirectPipelineResult otherResult = (DirectPipelineResult) p.run();
    otherResult.waitUntilFinish();
    assertThat("Each element should have been processed twice", changed.get(), equalTo(6));
}