java to python, and python to java

pdeyhim · Apr 15, 2014 · 67ba875 · 67ba875
1 parent bcc0f23
commit 67ba875
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
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: