add mapPartitionsWithIndex

Author: jey

src/Spock.jl

@@ -4,7 +4,7 @@ module Spock
   using Scotty
   using JavaCall
   import Base: Callable, map, collect, convert, count, reduce
-  export SparkContext, RDD, parallelize
+  export SparkContext, RDD, parallelize, transform
 
   const classpath = get(ENV, "CLASSPATH", "")
   JavaCall.init(["-ea", "-Xmx1024M", "-Djava.class.path=$(classpath)"])
@@ -13,7 +13,7 @@ module Spock
   JArrays = @jimport java.util.Arrays
   JList = @jimport java.util.List
   JFunction = @jimport org.apache.spark.api.java.function.Function
-  JFlatMapFunction = @jimport org.apache.spark.api.java.function.FlatMapFunction
+  JFunction2 = @jimport org.apache.spark.api.java.function.Function2
   JJavaRDD = @jimport org.apache.spark.api.java.JavaRDD
   JJavaSparkContext = @jimport org.apache.spark.api.java.JavaSparkContext
   JJuliaRDD = @jimport edu.berkeley.bids.spock.JuliaRDD
@@ -49,7 +49,7 @@ module Spock
   function jrdd(rdd::TransformedRDD)
     if rdd.jrdd === nothing
       jfunc = JJuliaFunction((JJuliaObject,), jbox(rdd.task))
-      rdd.jrdd = jcall(jrdd(rdd.parent), "mapPartitions", JJavaRDD, (JFlatMapFunction,), jfunc)
+      rdd.jrdd = jcall(jrdd(rdd.parent), "mapPartitionsWithIndex", JJavaRDD, (JFunction2, jboolean), jfunc, false)
     end
     rdd.jrdd::JJavaRDD
   end
@@ -78,7 +78,11 @@ module Spock
     deserialize(IOBuffer(payload))
   end
 
-  function transform(rdd::RDD, task)
+  # Analogous to `mapPartitionsWithIndex`.
+  #
+  # `task` will be called once per input partition with arguments of
+  # (partition_id, input_iter) and returns an iterable of new contents.
+  function transform(task::Function, rdd::RDD)
     if isa(rdd, TransformedRDD) && ispipelineable(rdd)
       TransformedRDD(rdd.parent, pipetask(task, rdd.task))
     else
@@ -87,11 +91,11 @@ module Spock
   end
 
   function map(f::Callable, rdd::RDD)
-    transform(rdd, maptask(f))
+    transform(maptask(f), rdd)
   end
 
   function reduce(f::Callable, rdd::RDD)
-    reduce(f, collect(transform(rdd, reducetask(f))))
+    reduce(f, collect(transform(reducetask(f), rdd)))
   end
 
   function collect(rdd::RDD)

src/edu/berkeley/cs/amplab/spock/JuliaFunction.java

@@ -1,7 +1,7 @@
 package edu.berkeley.bids.spock;
 import org.apache.spark.rdd.RDD;
 import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.function.FlatMapFunction;
+import org.apache.spark.api.java.function.Function2;
 import java.io.BufferedInputStream;
 import java.io.BufferedOutputStream;
 import java.io.DataInputStream;
@@ -10,7 +10,7 @@
 import java.util.Iterator;
 import java.util.LinkedList;
 
-public class JuliaFunction implements FlatMapFunction<Iterator<JuliaObject>, JuliaObject> {
+public class JuliaFunction implements Function2<Integer, Iterator<JuliaObject>, Iterator<JuliaObject>> {
   private static final long serialVersionUID = 1;
   final JuliaObject func;
 
@@ -23,7 +23,7 @@ String getScottyPath() {
   }
 
   @Override
-  public Iterable<JuliaObject> call(Iterator<JuliaObject> args) throws Exception {
+  public Iterator<JuliaObject> call(Integer partId, Iterator<JuliaObject> args) throws Exception {
     // launch worker
     ProcessBuilder pb = new ProcessBuilder("julia", "-L", getScottyPath(), "-e", "Scotty.worker()");
     pb.redirectError(ProcessBuilder.Redirect.INHERIT);
@@ -32,6 +32,7 @@ public Iterable<JuliaObject> call(Iterator<JuliaObject> args) throws Exception {
     // send input
     DataOutputStream out = new DataOutputStream(new BufferedOutputStream(worker.getOutputStream()));
     func.write(out);
+    out.writeInt(partId.intValue());
     while(args.hasNext()) {
       args.next().write(out);
     }
@@ -51,7 +52,7 @@ public Iterable<JuliaObject> call(Iterator<JuliaObject> args) throws Exception {
     if(worker.waitFor() != 0) {
       throw new RuntimeException(String.format("Spock worker died with exitValue=%d", worker.exitValue()));
     } else {
-      return results;
+      return results.iterator();
     }
   }
 }

src/scotty.jl

@@ -1,5 +1,5 @@
 module Scotty
-  export maptask, reducetask, pipetask, intask, outtask
+  export maptask, reducetask, pipetask
   const inf = STDIN
   const outf = STDOUT
 
@@ -24,23 +24,23 @@ module Scotty
     end
   end
 
-  function outtask(inp)
-    map(writeobj, inp)
+  function outtask(iter)
+    map(writeobj, iter)
   end
 
   function maptask(f)
-    (inp) -> begin
+    (partid, iter) -> begin
       Task() do
-        map(arg -> produce(f(arg)), inp)
+        map(arg -> produce(f(arg)), iter)
       end
     end
   end
 
   function reducetask(f)
-    (inp) -> begin
+    (partid, iter) -> begin
       Task() do
         accum = nothing
-        for arg in inp
+        for arg in iter
           if accum == nothing
             accum = arg
           else
@@ -53,16 +53,17 @@ module Scotty
   end
 
   function pipetask(t2, t1)
-    (inp) -> t2(t1(inp))
+    (partid, iter) -> t2(partid, t1(partid, iter))
   end
 
   function worker()
     try
       redirect_stdout(STDERR)
       close(redirect_stdin()[2])
       task = readobj()
+      partid = readint()
       try
-        outtask(task(intask()))
+        outtask(task(partid, intask()))
       catch exc
         writeint(0)
         writeint(2) # OOB 2: task error (fatal)

test/runtests.jl

@@ -2,16 +2,19 @@ using Spock
 using Base.Test
 
 sc = SparkContext()
-
 rdd1 = parallelize(sc, 1:10, 2)
+rdd2 = map(x -> x^2, rdd1)
+rddx = parallelize(sc, fill("x", 10), 2)
+allequal(val, rdd) = all(x -> x == val, collect(rdd))
+
+# test basics
 @assert 10 == count(rdd1)
 @assert 55 == sum(collect(rdd1))
 @assert 55 == reduce(+, rdd1)
-
-rdd2 = map(x -> x^2, rdd1)
 @assert 385 == sum(collect(rdd2))
 @assert 385 == reduce(+, rdd2)
 
+# test propagation of exceptions from workers
 @assert "moo" == begin
   try
     collect(map(x->throw("woof"), rdd1))
@@ -23,6 +26,7 @@ rdd2 = map(x -> x^2, rdd1)
 end
 
 let driverpid = getpid()
+  # test containment of worker I/O
   @assert 55 == reduce(rdd1) do x, y
     if getpid() != driverpid
       println("Ekke Ekke Ekke Ekke Ptangya Zoooooooom Boing Ni!")
@@ -31,13 +35,33 @@ let driverpid = getpid()
     x + y
   end
 
-  rdd3 = map(x->getpid(), rdd1)
-  @assert 2 == length(Set(collect(rdd3)))
-  @assert driverpid == reduce(rdd3) do l, r
+  # test pipelining
+  pids = map(x->getpid(), rdd1)
+  @assert 2 == length(Set(collect(pids)))
+  @assert driverpid == reduce(pids) do l, r
     mypid = getpid()
     @assert mypid == driverpid || (mypid == l && mypid == r)
     mypid
   end
 end
 
+# test mixing synchronous and asynchronous transforms
+frob = (partid, iter) -> ["f$(x)" for x in collect(iter)]
+brof = (partid, iter) -> begin
+  Task() do
+    for x in iter
+      produce("b$(x)")
+    end
+  end
+end
+@assert allequal("fbfbx", transform(frob, transform(brof, transform(frob, transform(brof, rddx)))))
+@assert allequal("bfbfx", transform(brof, transform(frob, transform(brof, transform(frob, rddx)))))
+@assert allequal("ffbfx", transform(frob, transform(frob, transform(brof, transform(frob, rddx)))))
+@assert allequal("fbbfx", transform(frob, transform(brof, transform(brof, transform(frob, rddx)))))
+@assert allequal("bbfbx", transform(brof, transform(brof, transform(frob, transform(brof, rddx)))))
+@assert allequal("bffbx", transform(brof, transform(frob, transform(frob, transform(brof, rddx)))))
+partids = collect(transform((partid, iter) -> [partid], rddx))
+@assert length(partids) == 2
+@assert Set([0, 1]) == Set(partids)
+
 println("Spock: all tests passed")