com.github.chen0040.sparkml.recommender.LinkRecommender Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of spark-ml-recommender Show documentation
Recommender algorithms implemented in Java and for Spark
There is a newer version: 1.0.5
package com.github.chen0040.sparkml.recommender;


import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.Function2;
import scala.Tuple2;

import java.util.*;


/**
 * Created by xschen on 22/2/2016.
 */
public class LinkRecommender {

   public JavaRDD fitAndTransform(JavaRDD rdd) {


      JavaPairRDD> rdd1 = rdd.flatMapToPair( s -> {

         String person = s.getPerson();
         List friendList = s.getHisFriends();
         List>> result = new ArrayList<>();

         for (int i = 0; i < friendList.size(); ++i) {
            result.add(new Tuple2<>(person, new Tuple2<>(friendList.get(i), "*")));
         }

         for (int i = 0; i < friendList.size() - 1; ++i) {
            for (int j = i + 1; j < friendList.size(); ++j) {
               result.add(new Tuple2<>(friendList.get(i), new Tuple2<>(friendList.get(j), person)));
               result.add(new Tuple2<>(friendList.get(j), new Tuple2<>(friendList.get(i), person)));
            }
         }

         return result;
      });

      JavaPairRDD>> rdd2 = rdd1.combineByKey((Function, Map>>) s1 -> {
         String toUser = s1._1();
         String commonFriend = s1._2();
         Map> map = new HashMap<>();

         if (commonFriend.equals("*")) {
            map.put(toUser, null);
         }
         else {
            Set set = new HashSet<>();
            set.add(commonFriend);
            map.put(toUser, set);
         }
         return map;
      }, (Function2>, Tuple2, Map>>) (map1, s2) -> {
         String toUser = s2._1();
         String commonFriend = s2._2();

         boolean alreadyFriend = commonFriend.equals("*");

         if (map1.containsKey(toUser)) {
            if (alreadyFriend) {
               map1.put(toUser, null);
            }
            else {
               if (map1.get(toUser) != null) {
                  map1.get(toUser).add(commonFriend);
               }
            }
         }
         else {
            if (alreadyFriend) {
               map1.put(toUser, null);
            }
            else {
               Set set = new HashSet<>();
               set.add(commonFriend);
               map1.put(toUser, set);
            }
         }

         return map1;

      }, (Function2>, Map>, Map>>) (map1, map2) -> {
         for (Map.Entry> entry2 : map2.entrySet()) {
            String toUser = entry2.getKey();
            Set set2 = entry2.getValue();

            if (set2 == null) {
               map1.put(toUser, null);
            }
            else {
               if (map1.containsKey(toUser)) {
                  if (map1.get(toUser) != null) {
                     map1.get(toUser).addAll(set2);
                  }
               }
               else {
                  map1.put(toUser, set2);
               }
            }
         }
         return map1;
      });

      return rdd2.flatMap(tuple2 -> {
         String person1 = tuple2._1();
         Map> links = tuple2._2();

         List result = new ArrayList<>();
         for(Map.Entry> entry : links.entrySet()){
            String person2 = entry.getKey();
            Set commonFriends = entry.getValue();
            FriendLinkRecommendation link = new FriendLinkRecommendation();
            link.setPerson1(person1);
            link.setPerson2(person2);
            link.setCommonFriends(commonFriends);
            result.add(link);
         }
         return result;
      });
   }
}