Java SparkRunner类代码示例

import org.apache.beam.runners.spark.SparkRunner; //导入依赖的package包/类
/**
 * Reject state and timers {@link DoFn}.
 *
 * @param doFn the {@link DoFn} to possibly reject.
 */
public static void rejectStateAndTimers(DoFn<?, ?> doFn) {
  DoFnSignature signature = DoFnSignatures.getSignature(doFn.getClass());

  if (signature.stateDeclarations().size() > 0) {
    throw new UnsupportedOperationException(
        String.format(
            "Found %s annotations on %s, but %s cannot yet be used with state in the %s.",
            DoFn.StateId.class.getSimpleName(),
            doFn.getClass().getName(),
            DoFn.class.getSimpleName(),
            SparkRunner.class.getSimpleName()));
  }

  if (signature.timerDeclarations().size() > 0) {
    throw new UnsupportedOperationException(
        String.format(
            "Found %s annotations on %s, but %s cannot yet be used with timers in the %s.",
            DoFn.TimerId.class.getSimpleName(),
            doFn.getClass().getName(),
            DoFn.class.getSimpleName(),
            SparkRunner.class.getSimpleName()));
  }
}
 

import org.apache.beam.runners.spark.SparkRunner; //导入依赖的package包/类
@Test
public void testTrackSingle() {
  options.setRunner(SparkRunner.class);
  JavaSparkContext jsc = SparkContextFactory.getSparkContext(options);
  JavaStreamingContext jssc = new JavaStreamingContext(jsc,
      new org.apache.spark.streaming.Duration(options.getBatchIntervalMillis()));

  Pipeline p = Pipeline.create(options);

  CreateStream<Integer> emptyStream =
      CreateStream.of(
          VarIntCoder.of(),
          Duration.millis(options.getBatchIntervalMillis())).emptyBatch();

  p.apply(emptyStream).apply(ParDo.of(new PassthroughFn<>()));

  p.traverseTopologically(new StreamingSourceTracker(jssc, p, ParDo.MultiOutput.class,  0));
  assertThat(StreamingSourceTracker.numAssertions, equalTo(1));
}
 

import org.apache.beam.runners.spark.SparkRunner; //导入依赖的package包/类
public static void rejectSplittable(DoFn<?, ?> doFn) {
  DoFnSignature signature = DoFnSignatures.getSignature(doFn.getClass());

  if (signature.processElement().isSplittable()) {
    throw new UnsupportedOperationException(
        String.format(
            "%s does not support splittable DoFn: %s", SparkRunner.class.getSimpleName(), doFn));
  }
}
 

import org.apache.beam.runners.spark.SparkRunner; //导入依赖的package包/类
@Override
public JavaStreamingContext call() throws Exception {
  LOG.info("Creating a new Spark Streaming Context");
  // validate unbounded read properties.
  checkArgument(
      options.getMinReadTimeMillis() < options.getBatchIntervalMillis(),
      "Minimum read time has to be less than batch time.");
  checkArgument(
      options.getReadTimePercentage() > 0 && options.getReadTimePercentage() < 1,
      "Read time percentage is bound to (0, 1).");

  SparkPipelineTranslator translator =
      new StreamingTransformTranslator.Translator(new TransformTranslator.Translator());
  Duration batchDuration = new Duration(options.getBatchIntervalMillis());
  LOG.info("Setting Spark streaming batchDuration to {} msec", batchDuration.milliseconds());

  JavaSparkContext jsc = SparkContextFactory.getSparkContext(options);
  JavaStreamingContext jssc = new JavaStreamingContext(jsc, batchDuration);

  // We must first init accumulators since translators expect them to be instantiated.
  SparkRunner.initAccumulators(options, jsc);

  EvaluationContext ctxt = new EvaluationContext(jsc, pipeline, options, jssc);
  // update cache candidates
  SparkRunner.updateCacheCandidates(pipeline, translator, ctxt);
  pipeline.traverseTopologically(new SparkRunner.Evaluator(translator, ctxt));
  ctxt.computeOutputs();

  checkpoint(jssc, checkpointDir);

  return jssc;
}
 

import org.apache.beam.runners.spark.SparkRunner; //导入依赖的package包/类
@Test
public void testTrackFlattened() {
  options.setRunner(SparkRunner.class);
  JavaSparkContext jsc = SparkContextFactory.getSparkContext(options);
  JavaStreamingContext jssc = new JavaStreamingContext(jsc,
      new org.apache.spark.streaming.Duration(options.getBatchIntervalMillis()));

  Pipeline p = Pipeline.create(options);

  CreateStream<Integer> queueStream1 =
      CreateStream.of(
          VarIntCoder.of(),
          Duration.millis(options.getBatchIntervalMillis())).emptyBatch();
  CreateStream<Integer> queueStream2 =
      CreateStream.of(
          VarIntCoder.of(),
          Duration.millis(options.getBatchIntervalMillis())).emptyBatch();

  PCollection<Integer> pcol1 = p.apply(queueStream1);
  PCollection<Integer> pcol2 = p.apply(queueStream2);
  PCollection<Integer> flattened =
      PCollectionList.of(pcol1).and(pcol2).apply(Flatten.<Integer>pCollections());
  flattened.apply(ParDo.of(new PassthroughFn<>()));

  p.traverseTopologically(new StreamingSourceTracker(jssc, p, ParDo.MultiOutput.class, 0, 1));
  assertThat(StreamingSourceTracker.numAssertions, equalTo(1));
}
 

import org.apache.beam.runners.spark.SparkRunner; //导入依赖的package包/类
private StreamingSourceTracker(
    JavaStreamingContext jssc,
    Pipeline pipeline,
    Class<? extends PTransform> transformClassToAssert,
    Integer... expected) {
  this.ctxt = new EvaluationContext(jssc.sparkContext(), pipeline, options, jssc);
  this.evaluator = new SparkRunner.Evaluator(
      new StreamingTransformTranslator.Translator(new TransformTranslator.Translator()), ctxt);
  this.transformClassToAssert = transformClassToAssert;
  this.expected = expected;
}
 

import org.apache.beam.runners.spark.SparkRunner; //导入依赖的package包/类
private Pipeline createSparkRunnerPipeline() {
    PipelineOptions o = PipelineOptionsFactory.create();
    SparkContextOptions options = o.as(SparkContextOptions.class);
    JavaSparkContext jsc = new JavaSparkContext("local[2]", "PubSubInput");
    options.setProvidedSparkContext(jsc);
    options.setUsesProvidedSparkContext(true);
    options.setRunner(SparkRunner.class);
    runtimeContainer = new BeamJobRuntimeContainer(options);
    return Pipeline.create(options);
}
 

import org.apache.beam.runners.spark.SparkRunner; //导入依赖的package包/类
private Pipeline createSparkRunnerPipeline() {
    JavaSparkContext jsc = new JavaSparkContext("local[2]", this.getClass().getName());
    PipelineOptions o = PipelineOptionsFactory.create();
    SparkContextOptions options = o.as(SparkContextOptions.class);
    options.setProvidedSparkContext(jsc);
    options.setUsesProvidedSparkContext(true);
    options.setRunner(SparkRunner.class);
    runtimeContainer = new BeamJobRuntimeContainer(options);
    return Pipeline.create(options);
}
 

import org.apache.beam.runners.spark.SparkRunner; //导入依赖的package包/类
@Before
public void setupLazyAvroCoder() {
    options = PipelineOptionsFactory.as(SparkPipelineOptions.class);
    options.setRunner(SparkRunner.class);
    options.setSparkMaster("local");
    options.setStreaming(false);
    pWrite = Pipeline.create(options);
    pRead = Pipeline.create(options);

}
 

import org.apache.beam.runners.spark.SparkRunner; //导入依赖的package包/类
/**
 * @return the options used to create this pipeline. These can be or changed before the Pipeline is created.
 */
public SparkContextOptions getOptions() {
    if (options == null) {
        options = PipelineOptionsFactory.as(SparkContextOptions.class);
        options.setRunner(SparkRunner.class);
    }
    return options;
}
 

import org.apache.beam.runners.spark.SparkRunner; //导入依赖的package包/类
private Pipeline createSparkRunnerPipeline() {
    PipelineOptions o = PipelineOptionsFactory.create();
    SparkContextOptions options = o.as(SparkContextOptions.class);
    options.setProvidedSparkContext(jsc);
    options.setUsesProvidedSparkContext(true);
    options.setRunner(SparkRunner.class);
    runtimeContainer = new BeamJobRuntimeContainer(options);
    return Pipeline.create(options);
}
 

import org.apache.beam.runners.spark.SparkRunner; //导入依赖的package包/类
private Pipeline createSparkRunnerPipeline() {
    PipelineOptions o = PipelineOptionsFactory.create();
    SparkContextOptions options = o.as(SparkContextOptions.class);

    SparkConf conf = new SparkConf();
    conf.setAppName("KinesisInput");
    conf.setMaster("local[2]");
    conf.set("spark.driver.allowMultipleContexts", "true");
    JavaSparkContext jsc = new JavaSparkContext(new SparkContext(conf));
    options.setProvidedSparkContext(jsc);
    options.setUsesProvidedSparkContext(true);
    options.setRunner(SparkRunner.class);

    return Pipeline.create(options);
}