Merge pull request #930 from holdenk/master

Add mapPartitionsWithIndex

1 files changed, 10 insertions, 0 deletions

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 7e6e691f11..7a3568c5ef 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
@@ -68,6 +68,16 @@ trait JavaRDDLike[T, This <: JavaRDDLike[T, This]] extends Serializable {
     new JavaRDD(rdd.map(f)(f.returnType()))(f.returnType())
 
   /**
+   * Return a new RDD by applying a function to each partition of this RDD, while tracking the index
+   * of the original partition.
+   */
+  def mapPartitionsWithIndex[R: ClassManifest](
+      f: JFunction2[Int, java.util.Iterator[T], java.util.Iterator[R]],
+      preservesPartitioning: Boolean = false): JavaRDD[R] =
+    new JavaRDD(rdd.mapPartitionsWithIndex(((a,b) => f(a,asJavaIterator(b))),
+					   preservesPartitioning))
+
+  /**
    * Return a new RDD by applying a function to all elements of this RDD.
    */
   def map[R](f: DoubleFunction[T]): JavaDoubleRDD =