hoshinojyunn commented on code in PR #64909: URL: https://github.com/apache/doris/pull/64909#discussion_r3542239303
########## be/src/storage/index/snii/writer/spimi_term_buffer.cpp: ########## @@ -0,0 +1,708 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "storage/index/snii/writer/spimi_term_buffer.h" + +#include <unistd.h> + +#include <algorithm> +#include <atomic> +#include <cstdio> +#include <cstdlib> +#include <memory> +#include <numeric> +#include <string> +#include <type_traits> +#include <utility> + +#include "storage/index/snii/encoding/varint.h" +#include "storage/index/snii/format/format_constants.h" +#include "storage/index/snii/writer/spill_run_codec.h" +#include "storage/index/snii/writer/temp_dir.h" + +#if defined(__GLIBC__) +#include <malloc.h> +#endif + +namespace doris::snii::writer { + +namespace { + +// Returns freed heap arenas to the OS (glibc only). The spill encode churns many +// small allocations whose freed chunks glibc retains in its arenas; trimming +// before the peak-RSS-defining merge phase recovers that retention. No-op (and +// harmless) on non-glibc libcs. +void TrimMalloc() { +#if defined(__GLIBC__) + ::malloc_trim(0); +#endif +} + +// Process-unique temp path for a spill run under `dir` (pid + monotonic counter so +// parallel builds / multiple buffers never collide). +std::string MakeRunPath(const std::string& dir) { + static std::atomic<uint64_t> counter {0}; + const uint64_t n = counter.fetch_add(1); + return dir + "/snii_spill_" + std::to_string(::getpid()) + "_" + std::to_string(n) + ".run"; +} + +// TEST-ONLY seam backing testing::vocab_string_materialization_count(). Bumped once +// per DISTINCT interned term (owned_vocab_.emplace_back), never per token. Relaxed: +// the build path is single-threaded, so only the COUNT matters, not ordering. +std::atomic<uint64_t> g_vocab_materializations {0}; + +} // namespace + +namespace testing { +uint64_t vocab_string_materialization_count() { + return g_vocab_materializations.load(std::memory_order_relaxed); +} +void reset_vocab_string_materialization_count() { + g_vocab_materializations.store(0, std::memory_order_relaxed); +} +} // namespace testing + +SpimiTermBuffer::SpimiTermBuffer(const std::vector<std::string>* vocab, bool has_positions, + size_t spill_threshold_bytes, MemoryReporter* reporter) + : vocab_(vocab), + // Bind the interning set's heterogeneous functors to &owned_vocab_ even in + // borrowed mode: the back-pointer is harmless here because + // add_token(string_view) rejects before touching intern_ (the functors are + // never dereferenced), and binding unconditionally keeps both constructors + // symmetric. Initialized in the member-init list (NOT the body): the functors + // are NESTED types, whose default-constructibility is not yet established at + // the point unordered_set's defaulted default ctor would be needed for a + // body assignment, so default-constructing intern_ is ill-formed. The + // (bucket_count, hash, equal) constructor sidesteps that entirely. + // owned_vocab_ is constructed before intern_ (declaration order) and the + // buffer is non-movable, so &owned_vocab_ is stable for the buffer's life. + intern_(0, OwnedVocabHash {&owned_vocab_}, OwnedVocabEq {&owned_vocab_}), + has_positions_(has_positions), + spill_threshold_bytes_(spill_threshold_bytes), + mem_reporter_(reporter) { + // Borrowed-vocab mode: only the 4 B/id slot-index array is sized to the + // vocabulary; the Term pool (slots_) grows with the LIVE touched count, so an + // all-but-empty vocabulary costs ~4 B/id instead of ~80 B/id. + slot_of_.assign(vocab_->size(), 0); + // The vocab-sized slot index is resident immediately and survives spills; report + // its initial positive delta now. + report_arena_delta(); +} + +SpimiTermBuffer::SpimiTermBuffer(bool has_positions, size_t spill_threshold_bytes, + MemoryReporter* reporter) + : vocab_(&owned_vocab_), + // Owned-vocab mode: bind the interning set's heterogeneous functors to + // &owned_vocab_ so a stored term-id dereferences to its string for content + // hashing and equality. Initialized in the member-init list (NOT the body): + // the functors are NESTED types whose default-constructibility is not yet + // established where unordered_set's defaulted default ctor would be needed for + // a body assignment, so the (bucket_count, hash, equal) constructor is used + // instead. owned_vocab_ is constructed before intern_ (declaration order) and + // the buffer is non-movable, so &owned_vocab_ is stable for the buffer's life. + intern_(0, OwnedVocabHash {&owned_vocab_}, OwnedVocabEq {&owned_vocab_}), + has_positions_(has_positions), + spill_threshold_bytes_(spill_threshold_bytes), + mem_reporter_(reporter) { + // Owned-vocab mode: the vocabulary grows as strings are interned in + // add_token(string_view, ...). +} + +SpimiTermBuffer::~SpimiTermBuffer() { + // Balance the writer-level / Doris tracker on the error path: if the buffer is + // destroyed while resident bytes were reported but not yet freed-and-reported + // (e.g. a build aborts before draining), return them here so nothing leaks. + if (mem_reporter_ != nullptr && reported_resident_ != 0) { + mem_reporter_->report(-reported_resident_); + reported_resident_ = 0; + } + cleanup_runs(); +} + +void SpimiTermBuffer::report_arena_delta() { + if (mem_reporter_ == nullptr) { + return; + } + // Diff the REAL resident bytes (arena + slot index) against the last reported + // total; emit the signed delta exactly once. + const auto now = static_cast<int64_t>(resident_bytes()); + mem_reporter_->report(now - reported_resident_); + reported_resident_ = now; +} + +size_t SpimiTermBuffer::unique_terms() const { + return live_term_count_; +} + +uint64_t SpimiTermBuffer::resident_bytes() const { + // REAL resident accumulator bytes: the posting arena plus the vocab-sized slot + // index (capacity, since the reserved-but-unused tail is still resident RSS and + // survives spills -- spill_to_run does NOT free slot_of_). This is the gate-2 + // spill trigger metric and the spill space-precheck figure -- NOT the old gated + // live_bytes_ estimate. + return pool_.arena_bytes() + static_cast<uint64_t>(slot_of_.capacity()) * sizeof(uint32_t); +} + +// Returns the live Term for `term_id`, claiming a pool slot on first touch (1 == +// new). Reuses a freed slot from free_slots_ when available; otherwise appends a +// fresh Term to slots_. slot_of_[term_id] holds (slot index + 1); 0 means empty. +SpimiTermBuffer::Term& SpimiTermBuffer::term_slot(uint32_t term_id, bool* new_term) { + uint32_t enc = slot_of_[term_id]; + if (enc != 0) { + *new_term = false; + return slots_[enc - 1]; + } + *new_term = true; + uint32_t slot; + if (!free_slots_.empty()) { + slot = free_slots_.back(); + free_slots_.pop_back(); + } else { + slot = static_cast<uint32_t>(slots_.size()); + slots_.emplace_back(); + } + slot_of_[term_id] = slot + 1; + return slots_[slot]; +} + +// Appends one byte to a term's chain, starting the chain lazily on first use. +void SpimiTermBuffer::put_byte(Term* t, uint8_t b) { + if (t->head == kNoChain) { + t->head = pool_.start_chain(&t->w, &t->level); + } + pool_.append_byte(&t->w, &t->level, b); +} + +void SpimiTermBuffer::put_varint(Term* t, uint64_t v) { + uint8_t tmp[10]; + const size_t n = encode_varint64(v, tmp); + for (size_t i = 0; i < n; ++i) { + put_byte(t, tmp[i]); + } +} + +void SpimiTermBuffer::accumulate(uint32_t term_id, uint32_t docid, uint32_t pos) { + bool new_term = false; + Term& t = term_slot(term_id, &new_term); + if (new_term) { + touched_ids_.push_back(term_id); + ++live_term_count_; + } + // A token starts a new doc unless it continues the most-recent doc for this term. + const bool new_doc = !t.started || t.cur_docid != docid; + // Tagged entry: varint((pos << 1) | new_doc). Positions are tagged 0 when + // disabled. The new_doc bit lets the decoder recover per-doc freqs by counting. + // Widen to 64-bit so a full 32-bit position survives the << 1 without truncation. + const uint64_t tagged = has_positions_ + ? ((static_cast<uint64_t>(pos) << 1) | (new_doc ? 1U : 0U)) + : (new_doc ? 1U : 0U); + put_varint(&t, tagged); + if (new_doc) { + // Out-of-order docids are tolerated (zigzag delta is signed) and reordered at + // finalize; flag them so to_postings sorts. The delta base is the previous + // distinct doc (cur_docid), which is 0 for the very first doc (started==false). + const int64_t base = t.started ? static_cast<int64_t>(t.cur_docid) : 0; + if (t.started && docid < t.cur_docid) { + t.sorted = false; + } + const int64_t delta = static_cast<int64_t>(docid) - base; + put_varint(&t, zigzag_encode(delta)); + t.cur_docid = docid; + t.started = true; + } + ++t.ntok; + ++total_tokens_; + + // Gate-2 spill: trigger on REAL resident bytes (arena + slot index), NOT the old + // gated live_bytes_ estimate. arena_bytes() is monotonic per fill and reset to 0 + // by spill_to_run()'s pool_.reset(), so the trigger self-rearms after each spill. + // The OTHER trigger is the hard arena safety stop (active even in unlimited mode): + // when the arena nears the 4 GiB uint32-offset limit -- without it, a single + // >4 GiB in-memory segment wraps alloc_run and silently corrupts data. A forced + // spill + final k-way merge stays byte-identical regardless of when it fires. + constexpr uint64_t kArenaSpillCap = 0xE0000000ULL; // 3.5 GiB, < UINT32_MAX margin + // Report this token's REAL resident growth FIRST so the writer's unified total + // (reporter_->current_bytes()) reflects it before the gate-2 check. Single-source + // diff: cheap (subtraction + relaxed atomic add; arena_bytes() is two field reads). + report_arena_delta(); + // Gate-2 spill (UNIFIED): when a reporter is attached, trigger on the writer's TOTAL + // build RAM (arena + slot index + dict) crossing the one configured cap -- the same + // total and cap every buffer of this writer shares, not a per-buffer threshold. Off + // Doris (no reporter) fall back to the local spill_threshold_bytes_. The hard arena + // safety stop (4 GiB uint32-offset limit) is always active. spill_to_run() resets the + // arena and reports its negative internally, so the unified total drops after a spill. + const bool over_cap = mem_reporter_ != nullptr ? mem_reporter_->over_cap() + : (spill_threshold_bytes_ != 0 && + resident_bytes() >= spill_threshold_bytes_); + const bool arena_near_limit = pool_.arena_bytes() >= kArenaSpillCap; + if ((over_cap || arena_near_limit) && spill_status_.ok()) { + spill_status_ = spill_to_run(); + } +} + +void SpimiTermBuffer::add_token(uint32_t term_id, uint32_t docid, uint32_t pos) { + // Hot path: a pooled slot lookup + a couple of pushes. No hashing, no string + // construction per token. Reject (and latch) an out-of-range id. + if (term_id >= slot_of_.size()) { + if (spill_status_.ok()) { + spill_status_ = Status::Error<ErrorCode::INVALID_ARGUMENT, false>( + "spimi: term_id out of vocab range"); + } + return; + } + accumulate(term_id, docid, pos); +} + +void SpimiTermBuffer::add_token(std::string_view term, uint32_t docid, uint32_t pos) { + // Compatibility path: intern the term into the owned vocabulary on first + // occurrence, then accumulate by its id. ONLY valid in OWNED-vocab mode. In + // BORROWED-vocab mode vocab_ points at the caller's vector, NOT &owned_vocab_: + // interning here would grow owned_vocab_ / intern_ / slot_of_ out of step with + // the active (borrowed) vocab, so the new id indexes the WRONG string and writes + // a slot_of_ entry the borrowed-vocab build never reconciles -- silent + // corruption. Reject (and latch) instead of forwarding by a bogus id. + if (vocab_ != &owned_vocab_) { + if (spill_status_.ok()) { + spill_status_ = Status::Error<ErrorCode::INVALID_ARGUMENT, false>( + "spimi: add_token(string_view) requires owned-vocab mode"); + } + return; + } + // F03 single-store invariant, fixed at compile time: the interning set keys on the + // 4-byte term-id, NEVER on a std::string, so each vocab string lives in exactly one + // place (owned_vocab_). A regression to a string key would reintroduce the + // double-store and fail this build. + static_assert(std::is_same_v<decltype(intern_)::key_type, uint32_t>, + "intern_ must key on term-id (single-store); a string key reintroduces F03"); + + // Heterogeneous probe with the string_view directly: NO per-token temporary + // std::string (F21). The set element is the term-id; its content is resolved + // through owned_vocab_ by the transparent functors. + auto it = intern_.find(term); + uint32_t term_id; + if (it == intern_.end()) { + term_id = static_cast<uint32_t>(owned_vocab_.size()); + // The SOLE materialization of this term's string (F03 single-store). Push + // FIRST so owned_vocab_[term_id] is valid before insert(term_id) hashes it. The + // set stores only the id, so this emplace_back's possible reallocation never + // invalidates existing entries (their functors re-read owned_vocab_[id]). + owned_vocab_.emplace_back(term); + g_vocab_materializations.fetch_add(1, std::memory_order_relaxed); + intern_.insert(term_id); + slot_of_.push_back(0); // vocab grows: new id starts with no live slot + } else { + term_id = *it; // the set element IS the term-id + } + accumulate(term_id, docid, pos); +} + +namespace { + +// Reorders a term's flat arrays into ascending-docid order, COALESCING any +// same-docid groups so the result has exactly one entry per docid -- matching the +// k-way-merge path's boundary-doc coalescing and the writer's strictly-ascending +// precondition. Only invoked for the rare term that received out-of-order docids +// (the common ascending path leaves t.sorted true and skips it). +// +// A docid may REVISIT (e.g. feed 5,1,5): the chain holds two separate doc-groups +// for doc 5. A STABLE sort keeps equal-docid groups in arrival order, then the +// coalesce pass sums their freqs and concatenates their positions in that same +// (document/arrival) order -- so the merged positions stay consistent with the +// merged freqs, exactly as the run-order merge would have produced. +void SortByDocid(std::vector<uint32_t>* docids, std::vector<uint32_t>* freqs, + std::vector<uint32_t>* positions_flat, bool has_positions) { + const size_t n = docids->size(); + std::vector<size_t> order(n); + std::iota(order.begin(), order.end(), 0); + // STABLE so equal docids keep arrival order: their positions then concatenate in + // document order, the same order the merge path's run concatenation yields. + std::ranges::stable_sort(order, + [&](size_t a, size_t b) { return (*docids)[a] < (*docids)[b]; }); + + std::vector<uint32_t> pos_off; + if (has_positions) { + pos_off.resize(n); + uint32_t running = 0; + for (size_t i = 0; i < n; ++i) { + pos_off[i] = running; + running += (*freqs)[i]; + } + } + std::vector<uint32_t> nd, nf, np; + nd.reserve(n); + nf.reserve(n); + if (has_positions) { + np.reserve(positions_flat->size()); + } + for (size_t k : order) { + // Coalesce a revisited docid into the previous entry (it sorts adjacent now): + // sum freqs and append this group's positions right after the prior group's, + // so flat doc order stays partitioned by the merged freqs. + if (!nd.empty() && nd.back() == (*docids)[k]) { + nf.back() += (*freqs)[k]; + } else { + nd.push_back((*docids)[k]); + nf.push_back((*freqs)[k]); + } + if (has_positions) { + np.insert(np.end(), positions_flat->begin() + pos_off[k], + positions_flat->begin() + pos_off[k] + (*freqs)[k]); + } + } + *docids = std::move(nd); + *freqs = std::move(nf); + if (has_positions) { + *positions_flat = std::move(np); + } +} + +} // namespace + +namespace { + +// Decodes one varint from a pool chain cursor. The chain was written by +// encode_varint*, so the same LEB128 continuation-bit loop reconstructs it. +uint64_t DecodeChainVarint(CompactPostingPool::Cursor* c) { + uint64_t result = 0; + int shift = 0; + for (;;) { + const uint8_t b = c->next(); + result |= static_cast<uint64_t>(b & 0x7F) << shift; + if ((b & 0x80) == 0) { + break; + } + shift += 7; + } + return result; +} + +void SkipChainVarint(CompactPostingPool::Cursor* c) { + DecodeChainVarint(c); +} + +} // namespace + +// Decodes a term's compact tagged chain back into a flat TermPostings (the exact +// docids/freqs/positions_flat the writer consumes), so the produced index is +// byte-identical to the legacy raw-uint32 accumulator. The chain holds one entry +// per token: varint((pos << 1) | new_doc); each new_doc entry is followed by a +// zigzag(docid-delta). A doc's freq is the run length of consecutive same-doc +// tokens; positions stream out in document order (empty when positions disabled). +// Stream positions for a sorted term whose token count exceeds this: such a term's +// flat positions buffer (uint32 per token) would be the peak-RSS transient (tens of +// MiB for the widest term). Below it, the flat buffer is cheap and simpler. +static constexpr uint32_t kStreamPositionsTokenThreshold = 1U << 16; // 65536 + +TermPostings SpimiTermBuffer::to_postings(std::string term, Term&& t, + bool allow_stream_positions) const { + TermPostings tp; + tp.term = std::move(term); + if (t.ntok == 0 || t.head == kNoChain) { + return tp; + } + + // Reserve docids/freqs by ntok (an upper bound on the doc count: ntok >= ndocs). + // The doc count is not stored separately to keep Term compact; since the corpus + // is freq~1 per (term, doc), ntok ~= ndocs so the over-reserve is negligible. + tp.docids.reserve(t.ntok); Review Comment: Is allocating the entire ntok size of memory directly to docids and freqs too aggressive? Wouldn't it be better to use an `nDocs` field in the `Term` to record the doc freq? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
