SparkRunnerDebugger extends PipelineRunner<SparkPipelineResult> { @Override public SparkPipelineResult run(Pipeline pipeline) { SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads(pipeline); JavaSparkContext jsc = new JavaSparkContext("local[1]", "Debug_Pipeline"); JavaStreamingContext jssc = new JavaStreamingContext(jsc, new org.apache.spark.streaming.Duration(1000)); SparkRunner.initAccumulators(options, jsc); TransformTranslator.Translator translator = new TransformTranslator.Translator(); SparkNativePipelineVisitor visitor; if (options.isStreaming() || (options instanceof TestSparkPipelineOptions && ((TestSparkPipelineOptions) options).isForceStreaming())) { SparkPipelineTranslator streamingTranslator = new StreamingTransformTranslator.Translator(translator); EvaluationContext ctxt = new EvaluationContext(jsc, pipeline, options, jssc); visitor = new SparkNativePipelineVisitor(streamingTranslator, ctxt); } else { EvaluationContext ctxt = new EvaluationContext(jsc, pipeline, options, jssc); visitor = new SparkNativePipelineVisitor(translator, ctxt); } pipeline.traverseTopologically(visitor); jsc.stop(); String debugString = visitor.getDebugString(); LOG.info("Translated Native Spark pipeline:\n" + debugString); return new DebugSparkPipelineResult(debugString); } private  SparkRunnerDebugger(SparkPipelineOptions options); static SparkRunnerDebugger fromOptions(PipelineOptions options); @Override SparkPipelineResult run(Pipeline pipeline);  }
@Test public void debugBatchPipeline() { PipelineOptions options = PipelineOptionsFactory.create().as(TestSparkPipelineOptions.class); options.setRunner(SparkRunnerDebugger.class); Pipeline pipeline = Pipeline.create(options); PCollection<String> lines = pipeline.apply(Create.of(Collections.<String>emptyList()).withCoder(StringUtf8Coder.of())); PCollection<KV<String, Long>> wordCounts = lines.apply(new WordCount.CountWords()); wordCounts.apply(GroupByKey.create()).apply(Combine.groupedValues(Sum.ofLongs())); PCollection<KV<String, Long>> wordCountsPlusOne = wordCounts.apply(MapElements.via(new PlusOne())); PCollectionList.of(wordCounts).and(wordCountsPlusOne).apply(Flatten.pCollections()); wordCounts .apply(MapElements.via(new WordCount.FormatAsTextFn())) .apply(TextIO.write().to("!!PLACEHOLDER-OUTPUT-DIR!!").withNumShards(3).withSuffix(".txt")); final String expectedPipeline = "sparkContext.<readFrom(org.apache.beam.sdk.transforms.Create$Values$CreateSource)>()\n" + "_.mapPartitions(" + "new org.apache.beam.runners.spark.examples.WordCount$ExtractWordsFn())\n" + "_.mapPartitions(new org.apache.beam.sdk.transforms.Contextful())\n" + "_.combineByKey(..., new org.apache.beam.sdk.transforms.Count$CountFn(), ...)\n" + "_.groupByKey()\n" + "_.map(new org.apache.beam.sdk.transforms.Sum$SumLongFn())\n" + "_.mapPartitions(new org.apache.beam.sdk.transforms.Contextful())\n" + "sparkContext.union(...)\n" + "_.mapPartitions(" + "new org.apache.beam.sdk.transforms.Contextful())\n" + "_.<org.apache.beam.sdk.io.TextIO$Write>"; SparkRunnerDebugger.DebugSparkPipelineResult result = (SparkRunnerDebugger.DebugSparkPipelineResult) pipeline.run(); assertThat( "Debug pipeline did not equal expected", result.getDebugString(), Matchers.equalTo(expectedPipeline)); }
@Test public void debugStreamingPipeline() { TestSparkPipelineOptions options = PipelineOptionsFactory.create().as(TestSparkPipelineOptions.class); options.setForceStreaming(true); options.setRunner(SparkRunnerDebugger.class); Pipeline pipeline = Pipeline.create(options); KafkaIO.Read<String, String> read = KafkaIO.<String, String>read() .withBootstrapServers("mykafka:9092") .withTopics(Collections.singletonList("my_input_topic")) .withKeyDeserializer(StringDeserializer.class) .withValueDeserializer(StringDeserializer.class); KafkaIO.Write<String, String> write = KafkaIO.<String, String>write() .withBootstrapServers("myotherkafka:9092") .withTopic("my_output_topic") .withKeySerializer(StringSerializer.class) .withValueSerializer(StringSerializer.class); KvCoder<String, String> stringKvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()); pipeline .apply(read.withoutMetadata()) .setCoder(stringKvCoder) .apply(Window.into(FixedWindows.of(Duration.standardSeconds(5)))) .apply(ParDo.of(new SparkRunnerDebuggerTest.FormatKVFn())) .apply(Distinct.create()) .apply(WithKeys.of(new SparkRunnerDebuggerTest.ArbitraryKeyFunction())) .apply(write); final String expectedPipeline = "KafkaUtils.createDirectStream(...)\n" + "_.map(new org.apache.beam.sdk.transforms.windowing.FixedWindows())\n" + "_.mapPartitions(new org.apache.beam.runners.spark." + "SparkRunnerDebuggerTest$FormatKVFn())\n" + "_.mapPartitions(new org.apache.beam.sdk.transforms.Contextful())\n" + "_.groupByKey()\n" + "_.map(new org.apache.beam.sdk.transforms.Combine$IterableCombineFn())\n" + "_.mapPartitions(new org.apache.beam.sdk.transforms.Distinct$3())\n" + "_.mapPartitions(new org.apache.beam.sdk.transforms.Contextful())\n" + "_.<org.apache.beam.sdk.io.kafka.AutoValue_KafkaIO_Write>"; SparkRunnerDebugger.DebugSparkPipelineResult result = (SparkRunnerDebugger.DebugSparkPipelineResult) pipeline.run(); assertThat( "Debug pipeline did not equal expected", result.getDebugString(), Matchers.equalTo(expectedPipeline)); }