Skip to content

Instantly share code, notes, and snippets.

What would you like to do?
A naive approximation of k-nn algorithm (k-nearest neighbors) in pyspark. Approximation quality can be controlled by number of repartitions and number of repartition
from __future__ import print_function
import sys
from math import sqrt
import argparse
from collections import defaultdict
from random import randint
from pyspark import SparkContext
def parseArguments(args):
parser = argparse.ArgumentParser(add_help=True,
description='Item similarity - " + \
"naive k-nn approximation')
parser.add_argument('--debug', action="store_true",
dest="debug", default=False,
help="Debug flag",
parser.add_argument('--num_reco', action="store",
dest="num_reco", default=10,
help="Number of recommendations per user", type=int,
parser.add_argument('--repartition', action="store",
dest="repartition", default=10,
help="Number of repartition", type=int,
parser.add_argument('--repetition', action="store",
dest="repetition", default=10,
help="Number of repetition", type=int,
parser.add_argument('--app', action="store",
help="App name (spark context)",
required=False, default="rec_streaming")
parser.add_argument('--input', action="store",
help="Input location",
arguments = parser.parse_args()
return arguments
def load_item(item):
item_data = item.split(",")
item_id = item_data[0]
item_vec = [float(s) for s in item_data[1:]]
item_len = sqrt(sum([a*a for a in item_vec]))
return (item_id, (item_vec, item_len))
def distance(item1, item2):
item1_id, (item1_ftrs, item1_len) = item1
item2_id, (item2_ftrs, item2_len) = item2
distance = sum([a * b for a, b in zip(item1_ftrs, item2_ftrs)])
return distance/(item1_len * item2_len)
def approx_distance_user(user_iter, num_reco):
users = list(user_iter)
users_dict = defaultdict(dict)
for idx1, user1 in enumerate(users):
for user2 in users[idx1+1:]:
dist = distance(user1, user2)
users_dict[user1[0]][user2[0]] = dist
users_dict[user2[0]][user1[0]] = dist
users_dict = {k : sorted(v.items(), key = lambda a: -a[1])[:num_reco]
for k,v in users_dict.iteritems()}
return users_dict.items()
def main(args):
arguments = parseArguments(args)
if arguments.debug:
sys.stdout.write("Arguments: %s\n"%arguments)
sc = SparkContext(
items = sc.textFile(arguments.input)
items = x: load_item(x)).cache()
distances = sc.emptyRDD()
for i in xrange(0, arguments.repetition):
curr_distances = items.repartitionAndSortWithinPartitions(\
arguments.repartition, \
lambda x: randint(1, arguments.repartition)).\
mapPartitions(lambda x:
approx_distance_user(x, arguments.num_reco))
distances = distances.union(curr_distances)
distances = distances.reduceByKey(lambda a, b : a + b)
distances = x: (x[0],
sorted(set(x[1]), key = lambda a: -a[1])))
distances = distances.collect()
for item_id, rec in distances:
print ("%s\t%s"%(item_id, rec))
if __name__ == "__main__":
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment