java to python, and python to java

ahirreddy · ahirreddy · commit 67ba875f262c · 2014-04-14T17:05:18.000-07:00
diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRDD.scala b/core/src/main/scala/org/apache/spark/api/python/PythonRDD.scala
@@ -289,8 +289,12 @@ private[spark] object PythonRDD {
   def pythonToJava(pyRDD: JavaRDD[Array[Byte]]): JavaRDD[_] = {
     pyRDD.rdd.mapPartitions { iter =>
       val unpickle = new Unpickler
-      iter.map { row =>
-        unpickle.loads(row)
+      iter.flatMap { row =>
+        unpickle.loads(row) match {
+          case objs: java.util.ArrayList[Any] => objs
+          // Incase the partition doesn't have a collection
+          case obj => Seq(obj)
+        }
       }
     }
   }
diff --git a/python/pyspark/context.py b/python/pyspark/context.py
@@ -174,6 +174,8 @@ def _ensure_initialized(cls, instance=None, gateway=None):
                 SparkContext._gateway = gateway or launch_gateway()
                 SparkContext._jvm = SparkContext._gateway.jvm
                 SparkContext._writeToFile = SparkContext._jvm.PythonRDD.writeToFile
+                SparkContext._pythonToJava = SparkContext._jvm.PythonRDD.pythonToJava
+                SparkContext._javaToPython = SparkContext._jvm.PythonRDD.javaToPython
 
             if instance:
                 if SparkContext._active_spark_context and SparkContext._active_spark_context != instance:

Original file line number	Diff line number	Diff line change
`@@ -289,8 +289,12 @@ private[spark] object PythonRDD {`
`289`	`289`	`def pythonToJava(pyRDD: JavaRDD[Array[Byte]]): JavaRDD[_] = {`
`290`	`290`	`pyRDD.rdd.mapPartitions { iter =>`
`291`	`291`	`val unpickle = new Unpickler`
`292`		`- iter.map { row =>`
`293`		`- unpickle.loads(row)`
	`292`	`+ iter.flatMap { row =>`
	`293`	`+ unpickle.loads(row) match {`
	`294`	`+ case objs: java.util.ArrayList[Any] => objs`
	`295`	`+ // Incase the partition doesn't have a collection`
	`296`	`+ case obj => Seq(obj)`
	`297`	`+ }`
`294`	`298`	`}`
`295`	`299`	`}`
`296`	`300`	`}`