Fix Java API for mapPartitionsWithIndex

holdenk · Mar 8, 2014 · 8d849a1 · 8d849a1
1 parent 0b7b7fd
commit 8d849a1
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 5 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
@@ -73,11 +73,13 @@ trait JavaRDDLike[T, This <: JavaRDDLike[T, This]] extends Serializable {
    * 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: ClassTag](
+  def mapPartitionsWithIndex[R](f: MapPartitionsWithIndexFunction[T, R],
-      f: JFunction2[java.lang.Integer, java.util.Iterator[T], java.util.Iterator[R]],
+    preservesPartitioning: Boolean = false): JavaRDD[R] = {
-      preservesPartitioning: Boolean = false): JavaRDD[R] =
+    import scala.collection.JavaConverters._
-    new JavaRDD(rdd.mapPartitionsWithIndex(((a,b) => f(a,asJavaIterator(b))),
+    def fn = (a: Int, b: Iterator[T]) => f.apply(a, asJavaIterator(b)).asScala
-        preservesPartitioning))
+    val newRdd = rdd.mapPartitionsWithIndex(fn, preservesPartitioning)(f.elementType())
+    new JavaRDD(newRdd)(f.elementType())
+  }
 
   /**
    * Return a new RDD by applying a function to all elements of this RDD.

diff --git a/core/src/test/java/org/apache/spark/JavaAPISuite.java b/core/src/test/java/org/apache/spark/JavaAPISuite.java
@@ -414,6 +414,26 @@ public void javaDoubleRDDHistoGram() {
    Assert.assertArrayEquals(expected_counts, histogram);
   }
 
+  @Test
+  public void mapPartitionsWithIndex() {
+    JavaRDD<Integer> rdd = sc.parallelize(Arrays.asList(1, 2, 3, 4, 5));
+    JavaRDD<Integer> rddByIndex =
+	rdd.mapPartitionsWithIndex(new MapPartitionsWithIndexFunction<Integer, Integer>() {
+      @Override
+      public Iterator<Integer> call(Integer start, java.util.Iterator<Integer> iter) {
+	  List<Integer> list = new ArrayList<Integer>();
+	int pos = start;
+	while (iter.hasNext()) {
+	  list.add(iter.next() * pos);
+	  pos += 1;
+	}
+	return list.iterator();
+      }
+    }, false);
+    Assert.assertEquals(0, rddByIndex.first().intValue());
+  }
+
+
   @Test
   public void map() {
     JavaRDD<Integer> rdd = sc.parallelize(Arrays.asList(1, 2, 3, 4, 5));