]> git.sesse.net Git - plocate/blobdiff - plocate-build.cpp
Add support for some basic options in plocate-build; specifically, block size.
[plocate] / plocate-build.cpp
index c6804497248f4a9a828200d784b9450f8986ad72..a077ac1944358ecee7e622e9846982d7ba39c064 100644 (file)
@@ -1,21 +1,21 @@
 #include "db.h"
-#include "vp4.h"
+#include "turbopfor-encode.h"
 
 #include <algorithm>
-#include <arpa/inet.h>
 #include <assert.h>
 #include <chrono>
-#include <endian.h>
-#include <fcntl.h>
+#include <iosfwd>
+#include <getopt.h>
 #include <math.h>
 #include <memory>
+#include <stdint.h>
 #include <stdio.h>
+#include <stdlib.h>
 #include <string.h>
 #include <string>
+#include <string_view>
 #include <sys/stat.h>
-#include <sys/types.h>
-#include <unistd.h>
-#include <unordered_map>
+#include <utility>
 #include <vector>
 #include <zstd.h>
 
@@ -23,6 +23,8 @@
 #define dprintf(...)
 //#define dprintf(...) fprintf(stderr, __VA_ARGS__);
 
+#define NUM_TRIGRAMS 16777216
+
 using namespace std;
 using namespace std::chrono;
 
@@ -70,7 +72,7 @@ struct db_directory {
 
 class PostingListBuilder {
 public:
-       void add_docid(uint32_t docid);
+       inline void add_docid(uint32_t docid);
        void finish();
 
        string encoded;
@@ -80,34 +82,31 @@ private:
        void write_header(uint32_t docid);
        void append_block();
 
-       vector<uint32_t> pending_docids;
+       vector<uint32_t> pending_deltas;
 
-       uint32_t last_block_end;
+       uint32_t last_block_end, last_docid = -1;
 };
 
 void PostingListBuilder::add_docid(uint32_t docid)
 {
        // Deduplicate against the last inserted value, if any.
-       if (pending_docids.empty()) {
-               if (encoded.empty()) {
-                       // Very first docid.
-                       write_header(docid);
-                       ++num_docids;
-                       last_block_end = docid;
-                       return;
-               } else if (docid == last_block_end) {
-                       return;
-               }
-       } else {
-               if (docid == pending_docids.back()) {
-                       return;
-               }
+       if (docid == last_docid) {
+               return;
+       }
+
+       if (num_docids == 0) {
+               // Very first docid.
+               write_header(docid);
+               ++num_docids;
+               last_block_end = last_docid = docid;
+               return;
        }
 
-       pending_docids.push_back(docid);
-       if (pending_docids.size() == 128) {
+       pending_deltas.push_back(docid - last_docid - 1);
+       last_docid = docid;
+       if (pending_deltas.size() == 128) {
                append_block();
-               pending_docids.clear();
+               pending_deltas.clear();
                last_block_end = docid;
        }
        ++num_docids;
@@ -115,7 +114,7 @@ void PostingListBuilder::add_docid(uint32_t docid)
 
 void PostingListBuilder::finish()
 {
-       if (pending_docids.empty()) {
+       if (pending_deltas.empty()) {
                return;
        }
 
@@ -123,37 +122,58 @@ void PostingListBuilder::finish()
 
        // No interleaving for partial blocks.
        unsigned char buf[P4NENC_BOUND(128)];
-       unsigned char *end = p4d1enc32(pending_docids.data(), pending_docids.size(), buf, last_block_end);
+       unsigned char *end = encode_pfor_single_block<128>(pending_deltas.data(), pending_deltas.size(), /*interleaved=*/false, buf);
        encoded.append(reinterpret_cast<char *>(buf), reinterpret_cast<char *>(end));
 }
 
 void PostingListBuilder::append_block()
 {
        unsigned char buf[P4NENC_BOUND(128)];
-       assert(pending_docids.size() == 128);
-       unsigned char *end = p4d1enc128v32(pending_docids.data(), 128, buf, last_block_end);
+       assert(pending_deltas.size() == 128);
+       unsigned char *end = encode_pfor_single_block<128>(pending_deltas.data(), 128, /*interleaved=*/true, buf);
        encoded.append(reinterpret_cast<char *>(buf), reinterpret_cast<char *>(end));
 }
 
 void PostingListBuilder::write_header(uint32_t docid)
 {
        unsigned char buf[P4NENC_BOUND(1)];
-       size_t bytes = p4nd1enc128v32(&docid, 1, buf);
-       encoded.append(reinterpret_cast<char *>(buf), bytes);
+       unsigned char *end = write_baseval(docid, buf);
+       encoded.append(reinterpret_cast<char *>(buf), end - buf);
 }
 
 class Corpus {
 public:
        Corpus(FILE *outfp, size_t block_size)
-               : outfp(outfp), block_size(block_size) {}
+               : invindex(new PostingListBuilder *[NUM_TRIGRAMS]), outfp(outfp), block_size(block_size)
+       {
+               fill(invindex.get(), invindex.get() + NUM_TRIGRAMS, nullptr);
+       }
+       ~Corpus()
+       {
+               for (unsigned i = 0; i < NUM_TRIGRAMS; ++i) {
+                       delete invindex[i];
+               }
+       }
+
        void add_file(string filename);
        void flush_block();
 
        vector<uint64_t> filename_blocks;
-       unordered_map<uint32_t, PostingListBuilder> invindex;
        size_t num_files = 0, num_files_in_block = 0, num_blocks = 0;
+       bool seen_trigram(uint32_t trgm)
+       {
+               return invindex[trgm] != nullptr;
+       }
+       PostingListBuilder &get_pl_builder(uint32_t trgm)
+       {
+               if (invindex[trgm] == nullptr) {
+                       invindex[trgm] = new PostingListBuilder;
+               }
+               return *invindex[trgm];
+       }
 
 private:
+       unique_ptr<PostingListBuilder *[]> invindex;
        FILE *outfp;
        string current_block;
        string tempbuf;
@@ -187,7 +207,7 @@ void Corpus::flush_block()
                if (s.size() >= 3) {
                        for (size_t j = 0; j < s.size() - 2; ++j) {
                                uint32_t trgm = read_trigram(s, j);
-                               invindex[trgm].add_docid(docid);
+                               get_pl_builder(trgm).add_docid(docid);
                        }
                }
                ptr += s.size() + 1;
@@ -209,7 +229,7 @@ void Corpus::flush_block()
 string read_cstr(FILE *fp)
 {
        string ret;
-       for ( ;; ) {
+       for (;;) {
                int ch = getc(fp);
                if (ch == -1) {
                        perror("getc");
@@ -309,7 +329,7 @@ uint32_t next_prime(uint32_t x)
        return x;
 }
 
-unique_ptr<Trigram[]> create_hashtable(const Corpus &corpus, const vector<uint32_t> &all_trigrams, uint32_t ht_size, uint32_t num_overflow_slots)
+unique_ptr<Trigram[]> create_hashtable(Corpus &corpus, const vector<uint32_t> &all_trigrams, uint32_t ht_size, uint32_t num_overflow_slots)
 {
        unique_ptr<Trigram[]> ht(new Trigram[ht_size + num_overflow_slots + 1]);  // 1 for the sentinel element at the end.
        for (unsigned i = 0; i < ht_size + num_overflow_slots + 1; ++i) {
@@ -319,7 +339,7 @@ unique_ptr<Trigram[]> create_hashtable(const Corpus &corpus, const vector<uint32
        }
        for (uint32_t trgm : all_trigrams) {
                // We don't know offset yet, so set it to zero.
-               Trigram to_insert{ trgm, uint32_t(corpus.invindex.find(trgm)->second.num_docids), 0 };
+               Trigram to_insert{ trgm, uint32_t(corpus.get_pl_builder(trgm).num_docids), 0 };
 
                uint32_t bucket = hash_trigram(trgm, ht_size);
                unsigned distance = 0;
@@ -354,6 +374,7 @@ void do_build(const char *infile, const char *outfile, int block_size)
        hdr.version = -1;  // Mark as broken.
        hdr.hashtable_size = 0;  // Not known yet.
        hdr.extra_ht_slots = num_overflow_slots;
+       hdr.num_docids = 0;
        hdr.hash_table_offset_bytes = -1;  // We don't know these offsets yet.
        hdr.filename_index_offset_bytes = -1;
        fwrite(&hdr, sizeof(hdr), 1, outfp);
@@ -393,7 +414,10 @@ void do_build(const char *infile, const char *outfile, int block_size)
        // Finish up encoding the posting lists.
        size_t trigrams = 0, longest_posting_list = 0;
        size_t bytes_for_posting_lists = 0;
-       for (auto &[trigram, pl_builder] : corpus.invindex) {
+       for (unsigned trgm = 0; trgm < NUM_TRIGRAMS; ++trgm) {
+               if (!corpus.seen_trigram(trgm))
+                       continue;
+               PostingListBuilder &pl_builder = corpus.get_pl_builder(trgm);
                pl_builder.finish();
                longest_posting_list = max(longest_posting_list, pl_builder.num_docids);
                trigrams += pl_builder.num_docids;
@@ -405,13 +429,13 @@ void do_build(const char *infile, const char *outfile, int block_size)
 
        dprintf("Building posting lists took %.1f ms.\n\n", 1e3 * duration<float>(steady_clock::now() - start).count());
 
-       // Sort the trigrams, mostly to get a consistent result every time
-       // (the hash table will put things in random order anyway).
+       // Find the used trigrams.
        vector<uint32_t> all_trigrams;
-       for (auto &[trigram, pl_builder] : corpus.invindex) {
-               all_trigrams.push_back(trigram);
+       for (unsigned trgm = 0; trgm < NUM_TRIGRAMS; ++trgm) {
+               if (corpus.seen_trigram(trgm)) {
+                       all_trigrams.push_back(trgm);
+               }
        }
-       sort(all_trigrams.begin(), all_trigrams.end());
 
        // Create the hash table.
        unique_ptr<Trigram[]> hashtable;
@@ -436,7 +460,7 @@ void do_build(const char *infile, const char *outfile, int block_size)
                        continue;
                }
 
-               const string &encoded = corpus.invindex[hashtable[i].trgm].encoded;
+               const string &encoded = corpus.get_pl_builder(hashtable[i].trgm).encoded;
                offset += encoded.size();
        }
 
@@ -450,7 +474,7 @@ void do_build(const char *infile, const char *outfile, int block_size)
                if (hashtable[i].num_docids == 0) {
                        continue;
                }
-               const string &encoded = corpus.invindex[hashtable[i].trgm].encoded;
+               const string &encoded = corpus.get_pl_builder(hashtable[i].trgm).encoded;
                fwrite(encoded.data(), encoded.size(), 1, outfp);
        }
 
@@ -471,8 +495,66 @@ void do_build(const char *infile, const char *outfile, int block_size)
        dprintf("\n");
 }
 
+void usage()
+{
+       printf(
+               "Usage: plocate-build MLOCATE_DB PLOCATE_DB\n"
+               "\n"
+               "Generate plocate index from mlocate.db, typically /var/lib/mlocate/mlocate.db.\n"
+               "Normally, the destination should be /var/lib/mlocate/plocate.db.\n"
+               "\n"
+               "  -b, --block-size SIZE  number of filenames to store in each block (default 32)\n"
+               "      --help             print this help\n"
+               "      --version          print version information\n");
+}
+
+void version()
+{
+       printf("plocate-build %s\n", PLOCATE_VERSION);
+       printf("Copyright 2020 Steinar H. Gunderson\n");
+       printf("License GPLv2+: GNU GPL version 2 or later <https://gnu.org/licenses/gpl.html>.\n");
+       printf("This is free software: you are free to change and redistribute it.\n");
+       printf("There is NO WARRANTY, to the extent permitted by law.\n");
+}
+
 int main(int argc, char **argv)
 {
-       do_build(argv[1], argv[2], 32);
+       static const struct option long_options[] = {
+               { "block-size", required_argument, 0, 'b' },
+               { "help", no_argument, 0, 'h' },
+               { "version", no_argument, 0, 'V' },
+               { 0, 0, 0, 0 }
+       };
+
+       int block_size = 32;
+
+       setlocale(LC_ALL, "");
+       for (;;) {
+               int option_index = 0;
+               int c = getopt_long(argc, argv, "b:hV", long_options, &option_index);
+               if (c == -1) {
+                       break;
+               }
+               switch (c) {
+               case 'b':
+                       block_size = atoi(optarg);
+                       break;
+               case 'h':
+                       usage();
+                       exit(0);
+               case 'v':
+                       version();
+                       exit(0);
+               default:
+                       exit(1);
+               }
+       }
+
+       if (argc - optind != 2) {
+               usage();
+               exit(1);
+       }
+
+       do_build(argv[optind], argv[optind + 1], block_size);
        exit(EXIT_SUCCESS);
 }