Maybe I totally reinvented the wheel, or maybe I've invented something new and useful. Can one of you tell me if there's a better way of doing this? Here's what I'm trying to do:

I want a generic reduceBy function, that works like an RDD's reduceByKey, but will let me group data by any column in a Spark DataFrame. You may say that we already have that, and it's called groupBy, but as far as I can tell, groupBy only lets you aggregate using some very limited options. I want to groupBy, and then run an arbitrary function to aggregate. Has anyone already done that?

Basically, I'm taking a Spark DataFrame that looks like this:

| birthdate|favecolor| name|twitterhandle|facebookpage|           favesong|
|2000-01-01|     blue|Alice|     allyblue|        null|               null|
|1999-12-31|     null|  Bob|         null|      BobbyG| Gangsters Paradise|
|      null|     null|Alice|         null|        null|Rolling in the Deep|

and reducing by the column 'name' with a custom function to get this:

| birthdate|favecolor|           favesong| name|twitterhandle|facebookpage|
|2000-01-01|     blue|Rolling in the Deep|Alice|     allyblue|        null|
|1999-12-31|     null| Gangsters Paradise|  Bob|         null|      BobbyG|

I just noticed the change in column order. I think I can fix that pretty quickly by taking note of the schema before beginning. But anyway, I had to write a ton of code to get that to work, and this seems like such a simple operation somebody else should have done it by now.

Here's the code, written with Python 3.5.1 and Spark 1.5.2:

 def addEmptyColumns(df, colNames):

     :param df: 
     :param colNames: 
     exprs = df.columns + ["null as " + colName for colName in colNames]
     return df.selectExpr(*exprs)

 def concatTwoDfs(left, right):

     :param left: 
     :param right: 
     # append columns from right df to left df
     missingColumnsLeft = set(right.columns) - set(left.columns)
     left = addEmptyColumns(left, missingColumnsLeft)

     # append columns from left df to right df
     missingColumnsRight = set(left.columns) - set(right.columns)
     right = addEmptyColumns(right, missingColumnsRight)

     # let's set the same order of columns
     right = right[left.columns]

      # finally, union them
     return left.unionAll(right)

 def reduce(function, iterable, initializer=None):
     A copy of the rough code from Python 2's reduce function documentation.  Why did Python 3 get rid of it?

     Apply function of two arguments cumulatively to the items of iterable, from left to right, so as to reduce the
     iterable to a single value. For example, reduce(lambda x, y: x+y, [1, 2, 3, 4, 5]) calculates ((((1+2)+3)+4)+5).
     The left argument, x, is the accumulated value and the right argument, y, is the update value from the iterable.
     If the optional initializer is present, it is placed before the items of the iterable in the calculation, and
     serves as a default when the iterable is empty. If initializer is not given and iterable contains only one item,
     the first item is returned.

     :param function: use this function to reduce the elements of iterable
     :param iterable:
     :param initializer:
     it = iter(iterable)
     if initializer is None:
             initializer = next(it)
         except StopIteration:
             raise TypeError('reduce() of empty sequence with no initial value')
     accum_value = initializer
     for x in it:
         accum_value = function(accum_value, x)
     return accum_value

 def concat(dfs):
     Concatenates two Spark dataframes intelligently, adding missing columns with 'null' entry where appropriate.

     :param dfs: a list or tuple of two Spark dataframes
     :return: single dataframe consisting of dfs' columns and data
     return reduce(concatTwoDfs, dfs)

 def combine_rows(row1, row2):
     Takes two rows assumed to have the same columns, combines them, using values from row1 when available, from row2

     :param row1: pyspark.sql.Row
     :param row2: pyspark.sql.Row
     :return: pyspark.sql.Row combined from row1 and row2
     from pyspark.sql import Row
     combined = {}
     for col in row1.asDict():
         if row1.asDict()[col] is not None:
             combined[col] = row1.asDict()[col]
             combined[col] = row2.asDict()[col]
     return Row(**combined)

 def remove_nones(row):
     Takes in a row, returns that same row minus all of the columns that have a None entry.  This is required in
     order to create a new DataFrame using only this row; DataFrame will not be created if it doesn't know what kind
     of value to expect in a column.

     :param row:
     from pyspark.sql import Row
     cleaned = {}
     for col in row.asDict():
         if row.asDict()[col] is not None:
             cleaned[col] = row.asDict()[col]
     return Row(**cleaned)

 def reduce_by(df, col, func):
     Does pretty much the same thing as an RDD's reduceByKey, but much more generic.  Kind of like a Spark DataFrame's
     groupBy, but lets you aggregate by any generic function.

     :param df: the DataFrame to be reduced
     :param col: the column you want to use for grouping in df
     :param func: the function you will use to reduce df
     :return: a reduced DataFrame
     first_loop = True
     unique_entries = df.select(col).distinct().collect()
     return_rdd = sc.parallelize([])
     for entry in unique_entries:
         if first_loop:
             return_df = sqlContext.createDataFrame( \
                                 sc.parallelize([remove_nones(df.filter(df[col] == entry[0]).rdd.reduce(func))]))
             first_loop = False
             return_df = concat((return_df, \
                                sqlContext.createDataFrame( \
                                 sc.parallelize([remove_nones(df.filter(df[col] == entry[0]).rdd.reduce(func))]))))
     return return_df

And you kick it all off by making a DataFrame called test_df, and running this:

reduce_by(test_df, 'name', combine_rows).show()

The function combine_rows - which isn't a very good name, considering I'd like to use any number of functions to combine the rows - pretty much just copies whatever isn't "null" into the new row, whenever it has the choice.

It seems to work OK, but I'm concerned about a few things and I find it hard to believe that no one has done this before, and if someone has something that works better, I'd like to use that.

If my code ends up being the best option for what I'm trying to do, I'd still like to improve on it. My biggest concern is that in reduce_by(). I'm collecting and iterating, both of which I try to avoid whenever possible.


1 Answer 1


The biggest issue here is a following chunk of code:

unique_entries = df.select(col).distinct().collect()

which unfortunately renders your approach useless in a general case:

  • it assumes that the result can fit into driver memory which may or may not be true.
  • finding distinct elements is an expensive process in a distributed application.
  • collect has to transfer data to the driver and pass data to the local Python interpreter converting from internal representation to PythonRDD somewhere on the way.

All of that can work pretty well if number of unique keys is small but can become prohibitively expensive otherwise. One possible improvement is to use to toLocalIterator instead of collect. It is more expensive because it triggers multiple jobs but fetches only a single partition at the time.

Another problem I see is a subsequent loop:

  • depending on a distribution of the keys reduce part can result in suboptimal resource usage up to the point when execution becomes completely sequential.
  • once again it has to collect data to the driver with all the related issues
  • iterative union can generate long lineages. It makes failure recovery expensive and can simply fail due to stack overflow
  • using parallelize on the small datasets (like a single row) is far from optimal especially combined with iterative union. It will result in a large number of empty partitions and growing number of total partitions and can show a similar behavior to the one described in Spark iteration time increasing exponentially when using join
  • last but not least iterative createDataFrame without specifying schema requires expensive schema inference.

These issues could be partially addressed (assuming data fits in the memory) by using only a single SparkContext.union, or even better single parallelize, followed by single a createDataFrame.

I want to groupBy, and then run an arbitrary function to aggregate. Has anyone already done that?

Kind of. Since 1.5.0 Spark supports UDAFs (User Defined Aggregate Functions) which can be used to apply any commutative and associative function. These can defined only using Scala / Java but with some effort can be used from Python. See How to map Python with Scala or Java User Defined Functions?.

If you're not very fond of an idea of writing Scala code an alternative approach is use RDD methods like this:

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))

DataFrame.reduce_by = reduce_by  # A quick monkey patch

Which can be used as follows:

def foo(row1, row2):
    """ A dummy function
    >>> foo(Row(x=1, y=None), Row(x=None, y=2))
    Row(x=1, y=2)
    return Row(**OrderedDict(zip(
      row1.__fields__, (x if x else y for (x, y) in zip(row1, row2))

# Example data
df = sc.parallelize([
    ("a", None, 1), ("a", None, 2), ("a", 3, None),
    ("b", None, 2), ("b", None, None), ("c", 1, -1)
]).toDF(["k", "v1", "v2"])

df.reduce_by(by=["k"], cols=["v1", "v2"], f=foo).show()

## +---+----+---+
## |  k|  v1| v2|
## +---+----+---+
## |  a|   3|  1|
## |  c|   1| -1|
## |  b|null|  2|
## +---+----+---+

It still has to move data between JVM an Python but doesn't suffer from other issues.

Finally to answer a question from the comments

reduce function (...) Why did Python 3 get rid of it?

Because Guido van Rossum hates reduce :) To quote All Things Pythonic [1]:

This is actually the one I've always hated most, because, apart from a few examples involving + or *, almost every time I see a reduce() call with a non-trivial function argument, I need to grab pen and paper to diagram what's actually being fed into that function before I understand what the reduce() is supposed to do.

Actually it is still out there but it has been moved to functools. Personally I would recommend toolz instead which provides a comprehensive set of functional utilities and as a bonus can serve as compatibility layer between Python 2.6+ and 3.3+.

  1. Guido van van Rossum (2005, March 10). All Things Pythonic. The fate of reduce() in Python 3000. Retrieved from http://www.artima.com/weblogs/viewpost.jsp?thread=98196
  • \$\begingroup\$ Thanks very much for such a detailed response. I learned a lot, and your code is better than anything I would have come up with. In the reduce_by() function though, you have a couple of return statements. Does the second one do anything? Also, is there a good place to learn all of this stuff at once, or do I just need to keep coding and pick things up as I go? Thanks again for your help! \$\endgroup\$
    – Nathaniel
    Commented Dec 28, 2015 at 21:56
  • \$\begingroup\$ You're right. The second return was just an artifact of the previous version of the code. Regarding your second question you can pick some solid basics sources (Learning Spark is pretty good from what I've seen, I've enjoyed edX courses as well) but beyond that this is just an open-ended quest :) \$\endgroup\$
    – zero323
    Commented Dec 29, 2015 at 9:03

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.