Skip to content

anda_db_tfs.md

Latest commit

 

History

History
381 lines (283 loc) · 27 KB

anda_db_tfs.md

File metadata and controls

381 lines (283 loc) · 27 KB
Raw

anda_db_tfs - Full-Text Search Engine for AndaDB

中文版

anda_db_tfs is the embedded full-text search component of AndaDB. It implements the classic Okapi BM25 ranking algorithm and is designed for long-term textual memory in AI agents. It is written in pure Rust, thread-safe, and has very few dependencies. It can be used on its own or reused directly by the upper-layer database as the anda_db::BM25 index.

1. Design Goals

Goal Implementation
Embedded, zero external services Pure Rust library; no Elasticsearch or Tantivy process required
Friendly to mixed Chinese and English text Pluggable Tokenizer pipeline; built-in Porter stemmer and jieba tokenization
High-concurrency reads and writes DashMap plus atomic counters; insert / remove / search can run concurrently across threads
Incremental persistence The inverted index is sharded into buckets; only dirty buckets are flushed
Small memory footprint Vec, FxHashMap, and compact CBOR encoding
Boolean queries AND / OR / NOT syntax with parentheses for agent retrieval

2. Algorithm: Okapi BM25 at a Glance

For a multi-term query q, the score of document d is:

$$ ext{score}(d, q) ;=; \sum_{t \in q} \text{idf}(t) \cdot \frac{tf_{t,d},(k_1 + 1)}{tf_{t,d} + k_1!\left(1 - b + b \cdot \dfrac{|d|}{\text{avgdl}}\right)} $$

  • $tf_{t,d}$: the number of times term $t$ appears in document $d$.
  • $|d|$: the total number of tokens in document $d$; $\text{avgdl}$ is the average document length across the corpus.
  • $\text{idf}(t) = \ln!\left(1 + \dfrac{N - df_t + 0.5}{df_t + 0.5}\right)$, the classic Okapi IDF smoothing formula.
  • The hyperparameters $k_1$ and $b$ are configured through BM25Params, with defaults k1=1.2 and b=0.75.

Before scoring, the library defensively clamps user-provided parameters: k1 is clamped to [0, 1000], and b to [0, 1]; non-finite inputs fall back to the defaults to avoid generating NaN or inf.

3. In-Memory Data Structures

BM25Index
├── doc_tokens       DashMap<doc_id, token_count>           // document length table
├── postings         DashMap<token, (bucket_id, Vec<(doc_id, tf)>)>  // inverted index
├── buckets          DashMap<bucket_id, Bucket>             // shard metadata
├── metadata         RwLock<BM25Metadata>                   // name / config / stats
└── atomic counters  max_bucket_id, max_document_id, total_tokens, ...

  • Posting (bucket_id, Vec<(doc_id, tf)>): bucket_id identifies the bucket currently owning the token; the list holds one entry per indexed document, so its length is the token's document frequency. remove keeps it that way: when the supplied text does not account for all of a document's entries, it sweeps the rest by id. Entries of absent documents, possible only in data written by earlier releases, are pruned on load.

  • Bucket: a serializable unit containing a group of tokens and the doc_ids they cover. A bucket tracks dirty state with dual version counters:

    Field Meaning
    dirty_version Incremented on every modification
    saved_version Version last persisted successfully
    size Approximate byte size under current CBOR encoding
    tokens The list of terms owned by this bucket
    doc_ids The set of document IDs seen in this bucket

is_dirty() is true if and only if dirty_version > saved_version. A flush records the captured dirty versions as saved only after the manifest commit succeeds; failures leave the buckets retryable. The caller must exclude mutations, compaction, and other flushes for the entire flush. Version tracking does not make concurrent mutation during flush safe.

4. Bucket Sharding Strategy

BM25Config.bucket_overload_size is the soft upper limit for the serialized size of a single bucket object, defaulting to 512 KiB. On insert:

  1. A token that already has a posting stays in the bucket that owns it; that bucket is only charged for the appended entry (and, the first time the document lands in it, for the document's token-count entry).
  2. A brand-new token is placed in the current tail if the bucket is empty or its estimated payload stays < limit, including the new document-length entry.
  3. Otherwise insertion advances the tail and retries that token. Its posting and final bucket registration are published together under the bucket lock; existing postings never have a temporary owner. Each token is placed once instead of repeatedly moving a pending token list.

Between flushes size is an estimate accumulated from those charges. Every flush and every reload replaces it with the exact length of the serialized bucket object — which includes the per-document token counts the object carries — so the estimate cannot drift for long. The same flush makes the bucket's doc_ids hint exact.

Note: a token that is already registered in a bucket is never moved by a later insert. This avoids the degenerate case where every insertion creates a new bucket (regression test test_no_excessive_small_buckets), but it also means the limit only applies at placement time: a bucket keeps growing by one entry for every later document that contains one of its tokens. Very frequent terms — stop words, with a tokenizer that keeps them — make their buckets grow with the corpus, and every flush that touches such a bucket rewrites it whole. Stop-word filtering in the tokenizer chain is the effective remedy; compact_buckets repacks whole tokens and never splits one.

4.1 Defragmentation: compact_buckets()

For fragmented indexes left over from earlier versions, or buckets hollowed out by long-running remove operations, compact_buckets() can repack everything in one pass:

  • Scan all postings and estimate each token's CBOR size.
  • Use the Best-Fit-Decreasing algorithm, with BTreeMap<remaining_capacity, bin_index> for O(n log n) packing.
  • Rebuild buckets: bucket IDs are reassigned from 0 upward, all buckets are marked dirty, and the next flush rewrites the full on-disk layout.

The return value (old_count, new_count) is useful for monitoring. Concurrent insert / remove calls are safe: compact_buckets takes the index's internal mutation gate exclusively, so no mutation can observe — or add to — the half-rebuilt bucket map. It must still not overlap a flush (see §5).

Packing uses a fixed conservative CBOR width for the bucket owner, so the old bucket IDs cannot influence the new layout. Repeating compaction without data changes, including after reload, leaves the version and dirty state unchanged. Maintenance also shrinks posting lists whose capacity exceeds four times their length (and 64 entries), retaining space for twice the live length. This applies to single-bucket indexes too and does not itself require a disk write.

5. Concurrency Model

  • Insert, remove, purge and search support concurrent calls. Mutations of the same document id are serialized with internal striped locks, spanning document membership, postings and bucket accounting. Batch purge deduplicates stripe indices and acquires them in ascending order.
  • The lock order is mutation gate → document stripes → maps. Posting creation and bucket registration share a bucket lock. Removal rechecks posting ownership while holding that same bucket lock before unlisting a token.
  • Compaction holds the mutation gate exclusively. The caller must exclude flush from mutations, compaction and other flushes for the entire async call; Collection's exclusive operation lease provides this. One writer per durable index remains the deployment contract.
  • Searches are best-effort concurrent reads, not transactional snapshots. A query samples corpus statistics once; counters converge after completed mutations. No average-length cache is maintained.
  • Metadata-only shells expose persisted statistics through metadata() and stats(), while len() counts loaded documents. Partially loaded or failed-load indexes refuse inserts and flushes; remove/purge and compaction do nothing. is_fully_loaded() indicates writability.
  • Flush captures dirty versions and metadata, then serializes and uploads one bucket at a time. Extra payload memory is bounded by the largest dirty bucket. No internal map guard spans an await, and no bucket is marked clean before the manifest commit succeeds.
  • Top-k selection uses select_nth_unstable_by, followed by sorting the selected results: O(matches + k log k).

6. Persistence Layout

Physically, one index consists of one metadata blob plus multiple bucket blobs. The upper layer is free to store them in files, object storage, a KV store, or any other backend.

<root>/
  metadata                # CBOR-encoded BM25Metadata (carries the bucket manifest)
  buckets/
    b_0_7                 # CBOR-encoded BucketOwned, addressed by (bucket_id, generation)
    b_1_9
    b_2                   # legacy (pre-manifest) object at generation 0, read-only
    ...

Bucket objects are immutable once referenced: every flush writes replaced buckets to a fresh (bucket_id, generation) object and the metadata's manifest is the loader's single source of truth.

6.1 BM25Metadata

pub struct BM25Metadata {
    pub name: String,
    pub config: BM25Config,
    pub stats: BM25Stats,          // version / counts / timestamps / watermarks
    /// Bucket manifest: bucket_id -> generation of the current durable object
    /// (0 = legacy un-suffixed object). Empty when loading pre-manifest data.
    pub buckets: BTreeMap<u32, u64>,
}

flush uses last_saved_version as an idempotency guard: when neither stats.version has increased nor any bucket is dirty, it returns immediately with saved == false.

6.2 Bucket CBOR

struct BucketOwned {
    #[serde(rename = "p")] postings:  FxHashMap<String, PostingValue>,
    #[serde(rename = "d")] doc_tokens: FxHashMap<u64,    usize>,
}

The short field names ("p", "d") are chosen to reduce CBOR size. Note that PostingValue embeds bucket_id, so loading does not require any extra bookkeeping.

6.3 Incremental Flush Flow

let outcome = index.flush(metadata_writer, now_ms, |object, bytes| {
    // write bytes to the storage object for (object.bucket_id, object.generation)
    std::future::ready(Ok(()))
}).await?;
for object in &outcome.obsolete {
    // best-effort: delete the objects the new manifest no longer references
}
  • flush captures dirty versions and the next manifest, then serializes/uploads dirty buckets one at a time. Only postings owned by that bucket are written; document lengths are derived from those postings. Buckets without tokens, other than the tail (max_bucket_id), are left out of the manifest instead of being written as empty objects, and their previous objects — including empty ones written by earlier releases — are reported obsolete. New tokens only go to the tail, so such buckets would otherwise stay empty and be fetched on every open. The tail is kept so the manifest never becomes empty.
  • Each dirty bucket is written to a new object keyed by (bucket_id, generation); the generation is this flush's metadata version, so committed objects are never mutated in place.
  • The metadata — whose manifest maps every live bucket id to its current generation — is written last. That single write is the atomic commit point: a crash or error before it leaves the previous snapshot fully intact (the new objects are unreferenced garbage); after it, the replaced objects are garbage and are returned in FlushOutcome::obsolete for best-effort deletion.
  • compact_buckets needs no special ordering: the repacked layout becomes visible atomically with the next manifest commit, and every pre-compaction object is reported obsolete.

6.4 Startup and Partial Loading

let idx = BM25Index::load_all(tokenizer, metadata_reader, async |object| {
    Ok(read_bucket(object).await?) // Ok(None) means this bucket is not loaded
}).await?;
  • load_metadata restores metadata only, which is useful for lightweight scenarios that need just statistics.
  • With a manifest present, load_buckets reads exactly the referenced (bucket_id, generation) objects. Metadata persisted by pre-manifest releases has no manifest; the loader falls back to scanning bucket ids 0..=max_bucket_id at generation 0 (the legacy un-suffixed objects), and the first flush upgrades the durable layout to the manifest format.
  • load_all / load_buckets permit explicit read-only partial loads through Ok(None). The missing bucket set is retained across incremental loads; completing it enables mutation. Failed or cancelled loads remain read-only and retain already loaded document lengths for retry.
  • Production startup uses load_all_strict; load_buckets_strict is also available. Missing manifest-referenced objects fail the load. Legacy layouts still tolerate holes in their bucket-id scan.
  • If the same token appears in more than one loaded bucket (possible only in legacy data written by the old multi-phase flush), the later bucket id wins. The loader removes that token from the older bucket, rebuilds bucket document-id sets from the winning postings, and marks repaired buckets dirty so the next flush removes stale on-disk ownership.

7. Query Language

search_advanced(query, top_k, params) accepts boolean expressions, and QueryType::parse turns the input string into an AST:

expr     := or_expr
or_expr  := and_expr ( " OR " and_expr )*
and_expr := not_expr ( " AND " not_expr )*
not_expr := "NOT " not_expr | term
term     := chunk ( whitespace chunk )*
chunk    := "(" or_expr ")" | word

Precedence is OR < AND < NOT. The operators are case-sensitive and must stand between spaces (NOT followed by a space): a and b is three search words, and NOT NOT a nests two negations (evaluated as a). Key properties:

  • Multi-term queries default to OR: "quick fox" and "quick OR fox" return the same results in search and search_advanced. A group next to words is one more OR operand: rust (async AND tokio) is rust OR (async AND tokio).
  • Score merging: AND sums the BM25 scores of its subqueries; OR does the same; NOT produces a zero-scored placeholder set used only for filtering, and in an AND context it removes matching items from the result set.
  • Robust parsing: unbalanced parentheses do not panic. They are treated as ordinary characters, which makes direct forwarding of user input safe. An empty group is an empty OR that matches nothing, so a AND () returns nothing.
  • Tokenization belongs to the tokenizer: the parser preserves operand case. Every boolean operand is tokenized independently with one tokenizer clone per query. An OR of terms merges and deduplicates the resulting tokens, never the raw text. Custom context-sensitive tokenizers may therefore distinguish plain text search from boolean word expressions.
  • Candidate execution: AND evaluates selective operands first and restricts later work to the surviving candidates. Each posting is scored in one pass; DF is the posting length, so candidate-restricted scoring keeps global DF/IDF. NOT filters use set membership without computing discarded BM25 scores. IDFs are cached within the query and computed using ln_1p for numerical stability.
  • NOT complement guard: a complement of the whole index (NOT a, a OR NOT b, NOT a AND NOT b) is limited to 10 000 documents. Inside an AND with a positive operand, or under try_search_in_ids, a NOT is taken relative to the documents already matched, which are in memory anyway, so a AND NOT (b AND NOT c) has no limit. Leading NOT pairs cancel at planning time: hello AND NOT (NOT world) is hello AND world, scored. try_search_advanced exposes resource errors; search_advanced returns an empty result on error. Strict parsing rejects combined parentheses/NOT budget exhaustion instead of silently changing query meaning.
  • Multi-byte safe: the delimiters " AND " and " OR " are ASCII, so byte-wise scanning remains safe under UTF-8. Mixed CJK text does not require extra handling.

Example:

let hits = index.search_advanced(
    "(brown AND fox) AND NOT (rare OR sleeps)",
    10,
    None,
);

OR normalization flattens nested ORs and merges their direct token operands even when AND/NOT branches are present. Duplicate normalized tokens contribute once within that OR; redundant OR parentheses do not change scores. Complex branches still contribute their own scores. Empty candidate sets stop before reading postings; other candidate searches continue to use global DF/IDF.

8. Tokenizers

All tokenizers implement tantivy_tokenizer_api::Tokenizer and are composed through TokenizerChain. Important APIs:

Function Role
TokenizerChain::builder(base).filter(f1).filter(f2).build() Compose a tokenization pipeline
default_tokenizer() SimpleTokenizer -> RemoveLongFilter(32) -> LowerCaser -> Stemmer (requires the tantivy feature)
jieba_tokenizer() Prepends JiebaMergeFilter to the pipeline above (requires the tantivy-jieba feature)
collect_tokens(tok, text, inclusive) Standalone tokenization helper for upper-layer preprocessing or non-BM25 matching statistics
flat_full_text_search(tok, query, text) Naive matching without building an index, useful for short passages
detect_script(text) Detects the dominant script by character frequency (Latin / Cyrillic / Arabic / Cjk / Other)

8.1 JiebaMergeFilter

In mixed-script text, SimpleTokenizer treats consecutive Chinese characters as one token. The public generic JiebaMergeFilter re-segments CJK spans and globally orders (offset_from, offset_to, position, text). The built-in jieba_tokenizer() uses the same segmentation and filters but buffers and sorts only one SimpleTokenizer span at a time. Regression tests compare the complete token stream, including offsets and positions, against the generic chain. Both preserve:

  • Chinese text is segmented correctly ("北京市东城区长安街" -> 北京, 东城区, 长安街);
  • English, Russian, Arabic, and other scripts retain the stemmed and lowercased output of the primary pipeline;
  • The resulting TokenStream is still monotonic, so downstream BM25 consumption remains correct.

RemoveLongFilter::limit(32) retains tokens with UTF-8 byte length less than 32, not 32 Unicode characters. collect_tokens copies a term only on its first occurrence in the frequency map.

Note: collect_tokens filters out single-byte tokens where token.text.len() <= 1 (length measured in bytes), which removes punctuation and isolated ASCII letters. Single Chinese characters are unaffected because their UTF-8 length is at least 3 bytes.

9. Error Handling

pub enum BM25Error {
    Generic       { name: String, source: BoxError },
    Serialization { name: String, source: BoxError },
    NotFound      { name: String, id: u64 },
    AlreadyExists { name: String, id: u64 },
    TokenizeFailed{ name: String, id: u64, text: String },
}

Generic is used for errors returned by I/O closures and by the query guards of try_search_advanced; Serialization wraps cbor2 failures; AlreadyExists and TokenizeFailed occur during insert (TokenizeFailed carries at most the first 256 bytes of the document text); NotFound is never raised by the crate and is left to upper-layer APIs for idempotent validation.

10. Configuration and Tuning

let cfg = BM25Config {
    bm25: BM25Params { k1: 1.5, b: 0.6 },
    bucket_overload_size: 1024 * 1024,     // 1 MiB
};
let index = BM25Index::new("mem".into(), jieba_tokenizer(), Some(cfg));

Tuning guidance:

  • k1 in [1.2, 2.0]: for shorter documents with denser keywords, increasing it slightly can amplify the effect of high-frequency terms.
  • b in [0.5, 0.9]: choose 0.75 to 0.9 when document lengths vary significantly; drop closer to 0.5 when most documents are similar in length.
  • bucket_overload_size:
    • Small (for example 64KiB): lower I/O amplification for incremental flushes, suitable for frequent checkpoints.
    • Large (for example 2MiB): fewer total buckets and faster full loads, suitable for read-heavy AI memory stores.
  • Periodic compaction: call compact_buckets() periodically in a background task to keep bucket counts stable. It is safe against concurrent insert / remove (the crate gates it internally); schedule it so it does not overlap a flush.

11. Testing and Benchmarks

Test coverage includes:

  • cargo test -p anda_db_tfs --all-features
  • cargo test -p anda_db_tfs --no-default-features
  • cargo test -p anda_db_tfs
  • cargo test -p anda_db_tfs --features full --example tfs_demo
  • Correctness of insert / remove / search, bucket serialization and partial loading, result invariance after compaction, the regression test test_no_excessive_small_buckets, UTF-8 query parsing, and more.

Benchmarks: cargo bench -p anda_db_tfs --features full --bench tfs_tokenizer and cargo bench -p anda_db_tfs --features full --bench tfs_index. The latter covers scoring, mutations, load/flush, compaction and mixed Chinese text. See measurements and the review checklist for reproducible comparisons and remaining storage tradeoffs.

12. Quick Start

use anda_db_tfs::{BM25Index, default_tokenizer};

let idx = BM25Index::new("notes".into(), default_tokenizer(), None);
idx.insert(1, "The quick brown fox jumps over the lazy dog", 0).unwrap();
idx.insert(2, "A fast brown fox runs past the lazy dog",     0).unwrap();
idx.insert(3, "The lazy dog sleeps all day",                 0).unwrap();

for (id, score) in idx.search("fox", 10, None) {
    println!("doc {id}: {score:.3}");
}

for (id, score) in idx.search_advanced("(brown AND fox) AND NOT sleeps", 10, None) {
    println!("doc {id}: {score:.3}");
}

Persisting to the local filesystem uses the tested write_atomic helper. It creates a temporary file beside the target, writes and syncs it, renames it atomically, and syncs the directory on Unix. If a crashed process left the first temporary name behind, the helper advances its sequence until it creates a fresh file; it never deletes an existing temporary file that could belong to another writer. The complete runnable adapter is in tfs_demo.

use std::{fs, path::{Path, PathBuf}};
use anda_db_tfs::{BoxError, BucketObject};
// Import write_atomic from the linked helper into your application.

fn bucket_path(object: BucketObject) -> PathBuf {
    if object.generation == 0 {
        format!("./idx/b_{}.cbor", object.bucket_id).into()
    } else {
        format!("./idx/b_{}_{}.cbor", object.bucket_id, object.generation).into()
    }
}

fs::create_dir_all("./idx")?;
let outcome = idx.flush_with(
    now_ms,
    |bytes| std::future::ready(
        write_atomic(Path::new("./idx/metadata.cbor"), &bytes)
            .map_err(BoxError::from)
    ),
    |object, bytes| std::future::ready(
        write_atomic(&bucket_path(object), &bytes).map_err(BoxError::from)
    ),
).await?;
for object in outcome.obsolete {
    let _ = fs::remove_file(bucket_path(object));
}

let idx = BM25Index::load_all_strict(
    default_tokenizer(), fs::File::open("./idx/metadata.cbor")?,
    async |object| match fs::read(bucket_path(object)) {
        Ok(bytes) => Ok(Some(bytes)),
        Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(None),
        Err(err) => Err(err.into()),
    },
).await?;

Never open the committed metadata with File::create before flush: a clean flush invokes no callback, so that would leave an empty file. flush<W: Write> writes and flushes its writer but cannot supply atomic replacement or filesystem durability on its own. A metadata callback with an uncertain outcome requires reopening from durable state.

13. Usage Notes

  1. Removal wants the original text: remove(id, text, now_ms) re-tokenizes the text to find the document's postings. If those entries do not add up to the document's token count (wrong text, or a changed tokenizer), remove sweeps every posting list for the id: still correct, but O(index size). When the text is genuinely unrecoverable — a repair path whose document bodies are gone — use purge_ids(&BTreeSet<u64>, now_ms): it sweeps every posting list once for the whole set, drops the ids from doc_tokens and total_tokens, and marks the affected buckets dirty. Pass all the dead ids in one call rather than looping.
  2. top_k = 0: kept for API compatibility. It returns an empty set and does not trigger sorting.
  3. Flush coordination: the crate does not serialize flushes internally. The caller must ensure a flush never overlaps mutations, compaction, or another flush (anda_db's Collection already guarantees this); a single writer per durable index is a deployment contract.
  4. Search semantics under partial loading: if load_buckets skips a posting bucket, terms owned by that bucket are unavailable. Loaded buckets also carry the document lengths needed to score their postings, so len() may include every document touched by those loaded terms even when other buckets are skipped. Search results remain the natural subset of the loaded postings.
  5. Embedded-only: this library is intended for in-process embedding inside AndaDB and does not provide HTTP or gRPC services. For remote access, use anda_db_server or anda_db_shard_proxy.

14. References

  • Robertson & Zaragoza. The Probabilistic Relevance Framework: BM25 and Beyond, 2009.
  • tantivy_tokenizer_api - tokenizer trait.
  • For regression cases and design discussion, see rs/anda_db_tfs/src/bm25/tests.rs, src/bm25/regression_tests.rs and tests/regressions.rs, and anda_db_btree.md for the related bucket strategy.