Utility function to get a setup & cleanup function for mapping each partition

core/src/main/scala/spark/RDD.scala

@@ -1,13 +1,10 @@
 package spark
 
-import java.net.URL
-import java.util.{Date, Random}
-import java.util.{HashMap => JHashMap}
+import java.util.Random
 
 import scala.collection.Map
 import scala.collection.JavaConversions.mapAsScalaMap
 import scala.collection.mutable.ArrayBuffer
-import scala.collection.mutable.HashMap
 
 import org.apache.hadoop.io.BytesWritable
 import org.apache.hadoop.io.NullWritable
@@ -34,6 +31,8 @@ import spark.rdd.ZippedRDD
 import spark.storage.StorageLevel
 
 import SparkContext._
+import spark.RDD.PartitionMapper
+import util.CleanupIterator
 
 /**
  * A Resilient Distributed Dataset (RDD), the basic abstraction in Spark. Represents an immutable,
@@ -339,6 +338,21 @@ abstract class RDD[T: ClassManifest](
     preservesPartitioning: Boolean = false): RDD[U] =
     new MapPartitionsWithSplitRDD(this, sc.clean(f), preservesPartitioning)
 
+  /**
+   * Return a new RDD by applying a function to every element in this RDD, with extra setup & cleanup
+   * at the beginning & end of processing every partition.
+   *
+   * This might be useful if you need to setup some resources per task & cleanup them up at the end, eg.
+   * a db connection
+   */
+  def mapWithSetupAndCleanup[U: ClassManifest](m: PartitionMapper[T,U]) : RDD[U] =
+    mapPartitionsWithSplit{
+      case(partition, itr) =>
+        m.setup(partition)
+        val subItr = itr.map{m.map}
+        CleanupIterator[U,Iterator[U]](subItr, m.cleanup _)
+    }
+
   /**
    * Zips this RDD with another one, returning key-value pairs with the first element in each RDD,
    * second element in each RDD, etc. Assumes that the two RDDs have the *same number of
@@ -680,3 +694,27 @@ abstract class RDD[T: ClassManifest](
     origin)
 
 }
+
+object RDD {
+
+  /**
+   * Defines a map function over elements of an RDD, but with extra setup and cleanup
+   * that happens
+   */
+  trait PartitionMapper[T,U] extends Serializable {
+    /**
+     * called at the start of processing of each partition
+     */
+    def setup(partiton:Int)
+
+    /**
+     * transform one element of the partition
+     */
+    def map(t: T) : U
+
+    /**
+     * called at the end of each partition
+     */
+    def cleanup
+  }
+}

core/src/main/scala/spark/util/CleanupIterator.scala

@@ -0,0 +1,25 @@
+package spark.util
+
+/**
+ * Wrapper around an iterator which calls a cleanup method when its finished iterating through its elements
+ */
+abstract class CleanupIterator[+A, +I <: Iterator[A]](sub: I) extends Iterator[A]{
+  def next = sub.next
+  def hasNext = {
+    val r = sub.hasNext
+    if (!r) {
+      cleanup
+    }
+    r
+  }
+
+  def cleanup
+}
+
+object CleanupIterator {
+  def apply[A, I <: Iterator[A]](sub: I, cleanupFunction: () => Unit) : CleanupIterator[A,I] = {
+    new CleanupIterator[A,I](sub) {
+      def cleanup = cleanupFunction()
+    }
+  }
+}

core/src/test/scala/spark/RDDSuite.scala

@@ -4,6 +4,8 @@ import scala.collection.mutable.HashMap
 import org.scalatest.FunSuite
 import spark.SparkContext._
 import spark.rdd.{CoalescedRDD, PartitionPruningRDD}
+import spark.RDD.PartitionMapper
+import collection._
 
 class RDDSuite extends FunSuite with LocalSparkContext {
 
@@ -173,4 +175,48 @@ class RDDSuite extends FunSuite with LocalSparkContext {
     assert(prunedData.size === 1)
     assert(prunedData(0) === 10)
   }
+
+  test("mapPartitionWithSetupAndCleanup") {
+    sc = new SparkContext("local[4]", "test")
+    val data = sc.parallelize(1 to 100, 4)
+    val acc = sc.accumulableCollection(new HashMap[Int,Set[Int]]())
+    val mapped = data.mapWithSetupAndCleanup(new PartitionMapper[Int,Int](){
+      var partition = -1
+      val values = mutable.Set[Int]()
+      def setup(partition:Int) {this.partition = partition}
+      def map(i:Int) = {values += i; i * 2}
+      def cleanup = {
+        //the purpose of this strange code is just to make sure this method is called
+        // after the data has been iterated through completely.
+        acc.localValue += (partition -> Set[Int](values.toSeq: _*))
+      }
+    }).collect
+
+    assert(mapped.toSet === (1 to 100).map{_ * 2}.toSet)
+    assert(acc.value.keySet == (0 to 3).toSet)
+    acc.value.foreach { case(partition, values) =>
+      assert(values.size === 25)
+    }
+
+
+    //the naive alternative doesn't work
+    val acc2 = sc.accumulableCollection(new HashMap[Int,Set[Int]]())
+    val m2 = data.mapPartitionsWithSplit{
+      case (partition, itr) =>
+        val values = mutable.Set[Int]()
+        val mItr = itr.map{i => values += i; i * 2}
+        //you haven't actually put anything into values yet, b/c itr.map defines another
+        // iterator, which is lazily computed.  so the Set is empty
+        acc2.localValue += (partition -> Set[Int](values.toSeq: _*))
+        mItr
+    }.collect
+
+    assert(m2.toSet === (1 to 100).map{_ * 2}.toSet)
+    assert(acc2.value.keySet == (0 to 3).toSet)
+    //this condition will fail
+//    acc2.value.foreach { case(partition, values) =>
+//      assert(values.size === 25)
+//    }
+
+  }
 }