[BEAM-1636] UnboundedDataset action() does not materialize RDD

mransley · Mar 7, 2017 · a889597 · a889597
1 parent 1fd52f5
commit a889597
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 7 deletions.
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/BoundedDataset.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/BoundedDataset.java
@@ -32,7 +32,6 @@
 import org.apache.spark.api.java.JavaRDD;
 import org.apache.spark.api.java.JavaRDDLike;
 import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.api.java.function.VoidFunction;
 import org.apache.spark.storage.StorageLevel;
 
 /**
@@ -106,11 +105,7 @@ public void cache(String storageLevel) {
  @Override
  public void action() {
  // Empty function to force computation of RDD.
- rdd.foreach(new VoidFunction<WindowedValue<T>>() {
- @Override public void call(WindowedValue<T> tWindowedValue) throws Exception {
- // Empty implementation.
- }
- });
+ rdd.foreach(TranslationUtils.<WindowedValue<T>>emptyVoidFunction());
  }
 
  @Override

diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TranslationUtils.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TranslationUtils.java
@@ -43,6 +43,7 @@
 import org.apache.spark.api.java.JavaSparkContext;
 import org.apache.spark.api.java.function.Function;
 import org.apache.spark.api.java.function.PairFunction;
+import org.apache.spark.api.java.function.VoidFunction;
 import org.apache.spark.streaming.api.java.JavaDStream;
 import org.apache.spark.streaming.api.java.JavaPairDStream;
 
@@ -267,4 +268,11 @@ public static void rejectStateAndTimers(DoFn<?, ?> doFn) {
  }
  }
 
+ public static <T> VoidFunction<T> emptyVoidFunction() {
+ return new VoidFunction<T>() {
+ @Override public void call(T t) throws Exception {
+ // Empty implementation.
+ }
+ };
+ }
 }
diff --git a/...k/src/main/java/org/apache/beam/runners/spark/translation/streaming/UnboundedDataset.java b/...k/src/main/java/org/apache/beam/runners/spark/translation/streaming/UnboundedDataset.java
@@ -21,7 +21,10 @@
 import java.util.ArrayList;
 import java.util.List;
 import org.apache.beam.runners.spark.translation.Dataset;
+import org.apache.beam.runners.spark.translation.TranslationUtils;
 import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.function.VoidFunction;
 import org.apache.spark.streaming.api.java.JavaDStream;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -68,11 +71,17 @@ public void cache(String storageLevel) {
  @Override
  public void action() {
  // Force computation of DStream.
- dStream.dstream().register();
+ dStream.foreachRDD(new VoidFunction<JavaRDD<WindowedValue<T>>>() {
+ @Override
+ public void call(JavaRDD<WindowedValue<T>> rdd) throws Exception {
+ rdd.foreach(TranslationUtils.<WindowedValue<T>>emptyVoidFunction());
+ }
+ });
  }
 
  @Override
  public void setName(String name) {
  // ignore
  }
+
 }