Search in sources :

Example 11 with LongOpenHashSet

use of it.unimi.dsi.fastutil.longs.LongOpenHashSet in project Anserini by castorini.

the class IndexTweets method main.

@SuppressWarnings("static-access")
public static void main(String[] args) throws Exception {
    Options options = new Options();
    options.addOption(new Option(HELP_OPTION, "show help"));
    options.addOption(new Option(OPTIMIZE_OPTION, "merge indexes into a single segment"));
    options.addOption(new Option(STORE_TERM_VECTORS_OPTION, "store term vectors"));
    options.addOption(OptionBuilder.withArgName("dir").hasArg().withDescription("source collection directory").create(COLLECTION_OPTION));
    options.addOption(OptionBuilder.withArgName("dir").hasArg().withDescription("index location").create(INDEX_OPTION));
    options.addOption(OptionBuilder.withArgName("file").hasArg().withDescription("file with deleted tweetids").create(DELETES_OPTION));
    options.addOption(OptionBuilder.withArgName("id").hasArg().withDescription("max id").create(MAX_ID_OPTION));
    CommandLine cmdline = null;
    CommandLineParser parser = new GnuParser();
    try {
        cmdline = parser.parse(options, args);
    } catch (ParseException exp) {
        System.err.println("Error parsing command line: " + exp.getMessage());
        System.exit(-1);
    }
    if (cmdline.hasOption(HELP_OPTION) || !cmdline.hasOption(COLLECTION_OPTION) || !cmdline.hasOption(INDEX_OPTION)) {
        HelpFormatter formatter = new HelpFormatter();
        formatter.printHelp(IndexTweets.class.getName(), options);
        System.exit(-1);
    }
    String collectionPath = cmdline.getOptionValue(COLLECTION_OPTION);
    String indexPath = cmdline.getOptionValue(INDEX_OPTION);
    final FieldType textOptions = new FieldType();
    textOptions.setIndexOptions(IndexOptions.DOCS_AND_FREQS_AND_POSITIONS);
    textOptions.setStored(true);
    textOptions.setTokenized(true);
    if (cmdline.hasOption(STORE_TERM_VECTORS_OPTION)) {
        textOptions.setStoreTermVectors(true);
    }
    LOG.info("collection: " + collectionPath);
    LOG.info("index: " + indexPath);
    LongOpenHashSet deletes = null;
    if (cmdline.hasOption(DELETES_OPTION)) {
        deletes = new LongOpenHashSet();
        File deletesFile = new File(cmdline.getOptionValue(DELETES_OPTION));
        if (!deletesFile.exists()) {
            System.err.println("Error: " + deletesFile + " does not exist!");
            System.exit(-1);
        }
        LOG.info("Reading deletes from " + deletesFile);
        FileInputStream fin = new FileInputStream(deletesFile);
        byte[] ignoreBytes = new byte[2];
        // "B", "Z" bytes from commandline tools
        fin.read(ignoreBytes);
        BufferedReader br = new BufferedReader(new InputStreamReader(new CBZip2InputStream(fin)));
        String s;
        while ((s = br.readLine()) != null) {
            if (s.contains("\t")) {
                deletes.add(Long.parseLong(s.split("\t")[0]));
            } else {
                deletes.add(Long.parseLong(s));
            }
        }
        br.close();
        fin.close();
        LOG.info("Read " + deletes.size() + " tweetids from deletes file.");
    }
    long maxId = Long.MAX_VALUE;
    if (cmdline.hasOption(MAX_ID_OPTION)) {
        maxId = Long.parseLong(cmdline.getOptionValue(MAX_ID_OPTION));
        LOG.info("index: " + maxId);
    }
    long startTime = System.currentTimeMillis();
    File file = new File(collectionPath);
    if (!file.exists()) {
        System.err.println("Error: " + file + " does not exist!");
        System.exit(-1);
    }
    StatusStream stream = new JsonStatusCorpusReader(file);
    Directory dir = FSDirectory.open(Paths.get(indexPath));
    final IndexWriterConfig config = new IndexWriterConfig(ANALYZER);
    config.setOpenMode(IndexWriterConfig.OpenMode.CREATE);
    IndexWriter writer = new IndexWriter(dir, config);
    int cnt = 0;
    Status status;
    try {
        while ((status = stream.next()) != null) {
            if (status.getText() == null) {
                continue;
            }
            // Skip deletes tweetids.
            if (deletes != null && deletes.contains(status.getId())) {
                continue;
            }
            if (status.getId() > maxId) {
                continue;
            }
            cnt++;
            Document doc = new Document();
            doc.add(new LongPoint(StatusField.ID.name, status.getId()));
            doc.add(new StoredField(StatusField.ID.name, status.getId()));
            doc.add(new LongPoint(StatusField.EPOCH.name, status.getEpoch()));
            doc.add(new StoredField(StatusField.EPOCH.name, status.getEpoch()));
            doc.add(new TextField(StatusField.SCREEN_NAME.name, status.getScreenname(), Store.YES));
            doc.add(new Field(StatusField.TEXT.name, status.getText(), textOptions));
            doc.add(new IntPoint(StatusField.FRIENDS_COUNT.name, status.getFollowersCount()));
            doc.add(new StoredField(StatusField.FRIENDS_COUNT.name, status.getFollowersCount()));
            doc.add(new IntPoint(StatusField.FOLLOWERS_COUNT.name, status.getFriendsCount()));
            doc.add(new StoredField(StatusField.FOLLOWERS_COUNT.name, status.getFriendsCount()));
            doc.add(new IntPoint(StatusField.STATUSES_COUNT.name, status.getStatusesCount()));
            doc.add(new StoredField(StatusField.STATUSES_COUNT.name, status.getStatusesCount()));
            long inReplyToStatusId = status.getInReplyToStatusId();
            if (inReplyToStatusId > 0) {
                doc.add(new LongPoint(StatusField.IN_REPLY_TO_STATUS_ID.name, inReplyToStatusId));
                doc.add(new StoredField(StatusField.IN_REPLY_TO_STATUS_ID.name, inReplyToStatusId));
                doc.add(new LongPoint(StatusField.IN_REPLY_TO_USER_ID.name, status.getInReplyToUserId()));
                doc.add(new StoredField(StatusField.IN_REPLY_TO_USER_ID.name, status.getInReplyToUserId()));
            }
            String lang = status.getLang();
            if (!lang.equals("unknown")) {
                doc.add(new TextField(StatusField.LANG.name, status.getLang(), Store.YES));
            }
            long retweetStatusId = status.getRetweetedStatusId();
            if (retweetStatusId > 0) {
                doc.add(new LongPoint(StatusField.RETWEETED_STATUS_ID.name, retweetStatusId));
                doc.add(new StoredField(StatusField.RETWEETED_STATUS_ID.name, retweetStatusId));
                doc.add(new LongPoint(StatusField.RETWEETED_USER_ID.name, status.getRetweetedUserId()));
                doc.add(new StoredField(StatusField.RETWEETED_USER_ID.name, status.getRetweetedUserId()));
                doc.add(new IntPoint(StatusField.RETWEET_COUNT.name, status.getRetweetCount()));
                doc.add(new StoredField(StatusField.RETWEET_COUNT.name, status.getRetweetCount()));
                if (status.getRetweetCount() < 0 || status.getRetweetedStatusId() < 0) {
                    LOG.warn("Error parsing retweet fields of " + status.getId());
                }
            }
            writer.addDocument(doc);
            if (cnt % 100000 == 0) {
                LOG.info(cnt + " statuses indexed");
            }
        }
        LOG.info(String.format("Total of %s statuses added", cnt));
        if (cmdline.hasOption(OPTIMIZE_OPTION)) {
            LOG.info("Merging segments...");
            writer.forceMerge(1);
            LOG.info("Done!");
        }
        LOG.info("Total elapsed time: " + (System.currentTimeMillis() - startTime) + "ms");
    } catch (Exception e) {
        e.printStackTrace();
    } finally {
        writer.close();
        dir.close();
        stream.close();
    }
}
Also used : IndexOptions(org.apache.lucene.index.IndexOptions) CBZip2InputStream(org.apache.tools.bzip2.CBZip2InputStream) StatusStream(io.anserini.document.twitter.StatusStream) JsonStatusCorpusReader(io.anserini.document.twitter.JsonStatusCorpusReader) LongOpenHashSet(it.unimi.dsi.fastutil.longs.LongOpenHashSet) Directory(org.apache.lucene.store.Directory) FSDirectory(org.apache.lucene.store.FSDirectory) Status(io.anserini.document.twitter.Status) InputStreamReader(java.io.InputStreamReader) FileInputStream(java.io.FileInputStream) IndexWriter(org.apache.lucene.index.IndexWriter) BufferedReader(java.io.BufferedReader) File(java.io.File) IndexWriterConfig(org.apache.lucene.index.IndexWriterConfig)

Example 12 with LongOpenHashSet

use of it.unimi.dsi.fastutil.longs.LongOpenHashSet in project FredBoat by Frederikam.

the class JDAUtil method countUniqueUsers.

/**
 * A count of unique users over the provided shards. This is an expensive operation given FredBoats scale.
 * <p>
 * Optionally pass in a value of value of previous counts / expected size to that we can initialize the set used
 * to count the unique values with an approriate size reducing expensive resizing operations.
 */
@CheckReturnValue
public static int countUniqueUsers(@Nonnull Collection<JDA> shards, @Nullable AtomicInteger expectedUserCount) {
    if (shards.size() == 1) {
        // a single shard provides a cheap call for getting user cardinality
        return Math.toIntExact(shards.iterator().next().getUserCache().size());
    }
    int expected = expectedUserCount != null && expectedUserCount.get() > 0 ? expectedUserCount.get() : LongOpenHashSet.DEFAULT_INITIAL_SIZE;
    // add 10k for good measure
    LongOpenHashSet uniqueUsers = new LongOpenHashSet(expected + 10000);
    TObjectProcedure<User> adder = user -> {
        uniqueUsers.add(user.getIdLong());
        return true;
    };
    Collections.unmodifiableCollection(shards).forEach(// this means however, that for the (small) duration, the map cannot be used by other threads (if there are any)
    shard -> ((JDAImpl) shard).getUserMap().forEachValue(adder));
    return uniqueUsers.size();
}
Also used : CheckReturnValue(javax.annotation.CheckReturnValue) LongOpenHashSet(it.unimi.dsi.fastutil.longs.LongOpenHashSet) JDAImpl(net.dv8tion.jda.core.entities.impl.JDAImpl) User(net.dv8tion.jda.core.entities.User) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) TObjectProcedure(gnu.trove.procedure.TObjectProcedure) Collection(java.util.Collection) JDA(net.dv8tion.jda.core.JDA) Nonnull(javax.annotation.Nonnull) Collections(java.util.Collections) Nullable(javax.annotation.Nullable) User(net.dv8tion.jda.core.entities.User) JDAImpl(net.dv8tion.jda.core.entities.impl.JDAImpl) LongOpenHashSet(it.unimi.dsi.fastutil.longs.LongOpenHashSet) CheckReturnValue(javax.annotation.CheckReturnValue)

Example 13 with LongOpenHashSet

use of it.unimi.dsi.fastutil.longs.LongOpenHashSet in project presto by prestodb.

the class ArrayDistinctFunction method bigintDistinct.

@SqlType("array(bigint)")
public static Block bigintDistinct(@SqlType("array(bigint)") Block array) {
    final int arrayLength = array.getPositionCount();
    if (arrayLength < 2) {
        return array;
    }
    BlockBuilder distinctElementBlockBuilder;
    LongSet set = new LongOpenHashSet(arrayLength);
    if (array.mayHaveNull()) {
        int position = 0;
        boolean containsNull = false;
        // Keep adding the element to the set as long as there are no dupes.
        while (position < arrayLength) {
            if (array.isNull(position)) {
                if (!containsNull) {
                    containsNull = true;
                } else {
                    // Second null.
                    break;
                }
            } else if (!set.add(BIGINT.getLong(array, position))) {
                // Dupe found.
                break;
            }
            position++;
        }
        if (position == arrayLength) {
            // All elements are distinct, so just return the original.
            return array;
        }
        distinctElementBlockBuilder = BIGINT.createBlockBuilder(null, arrayLength);
        for (int i = 0; i < position; i++) {
            BIGINT.appendTo(array, i, distinctElementBlockBuilder);
        }
        for (position++; position < arrayLength; position++) {
            if (array.isNull(position)) {
                if (!containsNull) {
                    BIGINT.appendTo(array, position, distinctElementBlockBuilder);
                }
            } else if (set.add(BIGINT.getLong(array, position))) {
                BIGINT.appendTo(array, position, distinctElementBlockBuilder);
            }
        }
    } else {
        int position = 0;
        // Keep adding the element to the set as long as there are no dupes.
        while (position < arrayLength && set.add(BIGINT.getLong(array, position))) {
            position++;
        }
        if (position == arrayLength) {
            // All elements are distinct, so just return the original.
            return array;
        }
        distinctElementBlockBuilder = BIGINT.createBlockBuilder(null, arrayLength);
        for (int i = 0; i < position; i++) {
            BIGINT.appendTo(array, i, distinctElementBlockBuilder);
        }
        for (position++; position < arrayLength; position++) {
            if (set.add(BIGINT.getLong(array, position))) {
                BIGINT.appendTo(array, position, distinctElementBlockBuilder);
            }
        }
    }
    return distinctElementBlockBuilder.build();
}
Also used : LongSet(it.unimi.dsi.fastutil.longs.LongSet) LongOpenHashSet(it.unimi.dsi.fastutil.longs.LongOpenHashSet) BlockBuilder(com.facebook.presto.common.block.BlockBuilder) SqlType(com.facebook.presto.spi.function.SqlType)

Aggregations

LongOpenHashSet (it.unimi.dsi.fastutil.longs.LongOpenHashSet)13 LongSet (it.unimi.dsi.fastutil.longs.LongSet)5 LongArrayList (it.unimi.dsi.fastutil.longs.LongArrayList)3 SqlType (com.facebook.presto.spi.function.SqlType)2 Supplier (com.google.common.base.Supplier)2 File (java.io.File)2 Collection (java.util.Collection)2 Map (java.util.Map)2 BlockBuilder (com.facebook.presto.common.block.BlockBuilder)1 BlockBuilder (com.facebook.presto.spi.block.BlockBuilder)1 BlockBuilderStatus (com.facebook.presto.spi.block.BlockBuilderStatus)1 JsonCreator (com.fasterxml.jackson.annotation.JsonCreator)1 JsonIgnore (com.fasterxml.jackson.annotation.JsonIgnore)1 JsonInclude (com.fasterxml.jackson.annotation.JsonInclude)1 JsonProperty (com.fasterxml.jackson.annotation.JsonProperty)1 VisibleForTesting (com.google.common.annotations.VisibleForTesting)1 Joiner (com.google.common.base.Joiner)1 Preconditions (com.google.common.base.Preconditions)1 Predicate (com.google.common.base.Predicate)1 Suppliers (com.google.common.base.Suppliers)1