java api zipWithIndex test

Author: witgo

core/src/main/scala/org/apache/spark/api/java/JavaRDDLike.scala

@@ -268,7 +268,9 @@ trait JavaRDDLike[T, This <: JavaRDDLike[T, This]] extends Serializable {
    * 2*n+k, ..., where n is the number of partitions. So there may exist gaps, but this method
    * won't trigger a spark job, which is different from [[org.apache.spark.rdd.RDD#zipWithIndex]].
    */
-  def zipWithUniqueId():JavaPairRDD[T,Long] = JavaPairRDD.fromRDD(rdd.zipWithUniqueId())
+  def zipWithUniqueId[Long](): JavaPairRDD[T, Long] = {
+    JavaPairRDD.fromRDD(rdd.zipWithUniqueId()).asInstanceOf[JavaPairRDD[T, Long]]
+  }
 
   /**
    * Zips this RDD with its element indices. The ordering is first based on the partition index
@@ -277,8 +279,9 @@ trait JavaRDDLike[T, This <: JavaRDDLike[T, This]] extends Serializable {
    * This is similar to Scala's zipWithIndex but it uses Long instead of Int as the index type.
    * This method needs to trigger a spark job when this RDD contains more than one partitions.
    */
-  def zipWithIndex():JavaPairRDD[T,Long] = JavaPairRDD.fromRDD(rdd.zipWithIndex())
-
+  def zipWithIndex[Long](): JavaPairRDD[T, Long] = {
+    JavaPairRDD.fromRDD(rdd.zipWithIndex()).asInstanceOf[JavaPairRDD[T, Long]]
+  }
 
   // Actions (launch a job to return a value to the user program)
 

core/src/test/java/org/apache/spark/JavaAPISuite.java

@@ -182,13 +182,39 @@ public void call(String s) {
     Assert.assertEquals(2, foreachCalls);
   }
 
-    @Test
-    public void toLocalIterator() {
-        List<Integer> correct = Arrays.asList(1, 2, 3, 4);
-        JavaRDD<Integer> rdd = sc.parallelize(correct);
-        List<Integer> result = Lists.newArrayList(rdd.toLocalIterator());
-        Assert.assertTrue(correct.equals(result));
-    }
+  @Test
+  public void toLocalIterator() {
+    List<Integer> correct = Arrays.asList(1, 2, 3, 4);
+    JavaRDD<Integer> rdd = sc.parallelize(correct);
+    List<Integer> result = Lists.newArrayList(rdd.toLocalIterator());
+    Assert.assertTrue(correct.equals(result));
+  }
+
+  @Test
+  public void zipWithUniqueId() {
+    List<Integer> correct = Arrays.asList(1, 2, 3, 4);
+    JavaPairRDD<Integer, Long> zip = sc.parallelize(correct).zipWithUniqueId();
+    JavaRDD<Long> indexes = zip.map(new Function<Tuple2<Integer, Long>, Long>() {
+      @Override
+      public Long call(Tuple2<Integer, Long> t) throws Exception {
+        return t._2();
+      }
+    });
+    Assert.assertTrue(new HashSet<Long>(indexes.collect()).size() == 4);
+  }
+
+  @Test
+  public void zipWithIndex() {
+    List<Integer> correct = Arrays.asList(1, 2, 3, 4);
+    JavaPairRDD<Integer, Long> zip = sc.parallelize(correct).zipWithIndex();
+    JavaRDD<Long> indexes = zip.map(new Function<Tuple2<Integer, Long>, Long>() {
+      @Override
+      public Long call(Tuple2<Integer, Long> t) throws Exception {
+        return t._2();
+      }
+    });
+    Assert.assertTrue(new HashSet<Long>(indexes.collect()).size() == 4);
+  }
 
   @SuppressWarnings("unchecked")
   @Test