Variable-BMW/shuffle_docids.cpp at master · rossanoventurini/Variable-BMW · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
#include <fstream>
#include <iostream>
#include <algorithm>
#include <thread>
#include <numeric>
#include <random>

#include <succinct/mapper.hpp>

#include "binary_freq_collection.hpp"
#include "index_build_utils.hpp"
#include "util.hpp"

using ds2i::logger;

void emit(std::ostream& os, const uint32_t* vals, size_t n)
{
    os.write(reinterpret_cast<const char*>(vals), sizeof(*vals) * n);
}

void emit(std::ostream& os, uint32_t val)
{
    emit(os, &val, 1);
}

int main(int argc, const char** argv)
{

    using namespace ds2i;

    if (argc < 3) {
        std::cerr << "Usage: " << argv[0]
                  << " <collection basename> <output basename>"
                  << std::endl;
        return 1;
    }

    constexpr uint32_t seed = 1729;
    std::mt19937 rng(seed);

    const std::string input_basename = argv[1];
    const std::string output_basename = argv[2];
    binary_freq_collection input(input_basename.c_str());

    logger() << "Computing random permutation" << std::endl;
    size_t num_docs = input.num_docs();
    std::vector<uint32_t> new_doc_id(num_docs);
    std::iota(new_doc_id.begin(), new_doc_id.end(), uint32_t());
    std::shuffle(new_doc_id.begin(), new_doc_id.end(), rng);

    {
        logger() << "Shuffling document sizes" << std::endl;
        binary_collection input_sizes((input_basename + ".sizes").c_str());
        auto sizes = *input_sizes.begin();
        if (sizes.size() != num_docs) {
            throw std::invalid_argument("Invalid sizes file");
        }

        std::vector<uint32_t> new_sizes(num_docs);
        for (size_t i = 0; i < num_docs; ++i) {
            new_sizes[new_doc_id[i]] = sizes.begin()[i];
        }

        std::ofstream output_sizes(output_basename + ".sizes");
        emit(output_sizes, sizes.size());
        emit(output_sizes, new_sizes.data(), num_docs);
    }

    logger() << "Shuffling posting lists" << std::endl;
    progress_logger plog;

    std::ofstream output_docs(output_basename + ".docs");
    std::ofstream output_freqs(output_basename + ".freqs");
    emit(output_docs, 1);
    emit(output_docs, num_docs);

    std::vector<std::pair<uint32_t, uint32_t>> pl;
    for (const auto& seq: input) {

        for (size_t i = 0; i < seq.docs.size(); ++i) {
            pl.emplace_back(new_doc_id[seq.docs.begin()[i]],
                            seq.freqs.begin()[i]);
        }

        std::sort(pl.begin(), pl.end());

        emit(output_docs, pl.size());
        emit(output_freqs, pl.size());
        for (const auto& posting: pl) {
            emit(output_docs, posting.first);
            emit(output_freqs, posting.second);
        }

        plog.done_sequence(seq.docs.size());
        pl.clear();
    }
    plog.log();
}