MaxScore TAAT and blocked accumulators

CMakeLists.txt

@@ -40,7 +40,6 @@ list(APPEND LCOV_REMOVE_PATTERNS "'${PROJECT_SOURCE_DIR}/external/*'")
 
 
 if (UNIX)
-
    # For hardware popcount and other special instructions
    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -march=native")
 
@@ -62,8 +61,6 @@ endif()
 
 set(THREADS_PREFER_PTHREAD_FLAG ON)
 find_package(Threads REQUIRED)
-link_libraries(Threads::Threads)
-
 
 include_directories(include)
 add_library(pisa INTERFACE)

include/pisa/accumulator/blocked_accumulator.hpp

@@ -0,0 +1,84 @@
+#pragma once
+
+namespace pisa {
+
+template <int block_size>
+struct Blocked_Accumulator {
+
+    struct Proxy_Element {
+        std::ptrdiff_t      document;
+        std::vector<float> &accumulators;
+        std::vector<float> &accumulators_max;
+
+        Proxy_Element &operator=(float score) {
+            accumulators[document] = score;
+            auto &block_max        = accumulators_max[document / block_size];
+            if (score > block_max) {
+                block_max = score;
+            }
+            return *this;
+        }
+        Proxy_Element &operator+=(float delta) {
+            accumulators[document] += delta;
+            auto const&score = accumulators[document];
+            auto &block_max = accumulators_max[document / block_size];
+            if (score > block_max) {
+                block_max = score;
+            }
+            return *this;
+        }
+
+        operator float() { return accumulators[document]; }
+    };
+
+    using reference = Proxy_Element;
+
+    static_assert(block_size > 0, "must be positive");
+
+    [[nodiscard]] constexpr static auto calc_block_count(std::size_t size) noexcept -> std::size_t {
+        return (size + block_size - 1) / block_size;
+    }
+
+    Blocked_Accumulator(std::size_t size)
+        : m_size(size),
+        m_block_count(calc_block_count(size)), m_accumulators(size),
+        m_accumulators_max(m_block_count) {}
+
+    void init() { std::fill(m_accumulators.begin(), m_accumulators.end(), 0.0); }
+
+    [[nodiscard]] auto operator[](std::ptrdiff_t document) -> Proxy_Element
+    {
+        return {document, m_accumulators, m_accumulators_max};
+    }
+
+    void accumulate(std::ptrdiff_t const document, float score_delta)
+    {
+        m_accumulators[document] += score_delta;
+        auto const &score = m_accumulators[document];
+        auto &block_max = m_accumulators_max[document / block_size];
+        if (score > block_max) {
+            block_max = score;
+        }
+    }
+
+    void aggregate(topk_queue &topk) {
+        for (size_t block = 0; block < m_block_count; ++block) {
+            if (not topk.would_enter(m_accumulators_max[block])) { continue; }
+            uint32_t doc = block * block_size;
+            uint32_t end = std::min((block + 1) * block_size, m_accumulators.size());
+            for (; doc < end; ++doc) {
+                topk.insert(m_accumulators[doc], doc);
+            }
+        }
+    }
+
+    [[nodiscard]] auto size() noexcept -> std::size_t { return m_size; }
+
+   private:
+    std::size_t        m_size;
+    std::size_t        m_block_count;
+    std::vector<float> m_accumulators;
+    std::vector<float> m_accumulators_max;
+};
+
+} // pisa

include/pisa/accumulator/lazy_accumulator.hpp

@@ -0,0 +1,95 @@
+#pragma once
+
+namespace pisa {
+
+template <int counter_bit_size, typename Descriptor = std::uint64_t>
+struct Lazy_Accumulator {
+    using reference = float &;
+
+    static_assert(std::is_integral_v<Descriptor> && std::is_unsigned_v<Descriptor>,
+                  "must be unsigned number");
+    constexpr static auto descriptor_size_in_bits = sizeof(Descriptor) * 8;
+    constexpr static auto counters_in_descriptor = descriptor_size_in_bits / counter_bit_size;
+    constexpr static auto cycle = (1u << counter_bit_size);
+    constexpr static Descriptor mask = (1u << counter_bit_size) - 1;
+
+    struct Block {
+        Descriptor                                descriptor{};
+        std::array<float, counters_in_descriptor> accumulators{};
+
+        [[nodiscard]] auto counter(int pos) const noexcept -> int {
+            return (descriptor >> (pos * counter_bit_size)) & mask;
+        }
+
+        void reset_counter(int pos, int counter)
+        {
+            auto const shift = pos * counter_bit_size;
+            descriptor &= ~(mask << shift);
+            descriptor |= static_cast<Descriptor>(counter) << shift;
+            accumulators[pos] = 0;
+        }
+    };
+
+    Lazy_Accumulator(std::size_t size)
+        : m_size(size), m_accumulators((size + counters_in_descriptor - 1) / counters_in_descriptor)
+    {}
+
+    void init()
+    {
+        if (m_counter == 0) {
+            auto first = reinterpret_cast<std::byte *>(&m_accumulators.front());
+            auto last =
+                std::next(reinterpret_cast<std::byte *>(&m_accumulators.back()), sizeof(Block));
+            std::fill(first, last, std::byte{0});
+        }
+    }
+
+    float &operator[](std::ptrdiff_t const document) {
+        auto const block        = document / counters_in_descriptor;
+        auto const pos_in_block = document % counters_in_descriptor;
+        if (//m_accumulators[block].accumulators[pos_in_block] > 0 &&
+            m_accumulators[block].counter(pos_in_block) != m_counter)
+        {
+            auto const shift = pos_in_block * counter_bit_size;
+            m_accumulators[block].descriptor &= ~(mask << shift);
+            m_accumulators[block].descriptor |= m_counter << shift;
+            m_accumulators[block].accumulators[pos_in_block] = 0;
+        }
+        return m_accumulators[block].accumulators[pos_in_block];
+    }
+
+    void accumulate(std::ptrdiff_t const document, float score)
+    {
+        auto const block = document / counters_in_descriptor;
+        auto const pos_in_block = document % counters_in_descriptor;
+        if (m_accumulators[block].counter(pos_in_block) != m_counter) {
+            m_accumulators[block].reset_counter(pos_in_block, m_counter);
+        }
+        m_accumulators[block].accumulators[pos_in_block] += score;
+    }
+
+    void aggregate(topk_queue &topk) {
+        uint64_t docid = 0u;
+        for (auto const &block : m_accumulators) {
+            int pos = 0;
+            for (auto const &score : block.accumulators) {
+                if (block.counter(pos++) == m_counter) {
+                    topk.insert(score, docid);
+                }
+                ++docid;
+            }
+        };
+        m_counter = (m_counter + 1) % cycle;
+    }
+
+    [[nodiscard]] auto size() const noexcept -> std::size_t { return m_size; }
+    [[nodiscard]] auto blocks() noexcept -> std::vector<Block> & { return m_accumulators; }
+    [[nodiscard]] auto counter() const noexcept -> int { return m_counter; }
+
+   private:
+    std::size_t        m_size;
+    std::vector<Block> m_accumulators;
+    int                m_counter{};
+};
+
+}

include/pisa/accumulator/simple_accumulator.hpp

@@ -0,0 +1,15 @@
+#pragma once
+
+namespace pisa {
+
+struct Simple_Accumulator : public std::vector<float> {
+    Simple_Accumulator(std::ptrdiff_t size) : std::vector<float>(size) {}
+    void init() { std::fill(begin(), end(), 0.0); }
+    void accumulate(uint32_t doc, float score) { operator[](doc) += score; }
+    void aggregate(topk_queue &topk) {
+        uint64_t docid = 0u;
+        std::for_each(begin(), end(), [&](auto score) { topk.insert(score, docid++); });
+    }
+};
+
+}

include/pisa/block_posting_list.hpp

@@ -82,6 +82,8 @@ namespace pisa {
 
         class document_enumerator {
         public:
+            using enumerator_category = pisa::block_enumerator_tag;
+
             document_enumerator(uint8_t const* data, uint64_t universe,
                                 size_t term_id = 0)
                 : m_n(0) // just to silence warnings
@@ -156,6 +158,25 @@ namespace pisa {
                 }
             }
 
+            // TODO(michal): I recommend using some view, like gsl::span or something
+            //               instead of a reference to a vector.
+            [[nodiscard]] auto document_buffer() -> std::vector<uint32_t> const & {
+                return m_docs_buf;
+            }
+
+            [[nodiscard]] auto frequency_buffer() -> std::vector<uint32_t> const & {
+                if (!m_freqs_decoded) {
+                    decode_freqs_block();
+                }
+                return m_freqs_buf;
+            }
+
+            void next_block()
+            {
+                m_pos_in_block = m_cur_block_size - 1;
+                next();
+            }
+
             uint64_t docid() const
             {
                 return m_cur_docid;

include/pisa/freq_index.hpp

@@ -76,6 +76,7 @@ namespace pisa {
 
         class document_enumerator {
         public:
+            using enumerator_category = pisa::input_enumerator_tag;
             void reset()
             {
                 m_cur_pos = 0;

include/pisa/query/algorithm/and_query.hpp

@@ -2,11 +2,12 @@
 
 namespace pisa {
 
-template <bool with_freqs>
+template <typename Index, bool with_freqs>
 struct and_query {
 
-    template <typename Index>
-    uint64_t operator()(Index const &index, term_id_vec terms) const {
+    and_query(Index const &index) : m_index(index) {}
+
+    uint64_t operator()(term_id_vec terms) const {
         if (terms.empty())
             return 0;
         remove_duplicate_terms(terms);
@@ -16,7 +17,7 @@ struct and_query {
         enums.reserve(terms.size());
 
         for (auto term : terms) {
-            enums.push_back(index[term]);
+            enums.push_back(m_index[term]);
         }
 
         // sort by increasing frequency
@@ -27,7 +28,7 @@ struct and_query {
         uint64_t results   = 0;
         uint64_t candidate = enums[0].docid();
         size_t   i         = 1;
-        while (candidate < index.num_docs()) {
+        while (candidate < m_index.num_docs()) {
             for (; i < enums.size(); ++i) {
                 enums[i].next_geq(candidate);
                 if (enums[i].docid() != candidate) {
@@ -52,6 +53,9 @@ struct and_query {
         }
         return results;
     }
+
+   private:
+    Index const &m_index;
 };
 
 } // namespace pisa

include/pisa/query/algorithm/block_max_maxscore_query.hpp

@@ -2,22 +2,22 @@
 
 namespace pisa {
 
-template <typename WandType>
+template <typename Index, typename WandType>
 struct block_max_maxscore_query {
 
     typedef bm25 scorer_type;
 
-    block_max_maxscore_query(WandType const &wdata, uint64_t k) : m_wdata(&wdata), m_topk(k) {}
+    block_max_maxscore_query(Index const &index, WandType const &wdata, uint64_t k)
+        : m_index(index), m_wdata(&wdata), m_topk(k) {}
 
-    template <typename Index>
-    uint64_t operator()(Index const &index, term_id_vec const &terms) {
+    uint64_t operator()(term_id_vec const &terms) {
         m_topk.clear();
         if (terms.empty())
             return 0;
 
         auto query_term_freqs = query_freqs(terms);
 
-        uint64_t                                        num_docs = index.num_docs();
+        uint64_t                                        num_docs = m_index.num_docs();
         typedef typename Index::document_enumerator     enum_type;
         typedef typename WandType::wand_data_enumerator wdata_enum;
 
@@ -32,7 +32,7 @@ struct block_max_maxscore_query {
         enums.reserve(query_term_freqs.size());
 
         for (auto term : query_term_freqs) {
-            auto list       = index[term.first];
+            auto list       = m_index[term.first];
             auto w_enum     = m_wdata->getenum(term.first);
             auto q_weight   = scorer_type::query_term_weight(term.second, list.size(), num_docs);
             auto max_weight = q_weight * m_wdata->max_term_weight(term.first);
@@ -66,10 +66,10 @@ struct block_max_maxscore_query {
                              })
                 ->docs_enum.docid();
 
-        while (non_essential_lists < ordered_enums.size() && cur_doc < index.num_docs()) {
+        while (non_essential_lists < ordered_enums.size() && cur_doc < m_index.num_docs()) {
             float    score    = 0;
             float    norm_len = m_wdata->norm_len(cur_doc);
-            uint64_t next_doc = index.num_docs();
+            uint64_t next_doc = m_index.num_docs();
             for (size_t i = non_essential_lists; i < ordered_enums.size(); ++i) {
                 if (ordered_enums[i]->docs_enum.docid() == cur_doc) {
                     score +=
@@ -129,6 +129,7 @@ struct block_max_maxscore_query {
     std::vector<std::pair<float, uint64_t>> const &topk() const { return m_topk.topk(); }
 
    private:
+    Index const &   m_index;
     WandType const *m_wdata;
     topk_queue      m_topk;
 };