This closes apache#2046

cosmoskitten · Feb 20, 2017 · 4e5a762 · 4e5a762
2 parents aa45ccb + a2f0615
commit 4e5a762
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 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
@@ -20,6 +20,7 @@
 
 import com.google.common.base.Function;
 import com.google.common.collect.Iterables;
+import java.util.Iterator;
 import java.util.List;
 import javax.annotation.Nullable;
 import org.apache.beam.runners.spark.coders.CoderHelpers;
@@ -32,6 +33,7 @@
 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;
 
 /**
@@ -104,7 +106,13 @@ public void cache(String storageLevel) {
 
   @Override
   public void action() {
-    rdd.count();
+    // Empty function to force computation of RDD.
+    rdd.foreachPartition(new VoidFunction<Iterator<WindowedValue<T>>>() {
+      @Override
+      public void call(Iterator<WindowedValue<T>> windowedValueIterator) throws Exception {
+        // Empty implementation.
+      }
+    });
   }
 
   @Override

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
@@ -32,7 +32,6 @@
 import org.apache.beam.sdk.coders.Coder;
 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.apache.spark.streaming.api.java.JavaStreamingContext;
 import org.slf4j.Logger;
@@ -115,12 +114,8 @@ public void cache(String storageLevel) {
 
   @Override
   public void action() {
-    dStream.foreachRDD(new VoidFunction<JavaRDD<WindowedValue<T>>>() {
-      @Override
-      public void call(JavaRDD<WindowedValue<T>> rdd) throws Exception {
-        rdd.count();
-      }
-    });
+    // Force computation of DStream.
+    dStream.dstream().register();
   }
 
   @Override