refactor, address comments

Author: davies

core/src/main/scala/org/apache/spark/api/python/PythonRDD.scala

@@ -705,25 +705,25 @@ private[spark] object PythonRDD extends Logging {
    * Convert an RDD of serialized Python tuple to Array (no recursive conversions).
    * It is only used by pyspark.sql.
    */
-  def pythonToJava(pyRDD: JavaRDD[Array[Byte]]): JavaRDD[Array[_]] = {
+  def pythonToJavaArray(pyRDD: JavaRDD[Array[Byte]], batched: Boolean): JavaRDD[Array[_]] = {
+
+    def toArray(obj: Any): Array[_] = {
+      obj match {
+        case objs: JArrayList[_] =>
+          objs.toArray
+        case obj if obj.getClass.isArray =>
+          obj.asInstanceOf[Array[_]].toArray
+      }
+    }
+
     pyRDD.rdd.mapPartitions { iter =>
       val unpickle = new Unpickler
       iter.flatMap { row =>
-        unpickle.loads(row) match {
-          // in case of objects are pickled in batch mode
-          case objs: JArrayList[_] => Try(objs.map(obj => obj match {
-            case list: JArrayList[_] => list.toArray // list
-            case obj if obj.getClass.isArray => // tuple
-              obj.asInstanceOf[Array[_]].toArray
-          })) match {
-            // objs is list of list or tuple
-            case Success(v) => v
-            // objs is a row, list of different objects
-            case Failure(e) => Seq(objs.toArray)
-          }
-          // not in batch mode
-          case obj if obj.getClass.isArray => // tuple
-            Seq(obj.asInstanceOf[Array[_]].toArray)
+        val obj = unpickle.loads(row)
+        if (batched) {
+          obj.asInstanceOf[JArrayList[_]].map(toArray)
+        } else {
+          Seq(toArray(obj))
         }
       }
     }.toJavaRDD()

python/pyspark/sql.py

@@ -640,7 +640,7 @@ def __init__(self, sparkContext, sqlContext=None):
         self._sc = sparkContext
         self._jsc = self._sc._jsc
         self._jvm = self._sc._jvm
-        self._pythonToJava = self._jvm.PythonRDD.pythonToJava
+        self._pythonToJava = self._jvm.PythonRDD.pythonToJavaArray
 
         if sqlContext:
             self._scala_SQLContext = sqlContext
@@ -686,10 +686,7 @@ def inferSchema(self, rdd):
 
         schema = _inferSchema(first)
         rdd = rdd.mapPartitions(lambda rows: _dropSchema(rows, schema))
-
-        jrdd = self._pythonToJava(rdd._jrdd)
-        srdd = self._ssql_ctx.applySchemaToPythonRDD(jrdd.rdd(), str(schema))
-        return SchemaRDD(srdd, self)
+        return self.applySchema(rdd, schema)
 
     def applySchema(self, rdd, schema):
         """Applies the given schema to the given RDD of L{dict}s.
@@ -719,7 +716,8 @@ def applySchema(self, rdd, schema):
         >>> srdd.collect()[0]
         (127, -32768, 1.0, datetime.datetime(2010, 1, 1, 1, 1, 1), 1, 2, [1, 2, 3], None)
         """
-        jrdd = self._pythonToJava(rdd._jrdd)
+        batched = isinstance(rdd._jrdd_deserializer, BatchedSerializer)
+        jrdd = self._pythonToJava(rdd._jrdd, batched)
         srdd = self._ssql_ctx.applySchemaToPythonRDD(jrdd.rdd(), str(schema))
         return SchemaRDD(srdd, self)