BestMatchedCitationPicker implements Serializable { public JavaPairRDD<MatchableEntity, IdWithSimilarity> pickBest(JavaPairRDD<MatchableEntity, MatchableEntity> citDocPairs) { JavaPairRDD<MatchableEntity, IdWithSimilarity> bestMatchedCitations = citDocPairs .mapPartitionsToPair(citAndDocParitition -> { SimilarityMeasurer similarityMeasurer = new SimilarityMeasurer(SimilarityMeasurer.advancedFvBuilder()); return calculateSimilarityAndFilter(similarityMeasurer, citAndDocParitition); }) .reduceByKey((docWithSim1, docWithSim2) -> (docWithSim1.getSimilarity() > docWithSim2.getSimilarity()) ? docWithSim1 : docWithSim2); return bestMatchedCitations; }  JavaPairRDD<MatchableEntity, IdWithSimilarity> pickBest(JavaPairRDD<MatchableEntity, MatchableEntity> citDocPairs);  }
@Test public void pickBest() { JavaPairRDD<MatchableEntity, MatchableEntity> citDocPairs = sparkContext.parallelizePairs(ImmutableList.of( new Tuple2<>(citation1, document1), new Tuple2<>(citation1, document2), new Tuple2<>(citation1, document3), new Tuple2<>(citation2, document4), new Tuple2<>(citation4, document4), new Tuple2<>(citation4, document5))); JavaPairRDD<MatchableEntity, IdWithSimilarity> actualMatchedCitations = bestMatchedCitationPicker.pickBest(citDocPairs); List<Tuple2<MatchableEntity, IdWithSimilarity>> expectedMatchedCitations = ImmutableList.of( new Tuple2<>(citation1, new IdWithSimilarity(document2.id(), 0.748069)), new Tuple2<>(citation4, new IdWithSimilarity(document5.id(), 0.739753))); assertMatchedEquals(actualMatchedCitations.collect(), expectedMatchedCitations); }