DirectRunner extends PipelineRunner<DirectPipelineResult> { @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); } pipeline.replaceAll(defaultTransformOverrides()); 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); SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReadsIfNecessary(pipeline); ExecutorService metricsPool = Executors.newCachedThreadPool( new ThreadFactoryBuilder() .setThreadFactory(MoreExecutors.platformThreadFactory()) .setDaemon(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); } } private  DirectRunner(DirectOptions options); static DirectRunner fromOptions(PipelineOptions options); @Override DirectPipelineResult run(Pipeline pipeline);  }
@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(); }
@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)); }
@Test public void splitsInputs() { Pipeline p = getPipeline(); PCollection<Long> longs = p.apply(Read.from(MustSplitSource.of(CountingSource.upTo(3)))); PAssert.that(longs).containsInAnyOrder(0L, 1L, 2L); p.run(); }
@Test public void cancelShouldStopPipeline() throws Exception { PipelineOptions opts = TestPipeline.testingPipelineOptions(); opts.as(DirectOptions.class).setBlockOnRun(false); opts.setRunner(DirectRunner.class); final Pipeline p = Pipeline.create(opts); p.apply(GenerateSequence.from(0).withRate(1L, Duration.standardSeconds(1))); final BlockingQueue<PipelineResult> resultExchange = new ArrayBlockingQueue<>(1); Runnable cancelRunnable = () -> { try { resultExchange.take().cancel(); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new IllegalStateException(e); } catch (IOException e) { throw new IllegalStateException(e); } }; Callable<PipelineResult> runPipelineRunnable = () -> { PipelineResult res = p.run(); try { resultExchange.put(res); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new IllegalStateException(e); } return res; }; ExecutorService executor = Executors.newCachedThreadPool(); Future<?> cancelResult = executor.submit(cancelRunnable); Future<PipelineResult> result = executor.submit(runPipelineRunnable); cancelResult.get(); result.get().waitUntilFinish(); }
@Test public void testWaitUntilFinishTimeout() throws Exception { DirectOptions options = PipelineOptionsFactory.as(DirectOptions.class); options.setBlockOnRun(false); options.setRunner(DirectRunner.class); Pipeline p = Pipeline.create(options); p.apply(Create.of(1L)) .apply( ParDo.of( new DoFn<Long, Long>() { @ProcessElement public void hang(ProcessContext context) throws InterruptedException { Thread.sleep(Long.MAX_VALUE); } })); PipelineResult result = p.run(); assertThat(result.getState(), is(State.RUNNING)); assertEquals(null, result.waitUntilFinish(Duration.millis(1L))); assertEquals(null, result.waitUntilFinish(Duration.millis(1L))); }
@Test public void tearsDownFnsBeforeFinishing() { TEARDOWN_CALL.set(-1); final Pipeline pipeline = getPipeline(); pipeline .apply(Create.of("a")) .apply( ParDo.of( new DoFn<String, String>() { @ProcessElement public void onElement(final ProcessContext ctx) { } @Teardown public void teardown() { try { Thread.sleep(1000); } catch (final InterruptedException e) { throw new AssertionError(e); } TEARDOWN_CALL.set(System.nanoTime()); } })); final PipelineResult pipelineResult = pipeline.run(); pipelineResult.waitUntilFinish(); final long doneTs = System.nanoTime(); final long tearDownTs = TEARDOWN_CALL.get(); assertThat(tearDownTs, greaterThan(0L)); assertThat(doneTs, greaterThan(tearDownTs)); }
@Test public void transformDisplayDataExceptionShouldFail() { DoFn<Integer, Integer> brokenDoFn = new DoFn<Integer, Integer>() { @ProcessElement public void processElement(ProcessContext c) throws Exception {} @Override public void populateDisplayData(DisplayData.Builder builder) { throw new RuntimeException("oh noes!"); } }; Pipeline p = getPipeline(); p.apply(Create.of(1, 2, 3)).apply(ParDo.of(brokenDoFn)); thrown.expectMessage(brokenDoFn.getClass().getName()); thrown.expectCause(ThrowableMessageMatcher.hasMessage(is("oh noes!"))); p.run(); }
@Test public void testMutatingOutputThenOutputDoFnError() throws Exception { Pipeline pipeline = getPipeline(); pipeline .apply(Create.of(42)) .apply( ParDo.of( new DoFn<Integer, List<Integer>>() { @ProcessElement public void processElement(ProcessContext c) { List<Integer> outputList = Arrays.asList(1, 2, 3, 4); c.output(outputList); outputList.set(0, 37); c.output(outputList); } })); thrown.expect(IllegalMutationException.class); thrown.expectMessage("output"); thrown.expectMessage("must not be mutated"); pipeline.run(); }
@Test public void testMutatingOutputWithEnforcementDisabledSucceeds() throws Exception { PipelineOptions options = PipelineOptionsFactory.create(); options.setRunner(DirectRunner.class); options.as(DirectOptions.class).setEnforceImmutability(false); Pipeline pipeline = Pipeline.create(options); pipeline .apply(Create.of(42)) .apply( ParDo.of( new DoFn<Integer, List<Integer>>() { @ProcessElement public void processElement(ProcessContext c) { List<Integer> outputList = Arrays.asList(1, 2, 3, 4); c.output(outputList); outputList.set(0, 37); c.output(outputList); } })); pipeline.run(); }
@Test public void testMutatingOutputThenTerminateDoFnError() throws Exception { Pipeline pipeline = getPipeline(); pipeline .apply(Create.of(42)) .apply( ParDo.of( new DoFn<Integer, List<Integer>>() { @ProcessElement public void processElement(ProcessContext c) { List<Integer> outputList = Arrays.asList(1, 2, 3, 4); c.output(outputList); outputList.set(0, 37); } })); thrown.expect(IllegalMutationException.class); thrown.expectMessage("output"); thrown.expectMessage("must not be mutated"); pipeline.run(); }
@Test public void testMutatingOutputCoderDoFnError() throws Exception { Pipeline pipeline = getPipeline(); pipeline .apply(Create.of(42)) .apply( ParDo.of( new DoFn<Integer, byte[]>() { @ProcessElement public void processElement(ProcessContext c) { byte[] outputArray = new byte[] {0x1, 0x2, 0x3}; c.output(outputArray); outputArray[0] = 0xa; c.output(outputArray); } })); thrown.expect(IllegalMutationException.class); thrown.expectMessage("output"); thrown.expectMessage("must not be mutated"); pipeline.run(); }
@Test public void testMutatingInputDoFnError() throws Exception { Pipeline pipeline = getPipeline(); pipeline .apply( Create.of(Arrays.asList(1, 2, 3), Arrays.asList(4, 5, 6)) .withCoder(ListCoder.of(VarIntCoder.of()))) .apply( ParDo.of( new DoFn<List<Integer>, Integer>() { @ProcessElement public void processElement(ProcessContext c) { List<Integer> inputList = c.element(); inputList.set(0, 37); c.output(12); } })); thrown.expect(IllegalMutationException.class); thrown.expectMessage("Input"); thrown.expectMessage("must not be mutated"); pipeline.run(); }
@Test public void testMutatingInputCoderDoFnError() throws Exception { Pipeline pipeline = getPipeline(); pipeline .apply(Create.of(new byte[] {0x1, 0x2, 0x3}, new byte[] {0x4, 0x5, 0x6})) .apply( ParDo.of( new DoFn<byte[], Integer>() { @ProcessElement public void processElement(ProcessContext c) { byte[] inputArray = c.element(); inputArray[0] = 0xa; c.output(13); } })); thrown.expect(IllegalMutationException.class); thrown.expectMessage("Input"); thrown.expectMessage("must not be mutated"); pipeline.run(); }
@Test public void testUnencodableOutputElement() throws Exception { Pipeline p = getPipeline(); PCollection<Long> pcollection = p.apply(Create.of((Void) null)) .apply( ParDo.of( new DoFn<Void, Long>() { @ProcessElement public void processElement(ProcessContext c) { c.output(null); } })) .setCoder(VarLongCoder.of()); pcollection.apply( ParDo.of( new DoFn<Long, Long>() { @ProcessElement public void unreachable(ProcessContext c) { fail("Pipeline should fail to encode a null Long in VarLongCoder"); } })); thrown.expectCause(isA(CoderException.class)); thrown.expectMessage("cannot encode a null Long"); p.run(); }
@Test public void testUnencodableOutputFromBoundedRead() throws Exception { Pipeline p = getPipeline(); p.apply(GenerateSequence.from(0).to(10)).setCoder(new LongNoDecodeCoder()); thrown.expectCause(isA(CoderException.class)); thrown.expectMessage("Cannot decode a long"); p.run(); }
@Test public void testUnencodableOutputFromUnboundedRead() { Pipeline p = getPipeline(); p.apply(GenerateSequence.from(0)).setCoder(new LongNoDecodeCoder()); thrown.expectMessage("Cannot decode a long"); p.run(); }
@Test public void testFromOptionsIfIgnoredFieldsGettingDropped() { TestSerializationOfOptions options = PipelineOptionsFactory.fromArgs( "--foo=testValue", "--ignoredField=overridden", "--runner=DirectRunner") .as(TestSerializationOfOptions.class); assertEquals("testValue", options.getFoo()); assertEquals("overridden", options.getIgnoredField()); Pipeline p = Pipeline.create(options); PCollection<Integer> pc = p.apply(Create.of("1")) .apply( ParDo.of( new DoFn<String, Integer>() { @ProcessElement public void processElement(ProcessContext c) { TestSerializationOfOptions options = c.getPipelineOptions().as(TestSerializationOfOptions.class); assertEquals("testValue", options.getFoo()); assertEquals("not overridden", options.getIgnoredField()); c.output(Integer.parseInt(c.element())); } })); PAssert.that(pc).containsInAnyOrder(1); p.run(); }
@Test(timeout = 10000) public void testTwoPOutputsInPipelineWithCascade() throws InterruptedException { StaticQueue<Integer> start = StaticQueue.of("start", VarIntCoder.of()); StaticQueue<Integer> messages = StaticQueue.of("messages", VarIntCoder.of()); Pipeline pipeline = getPipeline(false); pipeline.begin().apply("outputStartSignal", outputStartTo(start)); PCollection<Integer> result = pipeline .apply("processMessages", messages.read()) .apply( Window.<Integer>into(new GlobalWindows()) .triggering(AfterWatermark.pastEndOfWindow()) .discardingFiredPanes() .withAllowedLateness(Duration.ZERO)) .apply(Sum.integersGlobally()); PAssert.that(result).containsInAnyOrder(6); PipelineResult run = pipeline.run(); while (start.take() == null) {} messages.add(1).add(2).add(3).terminate(); run.waitUntilFinish(); }