Added mapPartitions method to JavaRDD.

1 files changed, 12 insertions, 1 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 6c91edaf5c..568ae1575b 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
@@ -21,8 +21,10 @@ import scala.reflect.ClassTag
 
 import org.apache.spark._
 import org.apache.spark.rdd.RDD
-import org.apache.spark.api.java.function.{Function => JFunction}
+import org.apache.spark.api.java.function.{Function => JFunction, FlatMapFunction => JFMap}
 import org.apache.spark.storage.StorageLevel
+import java.util.{Iterator => JIterator}
+import scala.collection.JavaConversions._
 
 class JavaRDD[T](val rdd: RDD[T])(implicit val classTag: ClassTag[T]) extends
 JavaRDDLike[T, JavaRDD[T]] {
@@ -138,6 +140,15 @@ JavaRDDLike[T, JavaRDD[T]] {
   def setGenerator(_generator: String) = {
     rdd.generator = _generator
   }
+
+  /**
+   * Return a new RDD by applying a function to each partition of this RDD.
+   */
+  def mapPartitions[U: ClassTag](
+      f: JFMap[JIterator[T], U], preservesPartitioning: Boolean = false): JavaRDD[U] = {
+    rdd.mapPartitions[U]((x => f(asJavaIterator(x)).iterator), preservesPartitioning)
+  }
+
 }
 
 object JavaRDD {