SPARK-2895: Add mapPartitionsWithContext related support on Spark Java API.

core/src/main/java/org/apache/spark/api/java/function/DoubleFlatMapFunction2.java

@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+
+/**
+ * A function that takes arguments of type T1 and T2, and returns zero or more records of type
+ * Double from each input record.
+ */
+public interface DoubleFlatMapFunction2<T1, T2> extends Serializable {
+  public Iterable<Double> call(T1 t1, T2 t2) throws Exception;
+}

core/src/main/java/org/apache/spark/api/java/function/PairFlatMapFunction2.java

@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import scala.Tuple2;
+
+import java.io.Serializable;
+
+/**
+ * A function that takes arguments of type T1 and T2, and returns zero or more key-value pair
+ * records from each input record. The key-value pairs are represented as scala.Tuple2 objects.
+ */
+public interface PairFlatMapFunction2<T1, T2 , K, V> extends Serializable {
+  public Iterable<Tuple2<K, V>> call(T1 t1, T2 t2) throws Exception;
+}

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

@@ -27,7 +27,7 @@ import com.google.common.base.Optional
 import org.apache.hadoop.io.compress.CompressionCodec
 
 import org.apache.spark.{FutureAction, Partition, SparkContext, TaskContext}
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{DeveloperApi, Experimental}
 import org.apache.spark.api.java.JavaPairRDD._
 import org.apache.spark.api.java.JavaSparkContext.fakeClassTag
 import org.apache.spark.api.java.function.{Function => JFunction, Function2 => JFunction2, _}
@@ -81,7 +81,7 @@ trait JavaRDDLike[T, This <: JavaRDDLike[T, This]] extends Serializable {
   def mapPartitionsWithIndex[R](
       f: JFunction2[java.lang.Integer, java.util.Iterator[T], java.util.Iterator[R]],
       preservesPartitioning: Boolean = false): JavaRDD[R] =
-    new JavaRDD(rdd.mapPartitionsWithIndex(((a,b) => f(a,asJavaIterator(b))),
+    new JavaRDD(rdd.mapPartitionsWithIndex(((a, b) => f(a, asJavaIterator(b))),
         preservesPartitioning)(fakeClassTag))(fakeClassTag)
 
   /**
@@ -185,6 +185,62 @@ trait JavaRDDLike[T, This <: JavaRDDLike[T, This]] extends Serializable {
       rdd.mapPartitions(fn, preservesPartitioning))(fakeClassTag[K2], fakeClassTag[V2])
   }
 
+  /**
+   * :: DeveloperApi ::
+   * Return a new RDD by applying a function to each partition of this RDD. This is a variant of
+   * mapPartitions that also passes the TaskContext into the closure.
+   *
+   * `preservesPartitioning` indicates whether the input function preserves the partitioner, which
+   * should be `false` unless this is a pair RDD and the input function doesn't modify the keys.
+   */
+  @DeveloperApi
+  def mapPartitionsWithContext[R](
+      f: JFunction2[TaskContext, java.util.Iterator[T], java.util.Iterator[R]],
+      preservesPartitioning: Boolean = false): JavaRDD[R] = {
+
+    new JavaRDD(rdd.mapPartitionsWithContext(
+      ((a, b) => f(a, asJavaIterator(b))), preservesPartitioning)(fakeClassTag))(fakeClassTag)
+  }
+
+
+  /**
+   * :: DeveloperApi ::
+   * Return a new JavaDoubleRDD by applying a function to each partition of this RDD. This is a
+   * variant of mapPartitions that also passes the TaskContext into the closure.
+   *
+   * `preservesPartitioning` indicates whether the input function preserves the partitioner, which
+   * should be `false` unless this is a pair RDD and the input function doesn't modify the keys.
+   */
+  @DeveloperApi
+  def mapPartitionsToDoubleWithContext(
+      f: DoubleFlatMapFunction2[TaskContext, java.util.Iterator[T]],
+      preservesPartitioning: Boolean): JavaDoubleRDD = {
+
+    def fn = (context: TaskContext, x: Iterator[T]) =>
+      asScalaIterator(f.call(context, asJavaIterator(x)).iterator())
+    new JavaDoubleRDD(
+      rdd.mapPartitionsWithContext(fn, preservesPartitioning).map(x => x.doubleValue()))
+  }
+
+  /**
+   * :: DeveloperApi ::
+   * Return a new JavaPairRDD by applying a function to each partition of this RDD. This is a
+   * variant of mapPartitions that also passes the TaskContext into the closure.
+   *
+   * `preservesPartitioning` indicates whether the input function preserves the partitioner, which
+   * should be `false` unless this is a pair RDD and the input function doesn't modify the keys.
+   */
+  @DeveloperApi
+  def mapPartitionsToPairWithContext[K2, V2](
+      f: PairFlatMapFunction2[TaskContext, java.util.Iterator[T], K2, V2],
+      preservesPartitioning: Boolean): JavaPairRDD[K2, V2] = {
+
+    def fn = (context: TaskContext, x: Iterator[T]) =>
+      asScalaIterator(f.call(context, asJavaIterator(x)).iterator())
+    JavaPairRDD.fromRDD(
+      rdd.mapPartitionsWithContext(fn, preservesPartitioning))(fakeClassTag[K2], fakeClassTag[V2])
+  }
+
   /**
    * Applies a function f to each partition of this RDD.
    */

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

@@ -707,6 +707,104 @@ public Iterable<Integer> call(Iterator<Integer> iter) {
     Assert.assertEquals("[3, 7]", partitionSums.collect().toString());
   }
 
+  @Test
+  public void mapPartitionsWithContext() {
+    JavaRDD<Integer> rdd = sc.parallelize(Arrays.asList(1, 2, 3, 4), 2);
+    JavaRDD<String> partitionSumsWithContext = rdd.mapPartitionsWithContext(
+      new Function2<TaskContext, Iterator<Integer>, Iterator<String>>() {
+        @Override
+        public Iterator<String> call(TaskContext context,
+          Iterator<Integer> iter) throws Exception {
+
+          int sum = 0;
+          while (iter.hasNext()) {
+            sum += iter.next();
+          }
+          return Collections.singletonList(sum + "-partition-" + context.partitionId()).iterator();
+        }
+      }, false);
+    Assert.assertEquals("[3-partition-0, 7-partition-1]",
+            partitionSumsWithContext.collect().toString());
+  }
+
+  @Test
+  public void mapPartitionsToPair() {
+    JavaRDD<Integer> rdd = sc.parallelize(Arrays.asList(1, 2, 3, 4), 2);
+    JavaPairRDD<Integer, String> pairRdd = rdd.mapPartitionsToPair(
+      new PairFlatMapFunction<Iterator<Integer>, Integer, String>() {
+        @Override
+        public Iterable<Tuple2<Integer, String>> call(Iterator<Integer> iter) throws Exception {
+          int sum = 0;
+          while (iter.hasNext()) {
+            sum += iter.next();
+          }
+          return Collections.singletonList(new Tuple2<Integer, String>(sum, "a"));
+        }
+      }
+    );
+
+    Assert.assertEquals("[(3,a), (7,a)]", pairRdd.collect().toString());
+  }
+
+  @Test
+  public void mapPartitionsToPairWithContext() {
+    JavaRDD<Integer> rdd = sc.parallelize(Arrays.asList(1, 2, 3, 4), 2);
+    JavaPairRDD<Integer, String> pairRdd = rdd.mapPartitionsToPairWithContext(
+      new PairFlatMapFunction2<TaskContext, Iterator<Integer>, Integer, String>() {
+        @Override
+        public Iterable<Tuple2<Integer, String>> call(TaskContext context, Iterator<Integer> iter)
+          throws Exception {
+
+          int sum = 0;
+          while (iter.hasNext()) {
+            sum += iter.next();
+          }
+          return Collections.singletonList(
+                   new Tuple2<Integer, String>(sum, "partition-" + context.partitionId()));
+        }
+      }, false);
+
+    Assert.assertEquals("[(3,partition-0), (7,partition-1)]", pairRdd.collect().toString());
+  }
+
+  @Test
+  public void mapPartitionsToDouble() {
+    JavaRDD<Integer> rdd = sc.parallelize(Arrays.asList(1, 2, 3, 4), 2);
+    JavaDoubleRDD pairRdd = rdd.mapPartitionsToDouble(
+      new DoubleFlatMapFunction<Iterator<Integer>>() {
+        @Override
+        public Iterable<Double> call(Iterator<Integer> iter) throws Exception {
+          int sum = 0;
+          while (iter.hasNext()) {
+            sum += iter.next();
+          }
+          return Collections.singletonList(Double.valueOf(sum));
+        }
+      }
+    );
+
+    Assert.assertEquals("[3.0, 7.0]", pairRdd.collect().toString());
+  }
+
+  @Test
+  public void mapPartitionsToDoubleWithContext() {
+    JavaRDD<Integer> rdd = sc.parallelize(Arrays.asList(1, 2, 3, 4), 2);
+    JavaDoubleRDD pairRdd = rdd.mapPartitionsToDoubleWithContext(
+      new DoubleFlatMapFunction2<TaskContext, Iterator<Integer>>() {
+      @Override
+      public Iterable<Double> call(TaskContext context, Iterator<Integer> iter) throws Exception {
+        int sum = 0;
+        while (iter.hasNext()) {
+          sum += iter.next();
+        }
+        sum += context.partitionId();
+        return Collections.singletonList(Double.valueOf(sum));
+      }
+    }, false);
+
+    Assert.assertEquals("[3.0, 8.0]", pairRdd.collect().toString());
+  }
+
   @Test
   public void repartition() {
     // Shrinking number of partitions