browskie browskie - 1 year ago 73
Python Question

Getting Top N items per group in pySpark

I am using Spark 1.6.2, I have the following data structure:

sample = sqlContext.createDataFrame([
(2,['vodka','beer','juice', 'vinegar'])


I would like to extract top N most frequent terms per cat. I have developed the following solution which seems to work, however I wanted to see if there is a better way to do this.

from collections import Counter
def get_top(it, terms=200):
c = Counter(it.__iter__())
return [x[0][1] for x in c.most_common(terms)]

('cat',sf.explode('terms')) x: (, x.col))
.groupBy(lambda x: x[0])
.map(lambda x: (x[0], get_top(x[1], 2)))

It provides the following output:

[(1, ['orange', 'potato']), (2, ['vodka', 'beer'])]

Which is in line with what I am looking for, but I really don't like the fact that I am resorting to using Counter. How can I do it with spark alone?


Answer Source

If this is working it is probably better to post this to Code Review.

Just as an exercise I did this without the Counter but largely you are just replicating the same functionality.

  • Count each occurrence of (cat, term)
  • Group by cat
  • Sort the values based on Count and slice to number of terms (2)


from operator import add

('cat', sf.explode('terms'))
 .map(lambda x: (x, 1))
 .groupBy(lambda x: x[0][0])
 .mapValues(lambda x: [r[1] for r, _ in sorted(x, key=lambda a: -a[1])[:2]])


[(1, ['orange', 'potato']), (2, ['vodka', 'beer'])]