PySpark Helper Function - perform reduceByKey on a dataframe

df_reduce_by_key.py

# Removes a lot of the boiler plate involved in converting a pyspark dataframe
# to and from an rdd, in order to do a reduceByKey operation. 
#
# Lifted from:
# - http://codereview.stackexchange.com/questions/115082/generic-reduceby-or-groupby-aggregate-functionality-with-spark-dataframe


from pyspark.sql import Row
from pyspark.sql.functions import struct
from pyspark.sql import DataFrame
from collections import OrderedDict

def reduce_by(self, by, cols, f, schema=None):
    """
    :param self DataFrame
    :param by a list of grouping columns 
    :param cols a list of columns to aggregate
    :param aggregation function Row => Row
    :return DataFrame
    """
    def merge_kv(kv):
        key, value = kv
        return Row(**OrderedDict(zip(
            key.__fields__ + value.__fields__, key + value)
        ))

    return (self
        .select(struct(*by), struct(*cols))
        .rdd
        .reduceByKey(f)
        .map(merge_kv)
        .toDF(schema))

DataFrame.reduce_by = reduce_by  # A quick monkey patch