beam-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From amits...@apache.org
Subject [7/9] beam git commit: Streaming sources tracking test.
Date Mon, 20 Feb 2017 09:58:03 GMT
Streaming sources tracking test.


Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/add87166
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/add87166
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/add87166

Branch: refs/heads/master
Commit: add87166fd445d634f6faf0f838f234c3f908c2e
Parents: 9784f20
Author: Sela <ansela@paypal.com>
Authored: Sun Feb 12 19:23:28 2017 +0200
Committer: Sela <ansela@paypal.com>
Committed: Mon Feb 20 11:30:16 2017 +0200

----------------------------------------------------------------------
 .../translation/streaming/UnboundedDataset.java |   6 +
 .../streaming/TrackStreamingSourcesTest.java    | 152 +++++++++++++++++++
 2 files changed, 158 insertions(+)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/beam/blob/add87166/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/UnboundedDataset.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/UnboundedDataset.java
b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/UnboundedDataset.java
index 80c0515..08d1ab6 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/UnboundedDataset.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/UnboundedDataset.java
@@ -18,6 +18,7 @@
 
 package org.apache.beam.runners.spark.translation.streaming;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Iterables;
 
 import java.util.ArrayList;
@@ -72,6 +73,11 @@ public class UnboundedDataset<T> implements Dataset {
     this.streamingSources.add(queuedStreamIds.decrementAndGet());
   }
 
+  @VisibleForTesting
+  public static void resetQueuedStreamIds() {
+    queuedStreamIds.set(0);
+  }
+
   @SuppressWarnings("ConstantConditions")
   JavaDStream<WindowedValue<T>> getDStream() {
     if (dStream == null) {

http://git-wip-us.apache.org/repos/asf/beam/blob/add87166/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/TrackStreamingSourcesTest.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/TrackStreamingSourcesTest.java
b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/TrackStreamingSourcesTest.java
new file mode 100644
index 0000000..f102ac8
--- /dev/null
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/TrackStreamingSourcesTest.java
@@ -0,0 +1,152 @@
+package org.apache.beam.runners.spark.translation.streaming;
+
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.core.IsEqual.equalTo;
+import static org.junit.Assert.assertThat;
+
+import java.util.Collections;
+import java.util.List;
+import org.apache.beam.runners.spark.ReuseSparkContext;
+import org.apache.beam.runners.spark.SparkPipelineOptions;
+import org.apache.beam.runners.spark.SparkRunner;
+import org.apache.beam.runners.spark.io.CreateStream;
+import org.apache.beam.runners.spark.translation.Dataset;
+import org.apache.beam.runners.spark.translation.EvaluationContext;
+import org.apache.beam.runners.spark.translation.SparkContextFactory;
+import org.apache.beam.runners.spark.translation.TransformTranslator;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.runners.TransformHierarchy;
+import org.apache.beam.sdk.transforms.AppliedPTransform;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.Flatten;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionList;
+import org.apache.beam.sdk.values.PValue;
+import org.apache.spark.SparkStatusTracker;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.streaming.Duration;
+import org.apache.spark.streaming.api.java.JavaStreamingContext;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+
+
+/**
+ * A test suite that tests tracking of the streaming sources created an
+ * {@link org.apache.beam.runners.spark.translation.streaming.UnboundedDataset}.
+ */
+public class TrackStreamingSourcesTest {
+
+  @Rule
+  public ReuseSparkContext reuseContext = ReuseSparkContext.yes();
+
+  private static final transient SparkPipelineOptions options =
+      PipelineOptionsFactory.create().as(SparkPipelineOptions.class);
+
+  @Before
+  public void before() {
+    UnboundedDataset.resetQueuedStreamIds();
+    StreamingSourceTracker.numAssertions = 0;
+  }
+
+  @Test
+  public void testTrackSingle() {
+    options.setRunner(SparkRunner.class);
+    JavaSparkContext jsc = SparkContextFactory.getSparkContext(options);
+    JavaStreamingContext jssc = new JavaStreamingContext(jsc,
+        new Duration(options.getBatchIntervalMillis()));
+
+    Pipeline p = Pipeline.create(options);
+
+    CreateStream.QueuedValues<Integer> queueStream =
+        CreateStream.fromQueue(Collections.<Iterable<Integer>>emptyList());
+
+    p.apply(queueStream).setCoder(VarIntCoder.of())
+        .apply(ParDo.of(new PassthroughFn<>()));
+
+    p.traverseTopologically(new StreamingSourceTracker(jssc, p, ParDo.Bound.class,  -1));
+    assertThat(StreamingSourceTracker.numAssertions, equalTo(1));
+  }
+
+  @Test
+  public void testTrackFlattened() {
+    options.setRunner(SparkRunner.class);
+    JavaSparkContext jsc = SparkContextFactory.getSparkContext(options);
+    JavaStreamingContext jssc = new JavaStreamingContext(jsc,
+        new Duration(options.getBatchIntervalMillis()));
+
+    Pipeline p = Pipeline.create(options);
+
+    CreateStream.QueuedValues<Integer> queueStream1 =
+        CreateStream.fromQueue(Collections.<Iterable<Integer>>emptyList());
+    CreateStream.QueuedValues<Integer> queueStream2 =
+        CreateStream.fromQueue(Collections.<Iterable<Integer>>emptyList());
+
+    PCollection<Integer> pcol1 = p.apply(queueStream1).setCoder(VarIntCoder.of());
+    PCollection<Integer> pcol2 = p.apply(queueStream2).setCoder(VarIntCoder.of());
+    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.Bound.class, -1, -2));
+    assertThat(StreamingSourceTracker.numAssertions, equalTo(1));
+  }
+
+  private static class PassthroughFn<T> extends DoFn<T, T> {
+    @ProcessElement
+    public void processElement(ProcessContext c) {
+      c.output(c.element());
+    }
+  }
+
+  private static class StreamingSourceTracker extends Pipeline.PipelineVisitor.Defaults {
+    private final EvaluationContext ctxt;
+    private final SparkRunner.Evaluator evaluator;
+    private final Class<? extends PTransform> transformClassToAssert;
+    private final Integer[] expected;
+
+    private static int numAssertions = 0;
+
+    private StreamingSourceTracker(
+        JavaStreamingContext jssc,
+        Pipeline pipeline,
+        Class<? extends PTransform> transformClassToAssert,
+        Integer... expected) {
+      this.ctxt = new EvaluationContext(jssc.sparkContext(), pipeline, jssc);
+      this.evaluator = new SparkRunner.Evaluator(
+          new StreamingTransformTranslator.Translator(new TransformTranslator.Translator()),
ctxt);
+      this.transformClassToAssert = transformClassToAssert;
+      this.expected = expected;
+    }
+
+    private void assertSourceIds(List<Integer> streamingSources) {
+      numAssertions++;
+      assertThat(streamingSources, containsInAnyOrder(expected));
+    }
+
+    @Override
+    public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {
+      return evaluator.enterCompositeTransform(node);
+    }
+
+    @Override
+    public void visitPrimitiveTransform(TransformHierarchy.Node node) {
+      PTransform transform = node.getTransform();
+      if (transform.getClass() == transformClassToAssert) {
+        AppliedPTransform<?, ?, ?> appliedTransform = node.toAppliedPTransform();
+        ctxt.setCurrentTransform(appliedTransform);
+        //noinspection unchecked
+        Dataset dataset = ctxt.borrowDataset((PTransform<? extends PValue, ?>) transform);
+        assertSourceIds(((UnboundedDataset<?>) dataset).getStreamingSources());
+        ctxt.setCurrentTransform(null);
+      } else {
+        evaluator.visitPrimitiveTransform(node);
+      }
+    }
+  }
+
+}


Mime
View raw message