Modifications as suggested in PR feedback-

- mapPartitions, foreachPartition moved to JavaRDDLike
- call scala rdd's setGenerator instead of setting directly in JavaRDD

2 files changed, 16 insertions, 17 deletions

diff --git a/core/src/main/scala/org/apache/spark/api/java/JavaRDD.scala b/core/src/main/scala/org/apache/spark/api/java/JavaRDD.scala
index e7b1ceebf6..e687bbdd99 100644
--- a/core/src/main/scala/org/apache/spark/api/java/JavaRDD.scala
+++ b/core/src/main/scala/org/apache/spark/api/java/JavaRDD.scala
@@ -138,24 +138,8 @@ JavaRDDLike[T, JavaRDD[T]] {
 
   /** Reset generator*/
   def setGenerator(_generator: String) = {
-    rdd.generator = _generator
+    rdd.setGenerator(_generator)
   }
-
-  /**
-   * Return a new RDD by applying a function to each partition of this RDD.
-   */
-  def mapPartitions[U](
-      f: JFMap[JIterator[T], U], preservesPartitioning: Boolean): JavaRDD[U] = {
-    rdd.mapPartitions[U]((x => f(asJavaIterator(x)).iterator), preservesPartitioning)
-  }
-
-  /**
-   * Applies a function f to each partition of this RDD.
-   */
-  def foreachPartition(f: VoidFunction[JIterator[T]]) {
-    rdd.foreachPartition((x => f(asJavaIterator(x))))
-  }
-
 }
 
 object JavaRDD {
diff --git a/core/src/main/scala/org/apache/spark/api/java/JavaRDDLike.scala b/core/src/main/scala/org/apache/spark/api/java/JavaRDDLike.scala
index ebbbbd8806..eb8e34e240 100644
--- a/core/src/main/scala/org/apache/spark/api/java/JavaRDDLike.scala
+++ b/core/src/main/scala/org/apache/spark/api/java/JavaRDDLike.scala
@@ -149,6 +149,21 @@ trait JavaRDDLike[T, This <: JavaRDDLike[T, This]] extends Serializable {
   }
 
   /**
+   * Return a new RDD by applying a function to each partition of this RDD.
+   */
+  def mapPartitions[U](
+      f: FlatMapFunction[java.util.Iterator[T], U], preservesPartitioning: Boolean): JavaRDD[U] = {
+    rdd.mapPartitions[U]((x => f(asJavaIterator(x)).iterator), preservesPartitioning)
+  }
+
+  /**
+   * Applies a function f to each partition of this RDD.
+   */
+  def foreachPartition(f: VoidFunction[java.util.Iterator[T]]) {
+    rdd.foreachPartition((x => f(asJavaIterator(x))))
+  }
+
+  /**
    * Return an RDD created by coalescing all elements within each partition into an array.
    */
   def glom(): JavaRDD[JList[T]] =