choose batch size automatically

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

@@ -775,17 +775,38 @@ private[spark] object PythonRDD extends Logging {
     }.toJavaRDD()
   }
 
+  private class AutoBatchedPickler(iter: Iterator[Any]) extends Iterator[Array[Byte]] {
+    private val pickle = new Pickler()
+    private var batch = 1
+    private val buffer = new mutable.ArrayBuffer[Any]
+
+    override def hasNext(): Boolean = iter.hasNext
+
+    override def next(): Array[Byte] = {
+      while (iter.hasNext && buffer.length < batch) {
+        buffer += iter.next()
+      }
+      val bytes = pickle.dumps(buffer)
+      val size = bytes.length
+      // let  1M < size < 10M
+      if (size < 1024 * 100) {
+        batch = (1024 * 100) / size  // fast grow
+      } else if (size < 1024 * 1024) {
+        batch *= 2
+      } else if (size > 1024 * 1024 * 10) {
+        batch /= 2
+      }
+      buffer.clear()
+      bytes
+    }
+  }
+
   /**
    * Convert and RDD of Java objects to and RDD of serialized Python objects, that is usable by
    * PySpark.
    */
   def javaToPython(jRDD: JavaRDD[Any]): JavaRDD[Array[Byte]] = {
-    jRDD.rdd.mapPartitions { iter =>
-      val pickle = new Pickler
-      iter.grouped(1024).map { rows =>
-        pickle.dumps(rows.toArray)
-      }
-    }
+    jRDD.rdd.mapPartitions { iter => new AutoBatchedPickler(iter) }
   }
 
   /**

python/pyspark/rdd.py

@@ -34,7 +34,7 @@
 
 from pyspark.serializers import NoOpSerializer, CartesianDeserializer, \
     BatchedSerializer, CloudPickleSerializer, PairDeserializer, \
-    PickleSerializer, pack_long, CompressedSerializer
+    PickleSerializer, pack_long, AutoBatchedSerializer
 from pyspark.join import python_join, python_left_outer_join, \
     python_right_outer_join, python_cogroup
 from pyspark.statcounter import StatCounter
@@ -1927,10 +1927,8 @@ def _to_java_object_rdd(self):
         It will convert each Python object into Java object by Pyrolite, whenever the
         RDD is serialized in batch or not.
         """
-        if not self._is_pickled():
-            self = self._reserialize(BatchedSerializer(PickleSerializer(), 1024))
-        batched = isinstance(self._jrdd_deserializer, BatchedSerializer)
-        return self.ctx._jvm.PythonRDD.pythonToJava(self._jrdd, batched)
+        rdd = self._reserialize(AutoBatchedSerializer(PickleSerializer()))
+        return self.ctx._jvm.PythonRDD.pythonToJava(rdd._jrdd, True)
 
     def countApprox(self, timeout, confidence=0.95):
         """

python/pyspark/serializers.py

@@ -68,6 +68,7 @@
 import types
 import collections
 import zlib
+import itertools
 
 from pyspark import cloudpickle
 
@@ -211,6 +212,42 @@ def __str__(self):
         return "BatchedSerializer<%s>" % str(self.serializer)
 
 
+class AutoBatchedSerializer(BatchedSerializer):
+    """
+    Choose the size of batch automatically based on the size of object
+    """
+
+    def __init__(self, serializer, bestSize=1 << 20):
+        BatchedSerializer.__init__(self, serializer, -1)
+        self.bestSize = bestSize
+
+    def dump_stream(self, iterator, stream):
+        batch, best = 1, self.bestSize
+        iterator = iter(iterator)
+        while True:
+            vs = list(itertools.islice(iterator, batch))
+            if not vs:
+                break
+
+            bytes = self.serializer.dumps(vs)
+            size = len(bytes)
+            if size < best / 10:
+                batch = best / 10 / size
+            elif size < best:
+                batch *= 2
+            elif size > best * 10:
+                batch /= 2
+            write_int(len(bytes), stream)
+            stream.write(bytes)
+
+    def __eq__(self, other):
+        return (isinstance(other, AutoBatchedSerializer) and
+                other.serializer == self.serializer)
+
+    def __str__(self):
+        return "BatchedSerializer<%s>" % str(self.serializer)
+
+
 class CartesianDeserializer(FramedSerializer):
 
     """