CitationAttacherWithMatchedLimiter implements Serializable { public JavaPairRDD<MatchableEntity, MatchableEntity> attachCitationsAndLimitDocs(JavaPairRDD<String, MatchableEntity> citIdDocPairs, JavaPairRDD<String, MatchableEntity> citations) { JavaPairRDD<MatchableEntity, MatchableEntity> citDocPairs = citIdDocPairs .groupByKey() .join(citations) .flatMapToPair(docsWithCit -> { MatchableEntity citation = docsWithCit._2._2; Collection<EntityWithSimilarity> docsFiltered = calculateSimilarityAndLimitDocs(citation, docsWithCit._2._1); return docsFiltered.stream() .map(x -> new Tuple2<MatchableEntity, MatchableEntity>(citation, x.getEntity())) .collect(Collectors.toList()); }); return citDocPairs; } CitationAttacherWithMatchedLimiter();  CitationAttacherWithMatchedLimiter(int sameCitationsLimit); JavaPairRDD<MatchableEntity, MatchableEntity> attachCitationsAndLimitDocs(JavaPairRDD<String, MatchableEntity> citIdDocPairs, JavaPairRDD<String, MatchableEntity> citations); int getSameCitationsLimit();  }
@Test public void attachCitationsAndLimitDocs() { JavaPairRDD<String, MatchableEntity> citIdDocPairs = sparkContext.parallelizePairs(ImmutableList.of( new Tuple2<>(citation1.id(), document1), new Tuple2<>(citation1.id(), document2), new Tuple2<>(citation2.id(), document3))); JavaPairRDD<String, MatchableEntity> citations = sparkContext.parallelizePairs(generateIdWithEntityTuples( Lists.newArrayList(citation1, citation2, citation3, citation4, citation5))); JavaPairRDD<MatchableEntity, MatchableEntity> actualCitDocPairs = citationAttacher.attachCitationsAndLimitDocs(citIdDocPairs, citations); List<Tuple2<MatchableEntity, MatchableEntity>> expectedCitDocPairs = ImmutableList.of( new Tuple2<>(citation1, document1), new Tuple2<>(citation1, document2), new Tuple2<>(citation2, document3)); assertCitDocPairsEquals(actualCitDocPairs.collect(), expectedCitDocPairs); }
@Test public void attachCitationsAndLimitDocs_LIMIT_REACHED() { JavaPairRDD<String, MatchableEntity> citIdDocPairs = sparkContext.parallelizePairs(ImmutableList.of( new Tuple2<>(citation1.id(), document1), new Tuple2<>(citation1.id(), document2), new Tuple2<>(citation1.id(), document3v2), new Tuple2<>(citation1.id(), document3), new Tuple2<>(citation1.id(), document4), new Tuple2<>(citation2.id(), document5))); JavaPairRDD<String, MatchableEntity> citations = sparkContext.parallelizePairs(generateIdWithEntityTuples( Lists.newArrayList(citation1, citation2, citation3, citation4, citation5))); JavaPairRDD<MatchableEntity, MatchableEntity> actualCitDocPairs = citationAttacher.attachCitationsAndLimitDocs(citIdDocPairs, citations); List<Tuple2<MatchableEntity, MatchableEntity>> expectedCitDocPairs = ImmutableList.of( new Tuple2<>(citation1, document1), new Tuple2<>(citation1, document2), new Tuple2<>(citation1, document3), new Tuple2<>(citation2, document5)); assertCitDocPairsEquals(actualCitDocPairs.collect(), expectedCitDocPairs); }
@Test public void attachCitationsAndLimitDocs_NO_CITATIONS() { JavaPairRDD<String, MatchableEntity> citIdDocPairs = sparkContext.parallelizePairs(ImmutableList.of( new Tuple2<>(citation1.id(), document1), new Tuple2<>(citation1.id(), document2), new Tuple2<>(citation2.id(), document4))); JavaPairRDD<String, MatchableEntity> citations = sparkContext.parallelizePairs(Lists.newArrayList()); JavaPairRDD<MatchableEntity, MatchableEntity> actualCitDocPairs = citationAttacher.attachCitationsAndLimitDocs(citIdDocPairs, citations); assertEquals(actualCitDocPairs.count(), 0); }