improve aggregate

Author: cloud-fan

python/pyspark/rdd.py

@@ -2330,11 +2330,9 @@ def _prepare_for_python_RDD(sc, command, obj=None):
     return pickled_command, broadcast_vars, env, includes
 
 
-def _wrap_function(sc, func, deserializer=None, serializer=None, profiler=None):
-    if deserializer is None:
-        deserializer = AutoBatchedSerializer(PickleSerializer())
-    if serializer is None:
-        serializer = AutoBatchedSerializer(PickleSerializer())
+def _wrap_function(sc, func, deserializer, serializer, profiler=None):
+    assert deserializer, "deserializer should not be empty"
+    assert serializer, "serializer should not be empty"
     command = (func, profiler, deserializer, serializer)
     pickled_command, broadcast_vars, env, includes = _prepare_for_python_RDD(sc, command)
     return sc._jvm.PythonFunction(bytearray(pickled_command), env, includes, sc.pythonExec,

python/pyspark/sql/dataframe.py

@@ -28,7 +28,8 @@
 
 from pyspark import since
 from pyspark.rdd import RDD, _load_from_socket, ignore_unicode_prefix
-from pyspark.serializers import BatchedSerializer, PickleSerializer, UTF8Deserializer
+from pyspark.serializers import AutoBatchedSerializer, BatchedSerializer, PickleSerializer, \
+    UTF8Deserializer, PairDeserializer
 from pyspark.storagelevel import StorageLevel
 from pyspark.traceback_utils import SCCallSiteSync
 from pyspark.sql.types import _parse_datatype_json_string
@@ -236,9 +237,14 @@ def collect(self):
         >>> df.collect()
         [Row(age=2, name=u'Alice'), Row(age=5, name=u'Bob')]
         """
+
+        if self._jdf.isPickled():
+            deserializer = PickleSerializer()
+        else:
+            deserializer = BatchedSerializer(PickleSerializer())
         with SCCallSiteSync(self._sc) as css:
             port = self._jdf.collectToPython()
-        return list(_load_from_socket(port, BatchedSerializer(PickleSerializer())))
+        return list(_load_from_socket(port, deserializer))
 
     @ignore_unicode_prefix
     @since(1.3)
@@ -282,13 +288,16 @@ def map(self, f):
     @since(2.0)
     def applySchema(self, schema=None):
         """ TODO """
-        if schema is None:
-            from pyspark.sql.types import _infer_type, _merge_type
-            # If no schema is specified, infer it from the whole data set.
-            jrdd = self._prev_jdf.javaToPython()
-            rdd = RDD(jrdd, self._sc, BatchedSerializer(PickleSerializer()))
-            schema = rdd.mapPartitions(self._func).map(_infer_type).reduce(_merge_type)
-        return PipelinedDataFrame(self, output_schema=schema)
+        if isinstance(self, PipelinedDataFrame):
+            if schema is None:
+                from pyspark.sql.types import _infer_type, _merge_type
+                # If no schema is specified, infer it from the whole data set.
+                jrdd = self._prev_jdf.javaToPython()
+                rdd = RDD(jrdd, self._sc, BatchedSerializer(PickleSerializer()))
+                schema = rdd.mapPartitions(self._func).map(_infer_type).reduce(_merge_type)
+            return PipelinedDataFrame(self, output_schema=schema)
+        else:
+            return self
 
     @ignore_unicode_prefix
     @since(2.0)
@@ -926,7 +935,7 @@ def groupByKey(self, key_func, key_type):
         wraped_func = _wrap_func(self._sc, self._jdf, f, False)
         jgd = self._jdf.pythonGroupBy(wraped_func, key_type.json())
         from pyspark.sql.group import GroupedData
-        return GroupedData(jgd, self.sql_ctx, key_func)
+        return GroupedData(jgd, self.sql_ctx, not isinstance(key_type, StructType))
 
     @since(1.4)
     def rollup(self, *cols):
@@ -1396,6 +1405,7 @@ def __init__(self, prev, func=None, output_schema=None):
         from pyspark.sql.group import GroupedData
 
         if output_schema is None:
+            # should get it from java side
             self._schema = StructType().add("binary", BinaryType(), False, {"pickled": True})
         else:
             self._schema = output_schema
@@ -1446,7 +1456,7 @@ def _jdf(self):
         return self._jdf_val
 
     def _create_jdf(self, func, schema=None):
-        wrapped_func = _wrap_func(self._sc, self._prev_jdf, func, schema is None)
+        wrapped_func = _wrap_func(self._sc, self._prev_jdf, func, schema is None, self._grouped)
         if schema is None:
             if self._grouped:
                 return self._prev_jdf.flatMapGroups(wrapped_func)
@@ -1460,16 +1470,18 @@ def _create_jdf(self, func, schema=None):
                 return self._prev_jdf.pythonMapPartitions(wrapped_func, schema_string)
 
 
-def _wrap_func(sc, jdf, func, output_binary):
-    if jdf.isPickled():
+def _wrap_func(sc, jdf, func, output_binary, input_grouped=False):
+    if input_grouped:
+        deserializer = PairDeserializer(PickleSerializer(), PickleSerializer())
+    elif jdf.isPickled():
         deserializer = PickleSerializer()
     else:
-        deserializer = None  # the framework will provide a default one
+        deserializer = AutoBatchedSerializer(PickleSerializer())
 
     if output_binary:
         serializer = PickleSerializer()
     else:
-        serializer = None  # the framework will provide a default one
+        serializer = AutoBatchedSerializer(PickleSerializer())
 
     from pyspark.rdd import _wrap_function
     return _wrap_function(sc, lambda _, iterator: func(iterator), deserializer, serializer)

python/pyspark/sql/group.py

@@ -15,6 +15,15 @@
 # limitations under the License.
 #
 
+import sys
+
+if sys.version >= '3':
+    basestring = unicode = str
+    long = int
+    from functools import reduce
+else:
+    from itertools import imap as map
+
 from pyspark import since
 from pyspark.rdd import ignore_unicode_prefix
 from pyspark.sql.column import Column, _to_seq, _to_java_column, _create_column_from_literal
@@ -54,25 +63,25 @@ class GroupedData(object):
     .. versionadded:: 1.3
     """
 
-    def __init__(self, jgd, sql_ctx, key_func=None):
+    def __init__(self, jgd, sql_ctx, flat_key=False):
         self._jgd = jgd
         self.sql_ctx = sql_ctx
-        if key_func is None:
-            self.key_func = lambda key: key
+        if flat_key:
+            self._key_converter = lambda key: key[0]
         else:
-            self.key_func = key_func
+            self._key_converter = lambda key: key
 
     @ignore_unicode_prefix
     @since(2.0)
     def flatMapGroups(self, func):
         """ TODO """
-        import itertools
-        key_func = self.key_func
+        key_converter = self._key_converter
 
-        def process(iterator):
-            first = iterator.next()
-            key = key_func(first)
-            return func(key, itertools.chain([first], iterator))
+        def process(inputs):
+            record_converter = lambda record: (key_converter(record[0]), record[1])
+            for key, values in GroupedIterator(map(record_converter, inputs)):
+                for output in func(key, values):
+                    yield output
 
         return PipelinedDataFrame(self, process)
 
@@ -217,6 +226,86 @@ def pivot(self, pivot_col, values=None):
         return GroupedData(jgd, self.sql_ctx)
 
 
+class GroupedIterator(object):
+    """ TODO """
+
+    def __init__(self, inputs):
+        self.inputs = BufferedIterator(inputs)
+        self.current_input = inputs.next()
+        self.current_key = self.current_input[0]
+        self.current_values = GroupValuesIterator(self)
+
+    def __iter__(self):
+        return self
+
+    def next(self):
+        if self.current_values is None:
+            self._fetch_next_group()
+
+        ret = (self.current_key, self.current_values)
+        self.current_values = None
+        return ret
+
+    def _fetch_next_group(self):
+        if self.current_input is None:
+            self.current_input = self.inputs.next()
+
+        # Skip to next group, or consume all inputs and throw StopIteration exception.
+        while self.current_input[0] == self.current_key:
+            self.current_input = self.inputs.next()
+
+        self.current_key = self.current_input[0]
+        self.current_values = GroupValuesIterator(self)
+
+
+class GroupValuesIterator(object):
+    """ TODO """
+
+    def __init__(self, outter):
+        self.outter = outter
+
+    def __iter__(self):
+        return self
+
+    def next(self):
+        if self.outter.current_input is None:
+            self._fetch_next_value()
+
+        value = self.outter.current_input[1]
+        self.outter.current_input = None
+        return value
+
+    def _fetch_next_value(self):
+        if self.outter.inputs.head()[0] == self.outter.current_key:
+            self.outter.current_input = self.outter.inputs.next()
+        else:
+            raise StopIteration
+
+
+class BufferedIterator(object):
+    """ TODO """
+
+    def __init__(self, iterator):
+        self.iterator = iterator
+        self.buffered = None
+
+    def __iter__(self):
+        return self
+
+    def next(self):
+        if self.buffered is None:
+            return self.iterator.next()
+        else:
+            item = self.buffered
+            self.buffered = None
+            return item
+
+    def head(self):
+        if self.buffered is None:
+            self.buffered = self.iterator.next()
+        return self.buffered
+
+
 def _test():
     import doctest
     from pyspark.context import SparkContext
@@ -237,13 +326,6 @@ def _test():
                                    Row(course="dotNET", year=2013, earnings=48000),
                                    Row(course="Java",   year=2013, earnings=30000)]).toDF()
 
-    ds = globs['sqlContext'].createDataFrame([(i, i) for i in range(100)], ("key", "value"))
-    grouped = ds.groupByKey(lambda row: row.key % 5, IntegerType())
-    value_sum = lambda rows: sum(map(lambda row: row.value, rows))
-    agged = grouped.mapGroups(lambda key, values: str(key) + ":" + str(value_sum(values)))
-    result = agged.applySchema(StringType()).collect()
-    raise ValueError(result[0][0])
-
     (failure_count, test_count) = doctest.testmod(
         pyspark.sql.group, globs=globs,
         optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF)

python/pyspark/sql/tests.py

@@ -1188,6 +1188,11 @@ def test_dataset(self):
         self.assertTrue(result[0][0] > 0)
         self.assertTrue(result[1][0] > 0)
 
+        # If no schema is given, collect will return custom objects instead of rows.
+        result = ds2.collect()
+        self.assertEqual(result[0], 0)
+        self.assertEqual(result[1], 3)
+
         # row count should be corrected even no schema is specified.
         self.assertEqual(ds2.count(), 100)
 
@@ -1198,6 +1203,40 @@ def test_dataset(self):
         self.assertEqual(result[0][0], 0)
         self.assertEqual(result[1][0], 3)
 
+    def test_typed_aggregate(self):
+        data = [(i, i * 2) for i in range(100)]
+        ds = self.sqlCtx.createDataFrame(data, ("key", "value"))
+        sum_tuple = lambda values: sum(map(lambda value: value[0] * value[1], values))
+
+        def get_python_result(data, key_func, agg_func):
+            data.sort(key=key_func)
+            expected_result = []
+            import itertools
+            for key, values in itertools.groupby(data, key_func):
+                expected_result.append(agg_func(key, values))
+            return expected_result
+
+        grouped = ds.groupByKey(lambda row: row.key % 5, IntegerType())
+        agg_func = lambda key, values: str(key) + ":" + str(sum_tuple(values))
+        result = sorted(grouped.mapGroups(agg_func).collect())
+        expected_result = get_python_result(data, lambda i: i[0] % 5, agg_func)
+        self.assertEqual(result, expected_result)
+
+        # We can also call groupByKey on a Dataset of custom objects.
+        ds2 = ds.map2(lambda row: row.key)
+        grouped = ds2.groupByKey(lambda i: i % 5, IntegerType())
+        agg_func = lambda key, values: str(key) + ":" + str(sum(values))
+        result = sorted(grouped.mapGroups(agg_func).collect())
+        expected_result = get_python_result(range(100), lambda i: i % 5, agg_func)
+        self.assertEqual(result, expected_result)
+
+        # We can also apply typed aggregate after structured groupBy, the key is row object.
+        grouped = ds.groupBy(ds.key % 2, ds.key % 3)
+        agg_func = lambda key, values: str(key[0]) + str(key[1]) + ":" + str(sum_tuple(values))
+        result = sorted(grouped.mapGroups(agg_func).collect())
+        expected_result = get_python_result(data, lambda i: (i[0] % 2, i[0] % 3), agg_func)
+        self.assertEqual(result, expected_result)
+
 
 class HiveContextSQLTests(ReusedPySparkTestCase):
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/package.scala

@@ -88,8 +88,6 @@ package object expressions  {
    */
   implicit class AttributeSeq(attrs: Seq[Attribute]) {
     /** Creates a StructType with a schema matching this `Seq[Attribute]`. */
-    def toStructType: StructType = {
-      StructType(attrs.map(a => StructField(a.name, a.dataType, a.nullable)))
-    }
+    def toStructType: StructType = StructType.fromAttributes(attrs)
   }
 }

sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala

@@ -1750,9 +1750,13 @@ class DataFrame private[sql](
    * Converts a JavaRDD to a PythonRDD.
    */
   protected[sql] def javaToPython: JavaRDD[Array[Byte]] = {
-    val structType = schema  // capture it for closure
-    val rdd = queryExecution.toRdd.map(EvaluatePython.toJava(_, structType))
-    EvaluatePython.javaToPython(rdd)
+    if (EvaluatePython.isPickled(schema)) {
+      queryExecution.toRdd.map(_.getBinary(0))
+    } else {
+      val structType = schema  // capture it for closure
+      val rdd = queryExecution.toRdd.map(EvaluatePython.toJava(_, structType))
+      EvaluatePython.javaToPython(rdd)
+    }
   }
 
   protected[sql] def collectToPython(): Int = {

sql/core/src/main/scala/org/apache/spark/sql/GroupedPythonDataset.scala

@@ -33,8 +33,7 @@ class GroupedPythonDataset private[sql](
 
   private def sqlContext = queryExecution.sqlContext
 
-  protected[sql] def isPickled(): Boolean =
-    EvaluatePython.isPickled(queryExecution.analyzed.output.toStructType)
+  protected[sql] def isPickled(): Boolean = EvaluatePython.isPickled(dataAttributes.toStructType)
 
   private def groupedData =
     new GroupedData(