Mirrors/the-algorithm
2
0
Fork 0
mirror of https://github.com/twitter/the-algorithm.git synced 2025-01-05 00:51:55 +01:00
Code Issues Packages Projects Releases Wiki Activity

Simplify getDocIdSetIterator

Browse Source 
* Rename tweetIDs to distinctTweetIDs
* Use Java Stream
This commit is contained in:
Shail Shah 2023-04-01 13:25:22 -07:00
parent ec83d01dca
commit 5ecf79363a
1 changed files with 15 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
src/java/com/twitter/search/earlybird/index/TweetIDQuery.java
View File

@ -3,6 +3,7 @@ package com.twitter.search.earlybird.index;
import java.io.IOException; import java.io.IOException;
import java.util.Arrays; import java.util.Arrays;
import java.util.Set; import java.util.Set;
import java.util.stream.Collectors;
import com.google.common.collect.Sets; import com.google.common.collect.Sets;
@ -20,11 +21,11 @@ import com.twitter.search.core.earlybird.index.EarlybirdIndexSegmentAtomicReader
import com.twitter.search.core.earlybird.index.EarlybirdIndexSegmentData; import com.twitter.search.core.earlybird.index.EarlybirdIndexSegmentData;
public class TweetIDQuery extends Query { public class TweetIDQuery extends Query {
private final Set<Long> tweetIDs = Sets.newHashSet(); private final Set<Long> distinctTweetIDs = Sets.newHashSet();
public TweetIDQuery(long... tweetIDs) { public TweetIDQuery(long... tweetIDs) {
for (long tweetID : tweetIDs) { for (long tweetID : tweetIDs) {
this.tweetIDs.add(tweetID); this.distinctTweetIDs.add(tweetID);
} }
} }
@ -37,32 +38,23 @@ public class TweetIDQuery extends Query {
((EarlybirdIndexSegmentAtomicReader) context.reader()).getSegmentData(); ((EarlybirdIndexSegmentAtomicReader) context.reader()).getSegmentData();
DocIDToTweetIDMapper docIdToTweetIdMapper = segmentData.getDocIDToTweetIDMapper(); DocIDToTweetIDMapper docIdToTweetIdMapper = segmentData.getDocIDToTweetIDMapper();
Set<Integer> set = Sets.newHashSet(); int[] distinctSortedDocIDs = distinctTweetIDs.stream()
for (long tweetID : tweetIDs) { .map(tweetID -> docIdToTweetIdMapper.getDocID(tweetID))
int docID = docIdToTweetIdMapper.getDocID(tweetID); .filter(docID -> docID != DocIDToTweetIDMapper.ID_NOT_FOUND)
if (docID != DocIDToTweetIDMapper.ID_NOT_FOUND) { .sorted()
set.add(docID); .mapToInt(Integer::intValue)
} .toArray();
}
if (set.isEmpty()) { return distinctSortedDocIDs.length == 0 ?
return DocIdSetIterator.empty(); DocIdSetIterator.empty() :
} new IntArrayDocIdSetIterator(distinctSortedDocIDs);
int[] docIDs = new int[set.size()];
int i = 0;
for (int docID : set) {
docIDs[i++] = docID;
}
Arrays.sort(docIDs);
return new IntArrayDocIdSetIterator(docIDs);
} }
}; };
} }
@Override @Override
public int hashCode() { public int hashCode() {
return tweetIDs.hashCode(); return distinctTweetIDs.hashCode();
} }
@Override @Override
@ -71,11 +63,11 @@ public class TweetIDQuery extends Query {
return false; return false;
} }
return tweetIDs.equals(TweetIDQuery.class.cast(obj).tweetIDs); return distinctTweetIDs.equals(TweetIDQuery.class.cast(obj).distinctTweetIDs);
} }
@Override @Override
public String toString(String field) { public String toString(String field) {
return "TWEET_ID_QUERY: " + tweetIDs; return "TWEET_ID_QUERY: " + distinctTweetIDs;
} }
} }